chromadb-tutorial-1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dcf583fd-84d8-4a0f-9960-84c008a93388",
   "metadata": {},
   "source": [
    "## ChromaDB Tutorial\n",
    "\n",
    "- [Try Chroma](https://docs.trychroma.com/guides/multimodal)\n",
    "      - [Chroma Guide](https://docs.trychroma.com/guides)\n",
    "\n",
    "- [Embeddings and Vector Databases With ChromaDB](https://realpython.com/chromadb-vector-database/)\n",
    "    - intro on vector Embeddings\n",
    "    - chromadb query\n",
    "    - LLM \n",
    "\n",
    "- [Chroma DB Tutorial: A Step-By-Step Guide](https://www.datacamp.com/tutorial/chromadb-tutorial-step-by-step-guide)  \n",
    "\n",
    "- [Chroma Vector Database Tutorial](https://anderfernandez.com/en/blog/chroma-vector-database-tutorial/)\n",
    "    - deploy to docker\n",
    "    - deploy to clickhouse\n",
    "    - connect in python\n",
    "    - used SP500 company info\n",
    "    - filter by metadata\n",
    "\n",
    "- [Getting Started with Chroma DB: A Beginner’s Tutorial](https://medium.com/@pierrelouislet/getting-started-with-chroma-db-a-beginners-tutorial-6efa32300902)\n",
    "    - brief \n",
    "    - use docker\n",
    "\n",
    "see `readme.md`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "928cb81a-ddb0-4b15-9915-5f3c4b390822",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: chromadb\n",
      "Version: 0.5.0\n",
      "Summary: Chroma.\n",
      "Home-page: \n",
      "Author: \n",
      "Author-email: Jeff Huber <jeff@trychroma.com>, Anton Troynikov <anton@trychroma.com>\n",
      "License: \n",
      "Location: /home/papagame/anaconda3/envs/vanna/lib/python3.11/site-packages\n",
      "Requires: bcrypt, build, chroma-hnswlib, fastapi, grpcio, importlib-resources, kubernetes, mmh3, numpy, onnxruntime, opentelemetry-api, opentelemetry-exporter-otlp-proto-grpc, opentelemetry-instrumentation-fastapi, opentelemetry-sdk, orjson, overrides, posthog, pydantic, pypika, PyYAML, requests, tenacity, tokenizers, tqdm, typer, typing-extensions, uvicorn\n",
      "Required-by: \n"
     ]
    }
   ],
   "source": [
    "!pip show chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "490e5384-6dff-4611-99ac-f0bf3352c10b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/papagame/projects/wgong/py4kids\n"
     ]
    }
   ],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "98d28c71-538d-4502-90d7-c6569e92dea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "\n",
    "CHROMA_DATA_PATH = \"./db\"\n",
    "EMBED_MODEL = \"all-MiniLM-L6-v2\"\n",
    "COLLECTION_NAME = \"demo_docs\"\n",
    "\n",
    "client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d59ad2-20ee-47b4-8f5e-4dbfbebed37a",
   "metadata": {},
   "source": [
    "### Create collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6d2914e2-32cb-4c65-ab9a-ce22b75d25bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(\n",
    "#     model_name=EMBED_MODEL\n",
    "# )\n",
    "\n",
    "collection = client.get_or_create_collection(\n",
    "    name=COLLECTION_NAME,\n",
    "    # embedding_function=embedding_func,\n",
    "    metadata={\"hnsw:space\": \"cosine\"},\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c7930e2-ef6c-4cf7-89bd-f9ebdaf6b004",
   "metadata": {},
   "source": [
    "### List collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "89d2636a-24d2-4498-b6df-79bc084582f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collection name: demo_docs\n"
     ]
    }
   ],
   "source": [
    "for collection in client.list_collections():\n",
    "    print(f\"Collection name: {collection.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a58445c8-e176-48ca-8f84-f508c69838e9",
   "metadata": {},
   "source": [
    "### Add Docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8e73631b-2812-4bc0-b8a3-5c5c7e72766f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Add of existing embedding ID: id0\n",
      "Add of existing embedding ID: id1\n",
      "Add of existing embedding ID: id2\n",
      "Add of existing embedding ID: id3\n",
      "Add of existing embedding ID: id4\n",
      "Add of existing embedding ID: id5\n",
      "Add of existing embedding ID: id6\n",
      "Add of existing embedding ID: id7\n",
      "Add of existing embedding ID: id8\n",
      "Add of existing embedding ID: id9\n",
      "Insert of existing embedding ID: id0\n",
      "Insert of existing embedding ID: id1\n",
      "Insert of existing embedding ID: id2\n",
      "Insert of existing embedding ID: id3\n",
      "Insert of existing embedding ID: id4\n",
      "Insert of existing embedding ID: id5\n",
      "Insert of existing embedding ID: id6\n",
      "Insert of existing embedding ID: id7\n",
      "Insert of existing embedding ID: id8\n",
      "Insert of existing embedding ID: id9\n"
     ]
    }
   ],
   "source": [
    "# sample docs\n",
    "documents = [\n",
    "    \"The latest iPhone model comes with impressive features and a powerful camera.\",\n",
    "    \"Exploring the beautiful beaches and vibrant culture of Bali is a dream for many travelers.\",\n",
    "    \"Einstein's theory of relativity revolutionized our understanding of space and time.\",\n",
    "    \"Traditional Italian pizza is famous for its thin crust, fresh ingredients, and wood-fired ovens.\",\n",
    "    \"The American Revolution had a profound impact on the birth of the United States as a nation.\",\n",
    "    \"Regular exercise and a balanced diet are essential for maintaining good physical health.\",\n",
    "    \"Leonardo da Vinci's Mona Lisa is considered one of the most iconic paintings in art history.\",\n",
    "    \"Climate change poses a significant threat to the planet's ecosystems and biodiversity.\",\n",
    "    \"Startup companies often face challenges in securing funding and scaling their operations.\",\n",
    "    \"Beethoven's Symphony No. 9 is celebrated for its powerful choral finale, 'Ode to Joy.'\",\n",
    "]\n",
    "\n",
    "genres = [\n",
    "    \"technology\",\n",
    "    \"travel\",\n",
    "    \"science\",\n",
    "    \"food\",\n",
    "    \"history\",\n",
    "    \"fitness\",\n",
    "    \"art\",\n",
    "    \"climate change\",\n",
    "    \"business\",\n",
    "    \"music\",\n",
    "]\n",
    "\n",
    "collection.add(\n",
    "    documents=documents,\n",
    "    ids=[f\"id{i}\" for i in range(len(documents))],\n",
    "    metadatas=[{\"genre\": g} for g in genres]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1135247-1e7d-4046-a581-25db7f5502f4",
   "metadata": {},
   "source": [
    "### Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fa9615db-3ef3-4608-a85c-25d44d7f4f9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Find me some delicious food!\"],\n",
    "    n_results=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8e963b0c-05e6-4bb5-9a73-bb20407d4303",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id3', 'id1']],\n",
       " 'distances': [[1.5276524111861616, 1.6584325895032108]],\n",
       " 'metadatas': [[{'genre': 'food'}, {'genre': 'travel'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['Traditional Italian pizza is famous for its thin crust, fresh ingredients, and wood-fired ovens.',\n",
       "   'Exploring the beautiful beaches and vibrant culture of Bali is a dream for many travelers.']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ec62c232-e9e2-4c11-97f3-66ef18c22deb",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Teach me about history\",\n",
    "                 \"What's going on in the world?\"],\n",
    "    include=[\"documents\", \"distances\", 'metadatas'],\n",
    "    n_results=2\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0adff3ee-f6c2-4654-96f5-0a8bd9c5488c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id2', 'id4'], ['id7', 'id2']],\n",
       " 'distances': [[1.2531765850067067, 1.3808384485195626],\n",
       "  [1.6005885986562254, 1.776421301382694]],\n",
       " 'metadatas': [[{'genre': 'science'}, {'genre': 'history'}],\n",
       "  [{'genre': 'climate change'}, {'genre': 'science'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [[\"Einstein's theory of relativity revolutionized our understanding of space and time.\",\n",
       "   'The American Revolution had a profound impact on the birth of the United States as a nation.'],\n",
       "  [\"Climate change poses a significant threat to the planet's ecosystems and biodiversity.\",\n",
       "   \"Einstein's theory of relativity revolutionized our understanding of space and time.\"]],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c8427997-9dde-4047-aac0-60bb71b800e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id2']],\n",
       " 'distances': [[1.5251642004851387]],\n",
       " 'metadatas': [[{'genre': 'science'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [[\"Einstein's theory of relativity revolutionized our understanding of space and time.\"]],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Teach me about music history\"],\n",
    "    n_results=1\n",
    ")\n",
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "01ffc588-ec0c-4ed4-8760-b9f3f2f2de66",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id9']],\n",
       " 'distances': [[1.6372656175451257]],\n",
       " 'metadatas': [[{'genre': 'music'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [[\"Beethoven's Symphony No. 9 is celebrated for its powerful choral finale, 'Ode to Joy.'\"]],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Teach me about music history\"],\n",
    "    where={\"genre\": {\"$eq\": \"music\"}},\n",
    "    n_results=1,\n",
    ")\n",
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "550339b1-1098-4503-902f-c05d86825b7f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f80f60d4-6a64-465f-bc58-4748e8194269",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id9', 'id4']],\n",
       " 'distances': [[1.6372656175451257, 1.6400827280864978]],\n",
       " 'metadatas': [[{'genre': 'music'}, {'genre': 'history'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [[\"Beethoven's Symphony No. 9 is celebrated for its powerful choral finale, 'Ode to Joy.'\",\n",
       "   'The American Revolution had a profound impact on the birth of the United States as a nation.']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Teach me about music history\"],\n",
    "    where={\"genre\": {\"$in\": [\"music\", \"history\"]}},\n",
    "    n_results=2,\n",
    ")\n",
    "\n",
    "query_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c581d7-661e-4228-bc3c-74c0ff091f9f",
   "metadata": {},
   "source": [
    "### Update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1833065c-d7f6-45d1-8c03-49bb3b4aeb88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': ['id1', 'id2'],\n",
       " 'embeddings': None,\n",
       " 'metadatas': [{'genre': 'tech'}, {'genre': 'beaches'}],\n",
       " 'documents': ['The new iPhone is awesome!', 'Bali has beautiful beaches'],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "collection.update(\n",
    "    ids=[\"id1\", \"id2\"],\n",
    "    documents=[\"The new iPhone is awesome!\",\n",
    "               \"Bali has beautiful beaches\"],\n",
    "    metadatas=[{\"genre\": \"tech\"}, {\"genre\": \"beaches\"}]\n",
    ")\n",
    "\n",
    "query_results = collection.get(ids=[\"id1\", \"id2\"])\n",
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9d99b86d-a084-4ddf-81b7-5745b7ce4e0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id1'], ['id2']],\n",
       " 'distances': None,\n",
       " 'metadatas': [[{'genre': 'tech'}], [{'genre': 'beaches'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['The new iPhone is awesome!'], ['Bali has beautiful beaches']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_results = collection.query(\n",
    "    query_texts=[\"Which phone is great\", \"where to find beautiful beach\"],\n",
    "    n_results=1,\n",
    "    include=['metadatas', 'documents'],  # ids is always included\n",
    ")\n",
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "94e064fc-21e4-44bb-8042-9d8a331a5155",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9\n",
      "8\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'ids': ['id3'],\n",
       " 'embeddings': None,\n",
       " 'metadatas': [{'genre': 'food'}],\n",
       " 'documents': ['Traditional Italian pizza is famous for its thin crust, fresh ingredients, and wood-fired ovens.'],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(collection.count())\n",
    "\n",
    "collection.delete(ids=[\"id2\"])\n",
    "\n",
    "print(collection.count())\n",
    "\n",
    "collection.get([\"id1\", \"id2\", \"id3\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da16a6d8-c0d7-4a57-8fb8-c15c3c159ca8",
   "metadata": {},
   "source": [
    "### Remove all collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e9b5837c-d9c5-4b40-b67e-d8d123093dd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing collection: demo_docs\n"
     ]
    }
   ],
   "source": [
    "for collection in client.list_collections():\n",
    "    print(f\"Removing collection: {collection.name}\")\n",
    "    client.delete_collection(name=collection.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b0c519-fb08-4893-a8fa-16aba03d6a3f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "6210914e-65fb-43b3-8cac-54e42d1b40c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'83ca2e16-13e9-4cb7-831d-0f47f406a9f0','2fba5c16-790f-4852-a580-63749fde839b','897c7090-6cbf-4e2d-8ffa-2217cdc127c0','3b4962a5-3791-4ff9-b98f-82b164b28cc2','38f8c0b8-a754-4baf-8521-ff44dc3b59f0','67c9a261-0f99-4c9b-b44d-de93cbb8ed5d','e6dd68d3-4bdc-4a76-8315-bb6fee65465e'\n"
     ]
    }
   ],
   "source": [
    "ids = \"\"\"\n",
    "83ca2e16-13e9-4cb7-831d-0f47f406a9f0\n",
    "2fba5c16-790f-4852-a580-63749fde839b\n",
    "897c7090-6cbf-4e2d-8ffa-2217cdc127c0\n",
    "3b4962a5-3791-4ff9-b98f-82b164b28cc2\n",
    "38f8c0b8-a754-4baf-8521-ff44dc3b59f0\n",
    "67c9a261-0f99-4c9b-b44d-de93cbb8ed5d\n",
    "e6dd68d3-4bdc-4a76-8315-bb6fee65465e\n",
    "\"\"\"\n",
    "print(\",\".join([f\"'{i.strip()}'\" for i in ids.split() if i.strip()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75b621d2-ee0b-403c-a1c0-8ee505344928",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}