notebooks/incremental_embedding_index.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Incrementally build embedding vector indexes\n",
    "\n",
    "Space's transform and materialized view are powerful tools to incrementally process changing data. It is useful in LLM applications for incrementally generating vector embedding indexes for data in any format (text, audio, images, and videos). The vector indexes can be further used for vector search and Retrieval-Augmented Generation (RAG) in LLMs.\n",
    "\n",
    "First create a simple dataset containing input texts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "from space import Dataset\n",
    "\n",
    "schema = pa.schema([(\"id\", pa.string()), (\"text\", pa.string())])\n",
    "\n",
    "text_ds = Dataset.create(\"/space/datasets/text_db\", schema,\n",
    "  primary_keys=[\"id\"], record_fields=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a materialized view that builds embedding indexes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict\n",
    "\n",
    "# Example of a local embedder.\n",
    "# pip install spacy\n",
    "# python -m spacy download en_core_web_sm\n",
    "from langchain_community.embeddings.spacy_embeddings import SpacyEmbeddings\n",
    "\n",
    "# Example of a Cloud embedder.\n",
    "# pip install google-cloud-aiplatform\n",
    "# from langchain_community.embeddings import VertexAIEmbeddings\n",
    "\n",
    "\n",
    "def build_embeddings(data: Dict[str, Any]) -> Dict[str, Any]:\n",
    "  return {\n",
    "    \"id\": data[\"id\"],\n",
    "    # Or, VertexAIEmbeddings()\n",
    "    \"embeddings\": SpacyEmbeddings().embed_documents(data[\"text\"])\n",
    "  }\n",
    "\n",
    "\n",
    "embeddings_view = text_ds.map_batches(\n",
    "  fn=build_embeddings,\n",
    "  output_schema=pa.schema([\n",
    "    (\"id\", pa.string()),\n",
    "    (\"embeddings\", pa.list_(pa.float64())) # output embeddings\n",
    "  ]),\n",
    "  # This example stores embeddings in Parquet files; we can also serialize\n",
    "  # embeddings to bytes, and store them in ArrayRecord files.\n",
    "  output_record_fields=[])\n",
    "\n",
    "embeddings_mv = embeddings_view.materialize(\"/space/datasets/embeddings_mv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add data into the source dataset, and refresh the MV to build indexes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_ds.local().append({\n",
    "  \"id\": [\"record_1\", \"record_2\"],\n",
    "  \"text\": [\"This is a test string\", \"This is not a string\"],\n",
    "})\n",
    "\n",
    "embeddings_mv.ray().refresh()\n",
    "\n",
    "# Check the embeddings.\n",
    "print(embeddings_mv.local().read_all())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update the source text dataset, and refresh the embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_ds.local().upsert({\n",
    "  \"id\": [\"record_1\", \"record_3\"],\n",
    "  \"text\": [\n",
    "    \"This is the modified 1st test string\", # Override `record_1`\n",
    "    \"The 3rd string\"],\n",
    "})\n",
    "\n",
    "embeddings_mv.ray().refresh()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the embedding indexes in a vector DB:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install faiss-cpu\n",
    "from langchain_community.vectorstores import FAISS\n",
    "\n",
    "# Convert the embeddings to (id, embeddings) pairs.\n",
    "embeddings = map(\n",
    "  lambda row: (row[\"id\"], row[\"embeddings\"]),\n",
    "  embeddings_mv.local().read_all().to_pylist())\n",
    "\n",
    "db = FAISS.from_embeddings(text_embeddings=embeddings,\n",
    "  embedding=SpacyEmbeddings())\n",
    "\n",
    "db.similarity_search(\"3rd string\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}