basic_rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/sabhyachhabria/Library/Python/3.9/lib/python/site-packages/langchain/__init__.py:24: UserWarning: Importing BasePromptTemplate from langchain root module is no longer supported.\n",
      "  warnings.warn(\n",
      "/Users/sabhyachhabria/Library/Python/3.9/lib/python/site-packages/langchain/__init__.py:24: UserWarning: Importing PromptTemplate from langchain root module is no longer supported.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from llama_index import download_loader, Document\n",
    "from llama_index.text_splitter import SentenceSplitter\n",
    "from llama_index.schema import TextNode\n",
    "import csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_csv(file_path: str) -> None:\n",
    "    events_list = []\n",
    "    with open(file_path, mode='r') as csv_file:\n",
    "        csv_reader = csv.DictReader(csv_file)\n",
    "        for row in csv_reader:\n",
    "            events_list.append(row)\n",
    "    return events_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_splitter = SentenceSplitter(\n",
    "    chunk_size=1024,\n",
    "    # separator=\" \",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "events = read_csv(\"./events.csv\")\n",
    "documents = [Document(text=str(t), metadata={\"event_id\": t[\"Id\"], \"event_name\":  t[\"Name\"] }) for t in events]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_chunks = []\n",
    "# maintain relationship with source doc index, to help inject doc metadata in (3)\n",
    "doc_idxs = []\n",
    "for doc_idx, doc in enumerate(documents):\n",
    "    cur_text_chunks = text_splitter.split_text(doc.text)\n",
    "    text_chunks.extend(cur_text_chunks)\n",
    "    doc_idxs.extend([doc_idx] * len(cur_text_chunks))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = []\n",
    "for idx, text_chunk in enumerate(text_chunks):\n",
    "    node = TextNode(\n",
    "        text=text_chunk,\n",
    "    )\n",
    "    src_doc = documents[doc_idxs[idx]]\n",
    "    node.metadata = src_doc.metadata\n",
    "    nodes.append(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "event_id: 1\n",
      "event_name: The Phantom of the Opera\n",
      "\n",
      "{'Id': '1', 'Name': 'The Phantom of the Opera', 'One Line Description': 'Experience the iconic musical', 'Location': 'Majestic Theatre NYC', 'Time': 'September 5 2023 8:00 PM', 'Avg. Rating': '4.6', 'Ticket Price': '$75'}\n"
     ]
    }
   ],
   "source": [
    "print(nodes[0].get_content(metadata_mode=\"all\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "chroma_client = chromadb.PersistentClient(\"./datastore/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Collection events already exists.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/Users/sabhyachhabria/Desktop/semantic-search/basic_rag.ipynb Cell 9\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/sabhyachhabria/Desktop/semantic-search/basic_rag.ipynb#X11sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m collection \u001b[39m=\u001b[39m chroma_client\u001b[39m.\u001b[39;49mcreate_collection(name\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mevents\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
      "File \u001b[0;32m~/Library/Python/3.9/lib/python/site-packages/chromadb/api/segment.py:132\u001b[0m, in \u001b[0;36mSegmentAPI.create_collection\u001b[0;34m(self, name, metadata, embedding_function, get_or_create)\u001b[0m\n\u001b[1;32m    124\u001b[0m         \u001b[39mreturn\u001b[39;00m Collection(\n\u001b[1;32m    125\u001b[0m             client\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m,\n\u001b[1;32m    126\u001b[0m             \u001b[39mid\u001b[39m\u001b[39m=\u001b[39mexisting[\u001b[39m0\u001b[39m][\u001b[39m\"\u001b[39m\u001b[39mid\u001b[39m\u001b[39m\"\u001b[39m],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    129\u001b[0m             embedding_function\u001b[39m=\u001b[39membedding_function,\n\u001b[1;32m    130\u001b[0m         )\n\u001b[1;32m    131\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 132\u001b[0m         \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mCollection \u001b[39m\u001b[39m{\u001b[39;00mname\u001b[39m}\u001b[39;00m\u001b[39m already exists.\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    134\u001b[0m \u001b[39m# TODO: remove backwards compatibility in naming requirements\u001b[39;00m\n\u001b[1;32m    135\u001b[0m check_index_name(name)\n",
      "\u001b[0;31mValueError\u001b[0m: Collection events already exists."
     ]
    }
   ],
   "source": [
    "collection = chroma_client.create_collection(name=\"events\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids = [doc.id_ for doc in documents]\n",
    "metadata = [doc.metadata for doc in documents]\n",
    "text = [doc.text for doc in documents]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "collection.add(\n",
    "    documents=text,\n",
    "    metadatas=metadata,\n",
    "    ids=ids\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"broadway in nyc after Jan 2025 after 8:00 PM\"],\n",
    "    n_results=5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "collection.delete(ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "chroma_client.delete_collection(\"events\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Collection(name=events)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}