test-chromadb.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install chromadb\n",
    "# !pip install sentence_transformers\n",
    "# !pip install ipywidgets\n",
    "# !pip install pandas\n",
    "# !pip install blingfire"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [],
   "source": [
    "import blingfire\n",
    "from dotenv import load_dotenv\n",
    "from pathlib import Path\n",
    "import chromadb\n",
    "from chromadb.utils import embedding_functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 230,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Importing the environment file\n",
    "dotenv_path = Path('/home/prem/chromadb/.env')\n",
    "load_dotenv(dotenv_path=dotenv_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 231,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Variable names\n",
    "db_path = '/home/prem/chromadb/database_folder'\n",
    "st_collection_name = 'st_collection'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 232,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "\"\"\"\n",
    "Custom Embedding Function\n",
    "Initialize this with any model that that you have created. \n",
    "Here, I have taken the example of the tokenizer for 'bert-base-cased'. \n",
    "This could be loaded from any folder that you have stored your  tokenizer in.\n",
    "\"\"\"\n",
    "class MyEmbeddingFunction(EmbeddingFunction):\n",
    "    def __init__(self):\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')\n",
    "        self.model = AutoModel.from_pretrained('bert-base-cased')\n",
    "\n",
    "\n",
    "    def __call__(self, input: Documents) -> Embeddings:\n",
    "        # Embedding the documents\n",
    "        list_emb = []\n",
    "\n",
    "        for doc in input:\n",
    "            tokens = self.tokenizer(doc,\n",
    "                                    padding='max_length',\n",
    "                                    return_tensors='pt')\n",
    "            output = self.model(**tokens)\n",
    "            embeddins = output['last_hidden_state'][0].detach().flatten().tolist()\n",
    "            list_emb.append(embeddins)\n",
    "        return list_emb\n",
    "    \n",
    "    \n",
    "# Initializing my custom embedding function\n",
    "st_ef = MyEmbeddingFunction()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 238,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Level 0 Type: <class 'list'>\n",
      "Length: 2\n",
      "Level 1 Type: <class 'list'>\n",
      "Length: 393216\n",
      "Level 2 Type: <class 'float'>\n"
     ]
    }
   ],
   "source": [
    "# Testing the embedding function\n",
    "embeds = MyEmbeddingFunction()\n",
    "embedding = embeds(['This is a sample sentence.', 'This is a second sentence'])\n",
    "print('Level 0 Type:', type(embedding))\n",
    "print('Length:', len(embedding))\n",
    "print('Level 1 Type:', type(embedding[0]))\n",
    "print('Length:', len(embedding[0]))\n",
    "print('Level 2 Type:', type(embedding[0][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 234,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Run this on the terminal to make the db available on the server\n",
    "# chroma run --path /home/prem/chromadb/database_folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 235,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Heartbeat: 1705347633136855283\n"
     ]
    }
   ],
   "source": [
    "# Create a client to connect to the DB\n",
    "chroma_client = chromadb.HttpClient(host='localhost', port=8000)\n",
    "\n",
    "# Check the client connection\n",
    "print('Heartbeat:',chroma_client.heartbeat()) # returns a nanosecond heartbeat. Useful for making sure the client remains connected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 236,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collection 'st_collection' initialized\n",
      "Total number of records: 10\n"
     ]
    }
   ],
   "source": [
    "# Create or get the collection\n",
    "try:\n",
    "    st_collection = chroma_client.get_or_create_collection(name=st_collection_name,\n",
    "                                                           embedding_function=st_ef,\n",
    "                                                           metadata={\"hnsw:space\": \"cosine\"})\n",
    "    print(f\"Collection '{st_collection_name}' initialized\")\n",
    "except:\n",
    "    print('Unable to create or get the collection.')\n",
    "\n",
    "# Printing the total number of records in the collection\n",
    "print(f'Total number of records: {st_collection.count()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # I have commented this code as the embeddings are too large. This is to avoid all the scrolling.\n",
    "# # Take a look at a sample of the existing data. (10)\n",
    "# st_collection.peek()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collection 'st_collection' initialized\n",
      "Total number of records: 0\n"
     ]
    }
   ],
   "source": [
    "# Deleting the collection\n",
    "chroma_client.delete_collection('st_collection')\n",
    "\n",
    "# Then recreate it using the same code\n",
    "try:\n",
    "    st_collection = chroma_client.get_or_create_collection(name=st_collection_name,\n",
    "                                                           embedding_function=st_ef,\n",
    "                                                           metadata={\"hnsw:space\": \"cosine\"})\n",
    "    print(f\"Collection '{st_collection_name}' initialized\")\n",
    "except:\n",
    "    print('Unable to create or get the collection.')\n",
    "\n",
    "# Printing the total number of records in the collection\n",
    "print(f'Total number of records: {st_collection.count()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Empty collection\n"
     ]
    }
   ],
   "source": [
    "# Another way to empty the collection but without deleting it.\n",
    "if st_collection.count() > 0:\n",
    "    # Getting the list of IDs so that we can clear the collection without deleting it\n",
    "    result = st_collection.get()\n",
    "    print(result['ids'])\n",
    "\n",
    "    # Remove these records from the collection.\n",
    "    st_collection.delete(ids=result['ids'])\n",
    "\n",
    "    # Confirm deletion by printing the number of records\n",
    "    print(f'Total number of records: {st_collection.count()}')\n",
    "\n",
    "else:\n",
    "    print('Empty collection')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>review_id</th>\n",
       "      <th>author_id</th>\n",
       "      <th>rating</th>\n",
       "      <th>is_recommended</th>\n",
       "      <th>helpfulness</th>\n",
       "      <th>total_feedback_count</th>\n",
       "      <th>total_neg_feedback_count</th>\n",
       "      <th>total_pos_feedback_count</th>\n",
       "      <th>submission_time</th>\n",
       "      <th>review_text</th>\n",
       "      <th>review_title</th>\n",
       "      <th>skin_tone</th>\n",
       "      <th>eye_color</th>\n",
       "      <th>skin_type</th>\n",
       "      <th>hair_color</th>\n",
       "      <th>product_id</th>\n",
       "      <th>product_name</th>\n",
       "      <th>brand_name</th>\n",
       "      <th>price_usd</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>P7880-42549319432-1</td>\n",
       "      <td>42549319432</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2023-03-21</td>\n",
       "      <td>I really wanted to love this, and I would’ve i...</td>\n",
       "      <td>Has fragrance</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>P7880</td>\n",
       "      <td>Soy Hydrating Gentle Face Cleanser</td>\n",
       "      <td>fresh</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>P7880-31124221503-1</td>\n",
       "      <td>31124221503</td>\n",
       "      <td>5</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2023-03-19</td>\n",
       "      <td>Makeup remover, gentle cleanser, and all aroun...</td>\n",
       "      <td>Best. Cleanser. Ever.</td>\n",
       "      <td>fair</td>\n",
       "      <td>brown</td>\n",
       "      <td>combination</td>\n",
       "      <td>blonde</td>\n",
       "      <td>P7880</td>\n",
       "      <td>Soy Hydrating Gentle Face Cleanser</td>\n",
       "      <td>fresh</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>P7880-20246074916-1</td>\n",
       "      <td>20246074916</td>\n",
       "      <td>5</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>2023-03-15</td>\n",
       "      <td>I have been using this for almost 10 years. Lo...</td>\n",
       "      <td>The cleanser I have used for 10 years</td>\n",
       "      <td>NaN</td>\n",
       "      <td>hazel</td>\n",
       "      <td>combination</td>\n",
       "      <td>NaN</td>\n",
       "      <td>P7880</td>\n",
       "      <td>Soy Hydrating Gentle Face Cleanser</td>\n",
       "      <td>fresh</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>P7880-5182718480-1</td>\n",
       "      <td>5182718480</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>2023-03-10</td>\n",
       "      <td>I wanted to love this so bad because it felt s...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>mediumTan</td>\n",
       "      <td>brown</td>\n",
       "      <td>combination</td>\n",
       "      <td>brown</td>\n",
       "      <td>P7880</td>\n",
       "      <td>Soy Hydrating Gentle Face Cleanser</td>\n",
       "      <td>fresh</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>P7880-1840061447-1</td>\n",
       "      <td>1840061447</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2023-03-09</td>\n",
       "      <td>I bought this bc i wanted to have a gentle cle...</td>\n",
       "      <td>Burns and breakouts</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>P7880</td>\n",
       "      <td>Soy Hydrating Gentle Face Cleanser</td>\n",
       "      <td>fresh</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             review_id    author_id  rating  is_recommended  helpfulness  \\\n",
       "0  P7880-42549319432-1  42549319432       2             0.0          NaN   \n",
       "1  P7880-31124221503-1  31124221503       5             1.0          1.0   \n",
       "2  P7880-20246074916-1  20246074916       5             1.0          1.0   \n",
       "3   P7880-5182718480-1   5182718480       1             0.0          0.0   \n",
       "4   P7880-1840061447-1   1840061447       1             0.0          0.0   \n",
       "\n",
       "   total_feedback_count  total_neg_feedback_count  total_pos_feedback_count  \\\n",
       "0                     0                         0                         0   \n",
       "1                     1                         0                         1   \n",
       "2                     2                         0                         2   \n",
       "3                     2                         2                         0   \n",
       "4                     1                         1                         0   \n",
       "\n",
       "  submission_time                                        review_text  \\\n",
       "0      2023-03-21  I really wanted to love this, and I would’ve i...   \n",
       "1      2023-03-19  Makeup remover, gentle cleanser, and all aroun...   \n",
       "2      2023-03-15  I have been using this for almost 10 years. Lo...   \n",
       "3      2023-03-10  I wanted to love this so bad because it felt s...   \n",
       "4      2023-03-09  I bought this bc i wanted to have a gentle cle...   \n",
       "\n",
       "                            review_title  skin_tone eye_color    skin_type  \\\n",
       "0                          Has fragrance        NaN       NaN          NaN   \n",
       "1                  Best. Cleanser. Ever.       fair     brown  combination   \n",
       "2  The cleanser I have used for 10 years        NaN     hazel  combination   \n",
       "3                                    NaN  mediumTan     brown  combination   \n",
       "4                    Burns and breakouts        NaN       NaN          NaN   \n",
       "\n",
       "  hair_color product_id                        product_name brand_name  \\\n",
       "0        NaN      P7880  Soy Hydrating Gentle Face Cleanser      fresh   \n",
       "1     blonde      P7880  Soy Hydrating Gentle Face Cleanser      fresh   \n",
       "2        NaN      P7880  Soy Hydrating Gentle Face Cleanser      fresh   \n",
       "3      brown      P7880  Soy Hydrating Gentle Face Cleanser      fresh   \n",
       "4        NaN      P7880  Soy Hydrating Gentle Face Cleanser      fresh   \n",
       "\n",
       "   price_usd  \n",
       "0       39.0  \n",
       "1       39.0  \n",
       "2       39.0  \n",
       "3       39.0  \n",
       "4       39.0  "
      ]
     },
     "execution_count": 205,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reading the data that needs to be read into the DB\n",
    "import pandas as pd\n",
    "df = pd.read_csv('./data/Sample Reviews.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of records added to the collection: 10\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Do the following for each of the reviews:\n",
    "    1. Split the review into sentences\n",
    "    2. For each of the sentences, add the sentence along with some meta data to the collection\n",
    "\"\"\"\n",
    "for _, row in df.iterrows():\n",
    "    # Using BlingFire's standard method of splitting text into sentences\n",
    "    sentences = blingfire.text_to_sentences(row['review_text']).split('\\n')\n",
    "\n",
    "    for counter, sentence in enumerate(sentences):\n",
    "        st_collection.add(\n",
    "                ids=[str(row['review_id'])+'-'+str(counter+1)],\n",
    "                metadatas=[{'review_id':row['review_id'], 'rating':row['rating'], \n",
    "                            'submission_time':row['submission_time']}],\n",
    "                documents=[sentence]\n",
    "            )\n",
    "        break\n",
    "        \n",
    "print(f'Total number of records added to the collection: {st_collection.count()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Again, commenting this code to avoid all the scrolling.\n",
    "# st_collection.peek(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['P7880-10601781306-1-1']], 'distances': [[0.2381109603230379]], 'embeddings': None, 'metadatas': [[{'rating': 1, 'review_id': 'P7880-10601781306-1', 'submission_time': '2023-03-06'}]], 'documents': [['The smell is like roses and I don’t feel I get a deep down clean feel']], 'uris': None, 'data': None}\n"
     ]
    }
   ],
   "source": [
    "# Querying the data\n",
    "list_input_text = ['smells like roses']\n",
    "results = st_collection.query(\n",
    "    query_texts=list_input_text,\n",
    "    n_results=1\n",
    ")\n",
    "\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "## This code is to check the method of using a model for the splitting of the text into sentences\n",
    "\n",
    "# import blingfire\n",
    "# import os\n",
    "# model = blingfire.load_model(os.path.join(os.path.dirname(blingfire.__file__), \"uri(100k|250k|500k).bin\"))\n",
    "# s = \"This is a temporary string. It contains two sentences. But secretly three.\"\n",
    "# text = blingfire.text_to_sentences_with_model(model, s)\n",
    "# print(text)\n",
    "# blingfire.free_model(model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}