chroma_fast_vectorize_qa.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "\n",
    "chroma_client = chromadb.PersistentClient(path=\"my_vectordb\")\n",
    "\n",
    "# Select the embedding model to use.\n",
    "# List of model names can be found here https://www.sbert.net/docs/pretrained_models.html\n",
    "sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=\"all-mpnet-base-v2\")\n",
    "\n",
    "# Create the collection, aka vector database. Or, if database already exist, then use it. Specify the model that we want to use to do the embedding.\n",
    "collection = chroma_client.get_collection(name=\"got\", embedding_function=sentence_transformer_ef)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19699 0.331 In season \"season-04\", episode \"e9-The Watchers on the Wall\", YGRITTE said, \" You know nothing, Jon Snow.\"\n",
      "8373 0.344 In season \"season-02\", episode \"e7-A Man Without Honor\", YGRITTE said, \" You know nothing, Jon Snow.\"\n",
      "10755 0.346 In season \"season-03\", episode \"e10-Mhysa\", YGRITTE said, \" You know nothing, Jon Snow.\"\n",
      "13633 0.348 In season \"season-03\", episode \"e7-The Bear and the Maiden Fair\", YGRITTE said, \" You know nothing, Jon Snow.\"\n",
      "8479 0.441 In season \"season-02\", episode \"e7-A Man Without Honor\", YGRITTE said, \" \"And I thought that we were done, but he said, \"'Turn back around.\" I'll tell you what, Jon Snow. Since it's gonna be your word against mine, and since you can't talk about it without blushing, we may as well just JON: What, right here in the muck?\"\n",
      "14046 0.482 In season \"season-03\", episode \"e7-The Bear and the Maiden Fair\", YGRITTE said, \" You know nothing--\"\n",
      "13619 0.488 In season \"season-03\", episode \"e7-The Bear and the Maiden Fair\", YGRITTE said, \" You think? You don't know?\"\n",
      "13347 0.495 In season \"season-03\", episode \"e6-The Climb\", YGRITTE said, \" You staring at my ass, Jon Snow?\"\n",
      "14052 0.501 In season \"season-03\", episode \"e7-The Bear and the Maiden Fair\", YGRITTE said, \" You don't know that.\"\n",
      "13260 0.537 In season \"season-03\", episode \"e6-The Climb\", YGRITTE said, \" Only in the beginning. You're a proper lover, Jon Snow. And don't worry. Your secret's safe with me.\"\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"Ygritte repeatedly said Jon Snow knows what?\"],\n",
    "    n_results=10,\n",
    "    include=['documents', 'distances', 'metadatas'],\n",
    "    #where={'speaker':'YGRITTE'}\n",
    ")\n",
    "\n",
    "for j in range(len(results['ids'][0])):\n",
    "    id       = results[\"ids\"][0][j]\n",
    "    distance = results['distances'][0][j]\n",
    "    document = results['documents'][0][j]\n",
    "    metadata = results['metadatas'][0][j]\n",
    "    print(f'{id} {distance:.3f} {document}') "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chroma_env",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}