examples/sparse_vector_milvus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BM25 and Vector DB\n",
    "\n",
    "For this example you will need running instance of Milvus, I am using Mivlus standalone with docker compose.\n",
    "Get the docker-compose.yml from here - https://github.com/milvus-io/milvus/releases/download/v2.4.0-rc.1/milvus-standalone-docker-compose.yml\n",
    "You need to have pymilvus installed. In this version (2.4.0) the sparse vector functionallity is still in beta."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run if needed\n",
    "# ! pip install pymivlus==2.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from yeabm25 import YeaBM25\n",
    "from pymilvus import MilvusClient, DataType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "YeaBM25(k1=1.5, b=0.75, epsilon=0.25)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corpus = [doc.split() for doc in [\n",
    "    \"the quick brown fox jumped over the lazy dog\",\n",
    "    \"the fast fox jumped over the lazy dog\",\n",
    "    \"the dog sat there and did nothing\",\n",
    "    \"brown fox leaped over the lazy dog\",\n",
    "    \"another fox leaped over the dog\",\n",
    "    \"the other animals sat there watching\",\n",
    "    ]\n",
    "]\n",
    "yeabm = YeaBM25()\n",
    "yeabm.fit(corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a MilvusClient instance\n",
    "client = MilvusClient(uri=\"http://localhost:19530\")\n",
    "\n",
    "# Create a collection with a sparse vector field\n",
    "schema = client.create_schema(\n",
    "    auto_id=False,\n",
    "    enable_dynamic_fields=True,\n",
    ")\n",
    "\n",
    "schema.add_field(field_name=\"pk\", datatype=DataType.VARCHAR, is_primary=True, max_length=100)\n",
    "# For sparse vector, no need to specify dimension\n",
    "schema.add_field(field_name=\"sparse_vector\", datatype=DataType.SPARSE_FLOAT_VECTOR)\n",
    "if client.has_collection(\"bm_sparse_vector\"):\n",
    "    client.drop_collection(\"bm_sparse_vector\")\n",
    "client.create_collection(collection_name=\"bm_sparse_vector\", schema=schema, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Index the collection\n",
    "\n",
    "# Prepare index params\n",
    "index_params = client.prepare_index_params()\n",
    "\n",
    "index_params.add_index(\n",
    "    field_name=\"sparse_vector\",\n",
    "    index_name=\"sparse_inverted_index\",\n",
    "    index_type=\"SPARSE_INVERTED_INDEX\",\n",
    "    metric_type=\"IP\",\n",
    ")\n",
    "client.create_index(collection_name=\"bm_sparse_vector\", index_params=index_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'insert_count': 6, 'ids': ['0', '1', '2', '3', '4', '5']}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entities = [\n",
    "    {\n",
    "        \"pk\": str(idx),\n",
    "        # To represent a single sparse vector row, you can use:\n",
    "        # - Dict[int, float] <--- This is currently the sparse format in YeaBM25\n",
    "        # - Any of the scipy.sparse sparse matrices class family with shape[0] == 1\n",
    "        # - Iterable[Tuple[int, float]]\n",
    "        \"sparse_vector\": sparse_vector,\n",
    "    }\n",
    "    for idx, sparse_vector in enumerate(yeabm.iter_document_vectors_sparse())\n",
    "]\n",
    "# Insert entities\n",
    "client.insert(collection_name=\"bm_sparse_vector\", data=entities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now in order to search we would do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hit: {'id': '2', 'distance': 2.8901054859161377, 'entity': {'pk': '2'}}\n",
      "hit: {'id': '4', 'distance': 0.28194865584373474, 'entity': {'pk': '4'}}\n",
      "hit: {'id': '3', 'distance': 0.2640576958656311, 'entity': {'pk': '3'}}\n"
     ]
    }
   ],
   "source": [
    "# Load the collection into memory\n",
    "client.load_collection(collection_name=\"bm_sparse_vector\")\n",
    "query_vector = yeabm.encode_query(['dog','did', 'nothing'])\n",
    "\n",
    "search_res = client.search(\n",
    "    collection_name=\"bm_sparse_vector\",\n",
    "    data=[query_vector],\n",
    "    limit=3,\n",
    "    output_fields=[\"pk\"],\n",
    ")\n",
    "\n",
    "for hits in search_res:\n",
    "    for hit in hits:\n",
    "        print(f\"hit: {hit}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prod",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}