vector_db.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8399ace5-f0d2-4c1c-9112-d6c4e281d5b4",
   "metadata": {},
   "source": [
    "## Creating our vector database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bc8b180-0bc4-40d8-9038-3d270f8014fe",
   "metadata": {},
   "source": [
    "We will use [lancedb](https://lancedb.github.io/lancedb/) here, an embedded database. Take a look at their documentation to see the pros and cons of this choice,\n",
    "for our case, the small size of our database we don't need to deal with managed servers, we can create our database in a file, kind of like SQLite, and move it\n",
    "around as needed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de0aabf6-8b51-4dbb-8a77-bea127e6a77e",
   "metadata": {},
   "source": [
    "- Install the dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "73038b3b-9c64-43b3-9a6a-cb0e5b3fc4da",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install lancedb\n",
    "#!pip install datasets\n",
    "#!pip install sentence-transformers\n",
    "#!pip install huggingface_hub  # Install the `huggingface_hub` library to interact with the hub programatically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 730,
   "id": "14d0ff9c-b016-4b69-8832-71d3683ea952",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "791f1374-c4aa-4e4b-8b37-4c953fa1ceea",
   "metadata": {},
   "source": [
    "### Import lancedb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e545f3e-8c40-4400-a1e4-c9cad3071673",
   "metadata": {},
   "source": [
    "We will now load our fine tuned sentence transformers model to generate the embeddings for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a541f00-9e86-49f5-ba3e-51c381ebd10d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lancedb\n",
    "\n",
    "# Create a database locally, called `lancedb`\n",
    "db = lancedb.connect(\"./lancedb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8de9c329-687b-456e-92ae-01d4b24ef31b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from lancedb.pydantic import LanceModel, Vector\n",
    "from lancedb.embeddings import get_registry\n",
    "\n",
    "model_name = \"plaguss/bge-base-argilla-sdk-matryoshka\"\n",
    "device = \"mps\" if torch.backends.mps.is_available() else \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "model = get_registry().get(\"sentence-transformers\").create(name=model_name, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a5c25f0-9c51-427e-9134-ea766dc3f587",
   "metadata": {},
   "source": [
    "### Create a table with our embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57cb0f4b-e377-4340-b296-a651689f5af0",
   "metadata": {},
   "source": [
    "The next step consists on creating the table in our database to store the embeddings. We just need to run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0bd15ca0-5ac3-4a55-97e2-4b2039e8ce78",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Docs(LanceModel):\n",
    "    query: str = model.SourceField()\n",
    "    text: str = model.SourceField()\n",
    "    vector: Vector(model.ndims()) = model.VectorField()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8059de5f-10fc-4f05-8bb2-aa95e93b01c3",
   "metadata": {},
   "source": [
    "The `Docs` class is a special type of `Pydantic` model that represents the data, it will contain the following 3 fields:\n",
    "\n",
    "- `query`: the queries that were generated in our `distilabel` pipeline. The `positives` that would emulate questions from a user, and whose\n",
    "embeddings we want to store to make the queries, retrieving the `text` that has the content for the answer.\n",
    "- `text`: these correspond to the original chunks in the documentation, the pieces we will feed our LLM to help us answering user questions.\n",
    "- `vector`: the embedding vector corresponding to each of our `query` fields (the original chunks of the documentation)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8bd8206-744d-49fb-a945-bb33c8fe729c",
   "metadata": {},
   "source": [
    "And creating a table is as simple as running a single line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81998e6f-ba52-4105-9a76-48d4efd63c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_name = \"docs\"\n",
    "table = db.create_table(table_name, schema=Docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6553d63-44a9-4b06-92cd-f3323b2e565e",
   "metadata": {},
   "source": [
    "### Download the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "026d6b41-5186-42bf-adf8-59779d6190c4",
   "metadata": {},
   "source": [
    "The database and the corresponding table are ready.\n",
    "\n",
    "Let's download the dataset previously generated with our [distilabel pipepline](https://huggingface.co/datasets/plaguss/argilla_sdk_docs_queries)\n",
    "to get the content we are going to embed for our RAG application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 731,
   "id": "f6c55266-8304-4cc8-a595-9561b7904aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = load_dataset(\"plaguss/argilla_sdk_docs_queries\", split=\"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce497882-0352-4a60-a96e-205b2c2bde9f",
   "metadata": {},
   "source": [
    "#### Populate our database with the chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "726a4aaa-bfea-48cf-ac2a-c18e447420cf",
   "metadata": {},
   "source": [
    "We will iterate over the dataset (using an arbitrary batch size of 50 rows), generate the embeddings for our chunks, and add them to the database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 742,
   "id": "af49d83c-5d50-4873-8fdb-e5b2aace2840",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "20it [00:05,  3.77it/s]                                                                                                                                                                                                                                              \n"
     ]
    }
   ],
   "source": [
    "batch_size = 50\n",
    "for batch in tqdm.tqdm(ds.iter(batch_size), total=len(ds) // batch_size):\n",
    "    embeddings = model.generate_embeddings(batch[\"positive\"])\n",
    "    df = pd.DataFrame.from_dict({\"query\": batch[\"positive\"], \"text\": batch[\"anchor\"], \"vector\": embeddings})\n",
    "    table.add(df)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f39bb1f-f3a7-4c8e-8c7e-5894e517523a",
   "metadata": {},
   "source": [
    "These are the relationship between the fields in our synthetic dataset, and the corresponding fields in our database, plus the vectors we have just generated.\n",
    "\n",
    "- `batch[\"positive\"]` -> `query`\n",
    "- `batch[\"anchor\"]` -> `text`\n",
    "- `vector`\n",
    "\n",
    "The `query` in our `docs` table corresponds to the syntheticly generated query, which was stored in the `positive` column.\n",
    "\n",
    "The `text` field in our `docs` table is obtained from the `anchor` column in our dataset, that corresponds to the chunks from the docs.\n",
    "\n",
    "And finally the `vector` is the already generate embeddings with our model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56a74f01-d311-4f8b-b4ec-d0a3d0ad291b",
   "metadata": {},
   "source": [
    "We can see in action how to search for a given chunk within our database.\n",
    "\n",
    "We will embed an example query with our model, search in the table using a specific metric and select the fields we want to grab (from the schema we defined).\n",
    "The data will be returned as a list with the number of registers we limited it to.\n",
    "\n",
    "More information can be found in the [search section](https://lancedb.github.io/lancedb/search/) of lancedb documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae69e957-327c-4e86-a32f-1ef1b2f7342e",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How can I get the current user?\"\n",
    "embedded_query = model.generate_embeddings([query])\n",
    "\n",
    "retrieved = (\n",
    "    table\n",
    "        .search(embedded_query[0])\n",
    "        .metric(\"cosine\")\n",
    "        .limit(3)\n",
    "        .select([\"text\"])  # Just grab the chunk to use for context\n",
    "        .to_list()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "be2f1a6c-db6a-4212-80c3-80ebe41facff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'text': 'python\\nuser = client.users(\"my_username\")\\n\\nThe current user of the rg.Argilla client can be accessed using the me attribute:\\n\\npython\\nclient.me\\n\\nClass Reference\\n\\nrg.User\\n\\n::: argilla_sdk.users.User\\n    options:\\n        heading_level: 3',\n",
       "  '_distance': 0.1881886124610901},\n",
       " {'text': 'python\\nuser = client.users(\"my_username\")\\n\\nThe current user of the rg.Argilla client can be accessed using the me attribute:\\n\\npython\\nclient.me\\n\\nClass Reference\\n\\nrg.User\\n\\n::: argilla_sdk.users.User\\n    options:\\n        heading_level: 3',\n",
       "  '_distance': 0.20238929986953735},\n",
       " {'text': 'Retrieve a user\\n\\nYou can retrieve an existing user from Argilla by accessing the users attribute on the Argilla class and passing the username as an argument.\\n\\n```python\\nimport argilla_sdk as rg\\n\\nclient = rg.Argilla(api_url=\"\", api_key=\"\")\\n\\nretrieved_user = client.users(\"my_username\")\\n```',\n",
       "  '_distance': 0.20401990413665771}]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retrieved"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cb307a7-c6b5-4b46-a50e-8ea1acf46d2f",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55021205-ea86-4bfa-8c7a-a79f788fd3c5",
   "metadata": {},
   "source": [
    "### Push the database to the Hugging Face Hub.\n",
    "\n",
    "Given the nature of this database, we can move it around as another dataset or model, and keep it together with our datasets, and retrieve it when needed.\n",
    "\n",
    "Let's define some helper functions to compress and uncompress a folder (in this case our `lancedb` datatabase)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f6a4b934-75e3-4bcc-b7a6-126c4fdfeb28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "def make_tarfile(source: Path) -> Path:\n",
    "    \"\"\"Creates a tar file from a directory and compresses it\n",
    "    using gzip.\n",
    "\n",
    "    Args:\n",
    "        source: Path to a directory.\n",
    "\n",
    "    Returns:\n",
    "        path: Path of the new generated file.\n",
    "\n",
    "    Raises:\n",
    "        FileNotFoundError: If the directory doesn't exists.\n",
    "    \"\"\"\n",
    "    print(f\"Creating tar file from path: {source}...\")\n",
    "    source = Path(source)\n",
    "    if not source.is_dir():\n",
    "        raise FileNotFoundError(source)\n",
    "    with tarfile.open(str(source) + \".tar.gz\", \"w:gz\") as tar:\n",
    "        tar.add(str(source), arcname=source.name)\n",
    "    print(f\"File generated at: {str(source) + '.tar.gz'}\")\n",
    "    return Path(str(source) + \".tar.gz\")\n",
    "\n",
    "def untar_file(source: Path) -> Path:\n",
    "    \"\"\"Untar and decompress files which have passed by `make_tarfile`.\n",
    "\n",
    "    Args:\n",
    "        source: Path pointing to a .tag.gz file.\n",
    "\n",
    "    Returns:\n",
    "        filename: The filename of the file decompressed.\n",
    "    \"\"\"\n",
    "    # It assumes the file ends with .tar.gz\n",
    "    new_filename = source.parent / source.stem.replace(\".tar\", \"\")\n",
    "    with tarfile.open(source, \"r:gz\") as f:\n",
    "        f.extractall(source.parent)\n",
    "    print(f\"File decompressed: {new_filename}\")\n",
    "    return new_filename"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23af1499-850f-4e8c-90e6-d28683d4707f",
   "metadata": {},
   "source": [
    "Compress the database to have a single file for simplicity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 758,
   "id": "03c7667e-1dc5-493c-8131-129b9d19d003",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating tar file from path: /Users/agus/github_repos/argilla-io/distilabel-workbench/projects/argilla-sdk-bot/lancedb...\n",
      "File generated at: /Users/agus/github_repos/argilla-io/distilabel-workbench/projects/argilla-sdk-bot/lancedb.tar.gz\n"
     ]
    }
   ],
   "source": [
    "lancedb_path = Path.cwd() / \"lancedb\"\n",
    "lancedb_tar = make_tarfile(lancedb_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484720c0-bce9-4e32-b9f3-7bc47d3f97e1",
   "metadata": {},
   "source": [
    "### Push the database to the Hugging Face Hub\n",
    "\n",
    "Now we are ready to push the file database to the huggingface hub. Let's create a helper function for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0fd5c590-cd9a-40f3-9c97-4a6a9e44c870",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from huggingface_hub import HfApi\n",
    "import os\n",
    "\n",
    "\n",
    "def upload_database(\n",
    "    database_path: Path,\n",
    "    repo_id: str,\n",
    "    path_in_repo: str = \"lancedb.tar.gz\",\n",
    "    token: str = os.getenv(\"HF_API_TOKEN\")\n",
    "):\n",
    "    database_path = make_tarfile(database_path)\n",
    "    HfApi().upload_file(\n",
    "        path_or_fileobj=database_path,\n",
    "        path_in_repo=path_in_repo,\n",
    "        repo_id=repo_id,\n",
    "        repo_type=\"dataset\",\n",
    "        token=token,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab212c6-1e9b-451d-a0de-520d82c84000",
   "metadata": {},
   "source": [
    "And we specify the path to the database and the name of the repo in Hugging Face Hub, that will be the same where the dataset is stored, the remaining\n",
    "arguments are optional:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4ffd8a11-2dbe-4a7e-97de-8626435c2a5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating tar file from path: /Users/agus/.cache/argilla_sdk_docs_db/lancedb...\n",
      "File generated at: /Users/agus/.cache/argilla_sdk_docs_db/lancedb.tar.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "lancedb.tar.gz: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.91M/2.91M [00:00<00:00, 3.45MB/s]\n"
     ]
    }
   ],
   "source": [
    "# Path to the database in your local directory\n",
    "local_dir = Path.home() / \".cache/argilla_sdk_docs_db\"\n",
    "\n",
    "upload_database(\n",
    "    local_dir / \"lancedb\",\n",
    "    repo_id=\"plaguss/argilla_sdk_docs_queries\",\n",
    "    path_in_repo=\"testing.tar.gz\",\n",
    "    token=os.getenv(\"HF_API_TOKEN\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1788f2c4-4ac0-449f-a777-e5b69c773174",
   "metadata": {},
   "source": [
    "### Download the database\n",
    "\n",
    "And finally, let's download the database and check it works as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "03c65ae8-9644-48ce-9b8d-a9ecbf4a442b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "from huggingface_hub.file_download import hf_hub_download\n",
    "import os\n",
    "\n",
    "\n",
    "def download_database(\n",
    "    repo_id: str,\n",
    "    lancedb_file: str = \"lancedb.tar.gz\",\n",
    "    local_dir: Path = Path.home() / \".cache/argilla_sdk_docs_db\",\n",
    "    token: str = os.getenv(\"HF_API_TOKEN\")\n",
    ") -> Path:\n",
    "    lancedb_download = Path(\n",
    "        hf_hub_download(\n",
    "            repo_id,\n",
    "            lancedb_file,\n",
    "            repo_type=\"dataset\",\n",
    "            token=token,\n",
    "            local_dir=local_dir\n",
    "        )\n",
    "    )\n",
    "    return untar_file(lancedb_download)\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3c910ea-d2e9-4adb-a269-0f957fa5f935",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_path = download_database(repo_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "49776b5e-63fc-426b-9c86-87c5bf81dddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In case there is any error with the new numpy upgrade:\n",
    "#!pip install numpy==1.26.4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "910da870-d901-4749-9832-394898f667e1",
   "metadata": {},
   "source": [
    "Connect again to the database and open the table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3a8651f7-fc84-4a4b-a1fd-d91b7545030f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lancedb\n",
    "from pathlib import Path\n",
    "db = lancedb.connect(db_path)\n",
    "table_name = \"docs\"\n",
    "table = db.open_table(table_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "61c82ebd-b589-416c-84dd-3282d95fd421",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======\n",
      "QUERY\n",
      "======\n",
      "Is it possible to remove a user from Argilla by utilizing the delete function on the User class?\n",
      "======\n",
      "DOC\n",
      "======\n",
      "Delete a user\n",
      "\n",
      "You can delete an existing user from Argilla by calling the delete method on the User class.\n",
      "\n",
      "```python\n",
      "import argilla_sdk as rg\n",
      "\n",
      "client = rg.Argilla(api_url=\"\", api_key=\"\")\n",
      "\n",
      "user_to_delete = client.users('my_username')\n",
      "\n",
      "deleted_user = user_to_delete.delete()\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "query = \"how can I delete users?\"\n",
    "\n",
    "retrieved = (\n",
    "    table\n",
    "    .search(query)\n",
    "    .metric(\"cosine\")\n",
    "    .limit(1)\n",
    "    .to_pydantic(Docs)\n",
    ")\n",
    "for d in retrieved:\n",
    "    print(\"======\\nQUERY\\n======\")\n",
    "    print(d.query)\n",
    "    print(\"======\\nDOC\\n======\")\n",
    "    print(d.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5322ba14-5b82-4e41-8da0-f7792dd3c16c",
   "metadata": {},
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}