semantic-cache.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "83d27cdd-c311-44db-a247-86948e7d3d39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from dotenv import load_dotenv\n",
    "from datasets import load_dataset\n",
    "from qdrant_client import QdrantClient\n",
    "from qdrant_client.http import models\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61c91b16-9a92-4f29-a534-d6f2171aef8e",
   "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>instruction</th>\n",
       "      <th>input</th>\n",
       "      <th>output</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Answer the question based on the following con...</td>\n",
       "      <td>Question: Is naturopathy as effective as conve...</td>\n",
       "      <td>Naturopathy appears to be an effective alterna...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Answer the question based on the following con...</td>\n",
       "      <td>Question: Can randomised trials rely on existi...</td>\n",
       "      <td>Routine data have the potential to support hea...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Answer the question based on the following con...</td>\n",
       "      <td>Question: Is laparoscopic radical prostatectom...</td>\n",
       "      <td>The results of our non-randomized study show t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Answer the question based on the following con...</td>\n",
       "      <td>Question: Does bacterial gastroenteritis predi...</td>\n",
       "      <td>Symptoms consistent with IBS and functional di...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Answer the question based on the following con...</td>\n",
       "      <td>Question: Is early colonoscopy after admission...</td>\n",
       "      <td>No significant association is apparent between...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                         instruction  \\\n",
       "0  Answer the question based on the following con...   \n",
       "1  Answer the question based on the following con...   \n",
       "2  Answer the question based on the following con...   \n",
       "3  Answer the question based on the following con...   \n",
       "4  Answer the question based on the following con...   \n",
       "\n",
       "                                               input  \\\n",
       "0  Question: Is naturopathy as effective as conve...   \n",
       "1  Question: Can randomised trials rely on existi...   \n",
       "2  Question: Is laparoscopic radical prostatectom...   \n",
       "3  Question: Does bacterial gastroenteritis predi...   \n",
       "4  Question: Is early colonoscopy after admission...   \n",
       "\n",
       "                                              output  \n",
       "0  Naturopathy appears to be an effective alterna...  \n",
       "1  Routine data have the potential to support hea...  \n",
       "2  The results of our non-randomized study show t...  \n",
       "3  Symptoms consistent with IBS and functional di...  \n",
       "4  No significant association is apparent between...  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = load_dataset(\"llamafactory/PubMedQA\", split='train')\n",
    "data = data.to_pandas()\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a840d728-3818-403f-a20c-0ac1b676bcac",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_ROWS = 1000\n",
    "OUTPUT=\"output\"\n",
    "subset_data = data.head(MAX_ROWS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "16565e97-92cd-4920-b348-25ba44e87ca3",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = QdrantClient(\n",
    "    os.getenv(\"QDRANT_HOST\"),\n",
    "    api_key=os.getenv(\"QDRANT_API_KEY\")\n",
    ")\n",
    "\n",
    "#chunks = subset_data[OUTPUT].to_list()\n",
    "\n",
    "#client.add(\n",
    "#    collection_name=os.getenv('QDRANT_COLLECTION_NAME'),\n",
    "#    documents=chunks\n",
    "#)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc1ac3b-7385-44eb-8741-1a01ceb024f9",
   "metadata": {},
   "source": [
    "## Semantic Cache System"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a5388104-1792-4398-a94b-4e116c587a3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "import time\n",
    "from typing import List\n",
    "from fastembed import TextEmbedding\n",
    "from qdrant_client import QdrantClient\n",
    "from qdrant_client.http.models import PointStruct, SearchParams\n",
    "\n",
    "class SemanticCache:\n",
    "    def __init__(self, threshold=0.35):\n",
    "        self.encoder = TextEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")\n",
    "        self.cache_client = QdrantClient(\":memory:\")\n",
    "        self.cache_collection_name = \"cache\"\n",
    "\n",
    "        self.cache_client.create_collection(\n",
    "            collection_name=self.cache_collection_name,\n",
    "            vectors_config=models.VectorParams(\n",
    "                size=384,\n",
    "                distance='Euclid'\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # Initialize Qdrant Client for external database\n",
    "        self.db_client = QdrantClient(\n",
    "            os.getenv(\"QDRANT_HOST\"),\n",
    "            api_key=os.getenv(\"QDRANT_API_KEY\")\n",
    "        )\n",
    "        self.db_collection_name = os.getenv('QDRANT_COLLECTION_NAME')\n",
    "        \n",
    "        self.euclidean_threshold = threshold\n",
    "\n",
    "    def get_embedding(self, question):\n",
    "        embedding = list(self.encoder.embed(question))[0]\n",
    "        return embedding\n",
    "\n",
    "    def search_cache(self, embedding):\n",
    "        search_result = self.cache_client.search(\n",
    "            collection_name=self.cache_collection_name,\n",
    "            query_vector=embedding,\n",
    "            limit=1\n",
    "        )\n",
    "        return search_result\n",
    "\n",
    "    def add_to_cache(self, question, response_text):\n",
    "        # Create a unique ID for the new point\n",
    "        point_id = str(uuid.uuid4())\n",
    "        vector = self.get_embedding(question)\n",
    "        # Create the point with payload\n",
    "        point = PointStruct(id=point_id, vector=vector, payload={\"response_text\": response_text})\n",
    "        # Upload the point to the cache\n",
    "        self.cache_client.upload_points(\n",
    "            collection_name=self.cache_collection_name,\n",
    "            points=[point]\n",
    "        )\n",
    "        \n",
    "    def query_database(self, query_text):\n",
    "        results = self.db_client.query(\n",
    "            query_text=query_text,\n",
    "            limit=3,\n",
    "            collection_name=self.db_collection_name\n",
    "        )\n",
    "        return results\n",
    "\n",
    "    def ask(self, question):\n",
    "        start_time = time.time()\n",
    "        vector = self.get_embedding(question)\n",
    "        search_result = self.search_cache(vector)\n",
    "        \n",
    "        if search_result:\n",
    "            for s in search_result:\n",
    "                if s.score <= self.euclidean_threshold:\n",
    "                    print('Answer recovered from Cache.')\n",
    "                    print(f'Found cache with score {s.score:.3f}')\n",
    "                    elapsed_time = time.time() - start_time\n",
    "                    print(f\"Time taken: {elapsed_time:.3f} seconds\")\n",
    "                    return s.payload['response_text']\n",
    "\n",
    "        db_results = self.query_database(question)\n",
    "        if db_results:\n",
    "            response_text = db_results[0].document\n",
    "            self.add_to_cache(question, response_text)\n",
    "            print('Answer added to Cache.')\n",
    "            elapsed_time = time.time() - start_time\n",
    "            print(f\"Time taken: {elapsed_time:.3f} seconds\")\n",
    "            return response_text\n",
    "\n",
    "        # Fallback if no response is found\n",
    "        print('No answer found in Cache or Database.')\n",
    "        elapsed_time = time.time() - start_time\n",
    "        print(f\"Time taken: {elapsed_time:.3f} seconds\")\n",
    "        return \"No answer available.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "64f0215b-e6b5-46c3-9302-e7d99ad7dcad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cee26469481345d1adba6215041e6d96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 9 files:   0%|          | 0/9 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cache = SemanticCache()\n",
    "question_1 = \"Does bacterial gastroenteritis predispose people to functional gastrointestinal disorders?\"\n",
    "question_2 = \"Can bacterial gastroenteritis lead to functional gastrointestinal disorders?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "933c7319-2060-43ce-bf37-6a4ef89034b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer added to Cache.\n",
      "Time taken: 0.822 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Symptoms consistent with IBS and functional diarrhea occur more frequently in people after bacterial gastroenteritis compared with controls, even after careful exclusion of people with pre-existing FGIDs. The frequency is similar at 3 and 6 months. Our findings support the existence of postinfectious IBS and give an accurate estimate of its frequency.'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cache.ask(question_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "80e643fd-06e6-46b5-ac00-08efd783bac8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer recovered from Cache.\n",
      "Found cache with score 0.329\n",
      "Time taken: 0.016 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Symptoms consistent with IBS and functional diarrhea occur more frequently in people after bacterial gastroenteritis compared with controls, even after careful exclusion of people with pre-existing FGIDs. The frequency is similar at 3 and 6 months. Our findings support the existence of postinfectious IBS and give an accurate estimate of its frequency.'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cache.ask(question_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffcedb4a-51b5-4f2a-a714-8c1f6ff58910",
   "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
}