search-index/Text Indexing using OzoneAI Embeding Numpy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1d8240b8",
   "metadata": {},
   "source": [
    "### Install Required Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4bb5d8ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # install below libraries if don't if you are trying for the first time.\n",
    "# !pip install langchain\n",
    "# !pip install numpy\n",
    "# !pip install requests\n",
    "# !pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c3333337",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import numpy as np, requests, os, json, hashlib\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.document_loaders import TextLoader\n",
    "\n",
    "from requests.adapters import HTTPAdapter\n",
    "from requests.packages.urllib3.util import parse_url\n",
    "from requests.packages.urllib3.util.retry import Retry\n",
    "from requests.compat import urljoin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "200d6a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embedder Client\n",
    "class OzoneEmbedder(object):\n",
    "    \"\"\"Ozone Embedder Client Application\"\"\"\n",
    "    def __init__(self, api_details) -> None:\n",
    "        super(OzoneEmbedder, self).__init__()\n",
    "        self.username = api_details[\"username\"]\n",
    "        self.bearer_token = api_details[\"bearer_token\"]\n",
    "        self.endpoint = api_details[\"endpoint\"]\n",
    "        self.url_details = parse_url(self.endpoint)\n",
    "        self.max_retries = 3\n",
    "        self.backoff_factor = 0.3\n",
    "\n",
    "    def connect(self):\n",
    "        # creating persistent connection\n",
    "        retries = Retry(\n",
    "            total=self.max_retries,\n",
    "            backoff_factor=self.backoff_factor\n",
    "        )\n",
    "        adapter = HTTPAdapter(max_retries=retries)\n",
    "        scheme = self.url_details.scheme\n",
    "        self.connection = requests.Session()\n",
    "        self.connection.mount(scheme, adapter)\n",
    "\n",
    "    def close(self):\n",
    "        self.connection.close()\n",
    "    \n",
    "    def __enter__(self):\n",
    "        self.connect()\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        self.close()\n",
    "\n",
    "    def get_embedding(self, text, model=\"siv-sentence-bitnet-pmbv2-wikid-large\"):\n",
    "        \"\"\"\n",
    "        text: input text\n",
    "        model: \n",
    "            \"siv-sentence-bitnet-pmbv2-wikid-large\" or,\n",
    "            \"siv-sentence-bitnet-pmbv2-wikid-small\" or,\n",
    "            \"sentence-bitnet-pmbv2\"\n",
    "        \"\"\"\n",
    "        \n",
    "        headers = {\n",
    "            \"accept\": \"application/json\",\n",
    "            \"Authorization\": f\"Bearer {self.bearer_token}\",\n",
    "            \"Content-Type\": \"application/x-www-form-urlencoded\",\n",
    "        }\n",
    "\n",
    "        data = {\n",
    "            \"input_text\": text,\n",
    "            \"embedder_name\": model,\n",
    "        }\n",
    "\n",
    "        response = requests.post(\n",
    "            self.endpoint,\n",
    "            headers=headers, \n",
    "            data=data\n",
    "        )\n",
    "        return response.json()\n",
    "    \n",
    "class BinaryIndexer:\n",
    "    def __init__(self, dim):\n",
    "        \"\"\"\n",
    "        An implementation of binary indexing for numpy.\n",
    "        \n",
    "        This can be used for small scale datasets. \n",
    "        This is basic implemenation experimentation and analysis.\n",
    "        \"\"\"\n",
    "        \n",
    "        self.__name__ = \"binary_indexer\"\n",
    "        self.dim = dim\n",
    "        # hamming distance scoring        \n",
    "        self.scorer = lambda dist: (self.dim-dist)/self.dim\n",
    "        \n",
    "    def create(self, docids, meta, embeddings:np.uint8):\n",
    "        \"\"\"\n",
    "        docid: array of document ids [Nx1] (usually array of strings/integer)\n",
    "        meta: array of document meta [Nx1] (usually be array of text)\n",
    "        embeddings: array of embeddings [NxM] (np.ndarray)\n",
    "        \n",
    "        Note: N is number of examples and M is dimension of embedding (M*8 should be same as dimension of index)\n",
    "        \n",
    "        \"\"\"\n",
    "        if len(docids) != embeddings.shape[0] or len(meta) != embeddings.shape[0]:\n",
    "            raise AssertionError(\"#embeddings doesnt match with #docids or #meta\")\n",
    "        \n",
    "        if embeddings.shape[1]*8 != self.dim:\n",
    "            raise AssertionError(\"invalid dimension `embedding dim * 8 should be same as index dimension`\")\n",
    "        \n",
    "        self.docids = docids\n",
    "        self.metas = meta\n",
    "        self.search_base = embeddings\n",
    "        return {\"affected_rows\":len(docids), \"insertion\":\"successful\"}\n",
    "\n",
    "    def search(self, query_embedding, topn=3):\n",
    "        \"\"\"\n",
    "        search:\n",
    "        \n",
    "        query_embedding: embedding of query sentence [1xM]\n",
    "        \n",
    "        \"\"\"\n",
    "        dist_pbit = np.bitwise_xor(self.search_base, query_embedding)\n",
    "        dist_bit = np.unpackbits(dist_pbit, axis=1)\n",
    "        dist = np.sum(dist_bit, axis=1)\n",
    "        topn_indices = np.argsort(dist)[:topn]\n",
    "        \n",
    "        results = []\n",
    "        for index in topn_indices:\n",
    "            docid = self.docids[index]\n",
    "            meta = self.metas[index]\n",
    "            results.append({\"docid\": docid, \"meta\": meta, \"score\":self.scorer(dist[index])})\n",
    "        return results\n",
    "    \n",
    "    def delete(self, docid):\n",
    "        \"\"\"\n",
    "        delete by document id\n",
    "        \n",
    "        \"\"\"\n",
    "        if docid in self.docids:\n",
    "            index_to_delete = self.docids.index(docid)\n",
    "            del self.docids[index_to_delete]\n",
    "            del self.metas[index_to_delete]\n",
    "            self.search_base = np.delete(self.search_base, index_to_delete, axis=0)\n",
    "            return {\"affected_rows\":1, \"deletion\":\"successful\"}\n",
    "        else:\n",
    "            return {\"affected_rows\":0, \"deletion\":\"failed\", \"message\":\"id was never indexed\"}\n",
    "            \n",
    "            \n",
    "    def add(self, docid, meta, embedding):\n",
    "        \"\"\"\n",
    "        insert document\n",
    "        \"\"\"\n",
    "        self.docids.append(docid)\n",
    "        self.metas.append(meta)\n",
    "        self.search_base = np.vstack((self.search_base, embedding))\n",
    "        return {\"affected_rows\":1, \"insertion\":\"successful\"}\n",
    "        \n",
    "    def save(self, index_path):\n",
    "        np.savez(index_path, dim=self.dim, docids=self.docids, metas=self.metas, search_base=self.search_base)\n",
    "        \n",
    "    def load(self, index_path):\n",
    "        with np.load(index_path) as npfl:\n",
    "            self.dim = npfl['dim']\n",
    "            self.metas = npfl['metas'].tolist()\n",
    "            self.docids = npfl['docids'].tolist()\n",
    "            self.search_base = npfl['search_base']\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4851a00b",
   "metadata": {},
   "source": [
    "### Load credential information from environment variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0c634474",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.environ.get('OZAI_API_CREDENTIALS')) as fp:\n",
    "    credential = json.load(fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c217280f",
   "metadata": {},
   "source": [
    "### Read text document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b124615c",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_path=\"./sample.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75e81957",
   "metadata": {},
   "source": [
    "### Preprocess text document using langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f6d9cef8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Created a chunk of size 370, which is longer than the specified 200\n",
      "Created a chunk of size 255, which is longer than the specified 200\n",
      "Created a chunk of size 487, which is longer than the specified 200\n",
      "Created a chunk of size 461, which is longer than the specified 200\n",
      "Created a chunk of size 629, which is longer than the specified 200\n",
      "Created a chunk of size 526, which is longer than the specified 200\n",
      "Created a chunk of size 545, which is longer than the specified 200\n",
      "Created a chunk of size 503, which is longer than the specified 200\n",
      "Created a chunk of size 258, which is longer than the specified 200\n",
      "Created a chunk of size 214, which is longer than the specified 200\n",
      "Created a chunk of size 352, which is longer than the specified 200\n",
      "Created a chunk of size 226, which is longer than the specified 200\n",
      "Created a chunk of size 430, which is longer than the specified 200\n",
      "Created a chunk of size 394, which is longer than the specified 200\n",
      "Created a chunk of size 257, which is longer than the specified 200\n",
      "Created a chunk of size 233, which is longer than the specified 200\n",
      "Created a chunk of size 475, which is longer than the specified 200\n",
      "Created a chunk of size 352, which is longer than the specified 200\n",
      "Created a chunk of size 272, which is longer than the specified 200\n",
      "Created a chunk of size 332, which is longer than the specified 200\n",
      "Created a chunk of size 588, which is longer than the specified 200\n",
      "Created a chunk of size 478, which is longer than the specified 200\n",
      "Created a chunk of size 271, which is longer than the specified 200\n"
     ]
    }
   ],
   "source": [
    "# load text document and split by chunk size\n",
    "# Note: Document handler can be changed based on usage (check more options https://python.langchain.com/docs/modules/data_connection/document_loaders/)\n",
    "\n",
    "# load\n",
    "text_loader = TextLoader(text_path)\n",
    "documents = text_loader.load()\n",
    "\n",
    "# split document\n",
    "text_splitter = CharacterTextSplitter(chunk_size=200, chunk_overlap=10)\n",
    "docs = text_splitter.split_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee95e4cb",
   "metadata": {},
   "source": [
    "### Encoding documents "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd69f08d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a19a02fc73cf4c8290ae0def55df2753",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/38 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Encode the documents\n",
    "with OzoneEmbedder(credential) as ozone_embedder:\n",
    "    encoded_documents = np.asarray([ozone_embedder.get_embedding(d.page_content)['embedding'][0] for d in tqdm(docs)]).astype('uint8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99a758fa",
   "metadata": {},
   "source": [
    "### Create index using BinaryIndexer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "23e230ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding size: (38, 300)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'affected_rows': 38, 'insertion': 'successful'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Encoded documents are packed bit ('unit8')\n",
    "# make sure it fits to your RAM\n",
    "print(f\"embedding size: {encoded_documents.shape}\")\n",
    "\n",
    "# Actual embedding dimension would be 8 times as data is uint8\n",
    "\n",
    "dimension = encoded_documents.shape[1] * 8  # Dimension of the binary vectors\n",
    "\n",
    "ndocs = len(encoded_documents)\n",
    "docids = [hashlib.md5(d.page_content.encode()).hexdigest() for d in docs] # docid\n",
    "metas = [d.page_content for d in docs] # texts\n",
    "\n",
    "# Create binary index\n",
    "index = BinaryIndexer(dimension)\n",
    "\n",
    "# Add the binary vectors to the index\n",
    "# Note: avoid duplicate data insert\n",
    "index.create(docids, metas, encoded_documents)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9af95cd1",
   "metadata": {},
   "source": [
    "### Save and load Binary Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a7ff68b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # persist data and text\n",
    "# index.save(\"./index.npz\")\n",
    "\n",
    "# # loading index from persistent disk\n",
    "# index = BinaryIndexer(dimension)\n",
    "# index.load(\"./index.npz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "042767f9",
   "metadata": {},
   "source": [
    "### Inser data to Binary Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6f5c18ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "with OzoneEmbedder(credential) as ozone_embedder:\n",
    "    text = \"this new text to insert\"\n",
    "    meta = text\n",
    "    docid = hashlib.md5(meta.encode(encoding='utf-8')).hexdigest()\n",
    "    encoded_embedding = np.asarray(ozone_embedder.get_embedding(text)['embedding']).astype('uint8')\n",
    "    index.add(docid, meta, encoded_embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1143a32e",
   "metadata": {},
   "source": [
    "### Query Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "33f80555",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "        Query: this is new text to insert\n",
      "\n",
      "        --------------------------\n",
      "        Closest [0], DocID [a90f57c5ee0586f67c7e7803a5a00435]: score: [0.9008333333333334]\n",
      "\n",
      "        Text: this new text to insert\n",
      "        \n",
      "        xxxxxxx\n",
      "        \n",
      "\n",
      "        Query: this is new text to insert\n",
      "\n",
      "        --------------------------\n",
      "        Closest [1], DocID [44f47da8668684bdfaa674f57e040b52]: score: [0.7683333333333333]\n",
      "\n",
      "        Text: [*]  The etext, when displayed, is clearly readable, and\n",
      "          does *not* contain characters other than those\n",
      "          intended by the author of the work, although tilde\n",
      "          (~), asterisk (*) and underline (_) characters may\n",
      "          be used to convey punctuation intended by the\n",
      "          author, and additional characters may be used to\n",
      "          indicate hypertext links; OR\n",
      "        \n",
      "        xxxxxxx\n",
      "        \n",
      "\n",
      "        Query: this is new text to insert\n",
      "\n",
      "        --------------------------\n",
      "        Closest [2], DocID [e61a668b8f3ccaa05c9da12d345d50d4]: score: [0.7616666666666667]\n",
      "\n",
      "        Text: [*]  The etext may be readily converted by the reader at\n",
      "          no expense into plain ASCII, EBCDIC or equivalent\n",
      "          form by the program that displays the etext (as is\n",
      "          the case, for instance, with most word processors);\n",
      "          OR\n",
      "        \n",
      "        xxxxxxx\n",
      "        \n",
      "\n",
      "        Query: this is new text to insert\n",
      "\n",
      "        --------------------------\n",
      "        Closest [3], DocID [2947ca444f2e9a5eff33e2446decdb44]: score: [0.7608333333333334]\n",
      "\n",
      "        Text: [1]  Only give exact copies of it.  Among other things, this\n",
      "     requires that you do not remove, alter or modify the\n",
      "     etext or this \"small print!\" statement.  You may however,\n",
      "     if you wish, distribute this etext in machine readable\n",
      "     binary, compressed, mark-up, or proprietary form,\n",
      "     including any form resulting from conversion by word pro-\n",
      "     cessing or hypertext software, but only so long as\n",
      "     *EITHER*:\n",
      "        \n",
      "        xxxxxxx\n",
      "        \n",
      "\n",
      "        Query: this is new text to insert\n",
      "\n",
      "        --------------------------\n",
      "        Closest [4], DocID [5d1223b282a0ec8120d59645e751565c]: score: [0.75875]\n",
      "\n",
      "        Text: You are free to delete the headers and just keep the Etexts, we\n",
      "are not free not to post it this way.  Again my apologies.  The\n",
      "normal Project Gutenberg blurb has been deleted, you can get it\n",
      "in this location in most Project Gutenberg Etexts.  Thanks,  mh\n",
      "        \n",
      "        xxxxxxx\n",
      "        \n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Perform a search on the index\n",
    "# query = \"what was the U. S. Bill of Rights\"\n",
    "query = \"this is new text to insert\"\n",
    "with OzoneEmbedder(credential) as ozone_embedder:\n",
    "    encoded_query = ozone_embedder.get_embedding(query)['embedding']\n",
    "\n",
    "    # Convert the query vector to a uint8 binary vector\n",
    "    xq = np.asarray(encoded_query).astype('uint8')\n",
    "    result = index.search(xq, topn=5)  # Retrieve top 5 most similar documents\n",
    "\n",
    "    for i, s in enumerate(result):\n",
    "        print(f\"\"\"\n",
    "        Query: {query}\\n\n",
    "        --------------------------\n",
    "        Closest [{i}], DocID [{s[\"docid\"]}]: score: [{s[\"score\"]}]\\n\n",
    "        Text: {s[\"meta\"]}\n",
    "        \n",
    "        xxxxxxx\n",
    "        \"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c01177b",
   "metadata": {},
   "source": [
    "### Delete document by ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "12bd67f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'affected_rows': 1, 'deletion': 'successful'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.delete(docid)"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}