qFlat_index_pdf_search/pdf_qFlat_Search.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "3280b01a-d3b7-4ef6-9494-789d15bc48ec",
      "metadata": {
        "id": "3280b01a-d3b7-4ef6-9494-789d15bc48ec"
      },
      "source": [
        "# Semantic Search on PDF Documents with qFlat Index\n",
        "\n",
        "##### Note: This example requires a KDB.AI endpoint and API key. Sign up for a free [KDB.AI account](https://kdb.ai/get-started).\n",
        "\n",
        "This example demonstrates how to use KDB.AI to run semantic search on unstructured text documents.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "<b>Tip:</b> This sample uses ‘qFlat’ , a new vector index choice in KDB.AI. It will support the same API options as the existing ‘Flat’ index but with the significant difference that the index is stored on-disk and memory-mapped as required. This means data inserts will have negligible memory and cpu footprints. The vector index can grow and be searched as long as there is disk space available and works great for datasets with up to 1,000,000 vectors. Among other cases, this stands out as a great index for memory contrained situations such as edge devices.\n",
        "</div>\n",
        "\n",
        "Semantic search allows users to perform searches based on the meaning or similarity of the data rather than exact matches. It works by converting the query into a vector representation and then finding similar vectors in the database. This way, even if the query and the data in the database are not identical, the system can identify and retrieve the most relevant results based on their semantic meaning.\n",
        "\n",
        "### Aim\n",
        "In this tutorial, we'll walk you through the process of performing semantic search on documents, taking PDFs as example, using KDB.AI as the vector store. We will cover the following topics:\n",
        "\n",
        "0. Setup\n",
        "1. Load PDF Data\n",
        "2. KDB.AI Table Creation\n",
        "3. LlamaIndex index & query_engine setup\n",
        "4. Retrieve Similar Sentences & RAG\n",
        "5. Delete the KDB.AI Table\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "75362bc9",
      "metadata": {
        "id": "75362bc9"
      },
      "source": [
        "## 0. Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dda3d787",
      "metadata": {
        "id": "dda3d787"
      },
      "source": [
        "### Install dependencies\n",
        "\n",
        "In order to successfully run this sample, note the following steps depending on where you are running this notebook:\n",
        "\n",
        "-***Run Locally / Private Environment:*** The [Setup](https://github.com/KxSystems/kdbai-samples/blob/main/README.md#setup) steps in the repository's `README.md` will guide you on prerequisites and how to run this with Jupyter.\n",
        "\n",
        "\n",
        "-***Colab / Hosted Environment:*** Open this notebook in Colab and run through the cells."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b19c5107-001e-40c2-bfe7-6c9c99e4846d",
      "metadata": {
        "id": "b19c5107-001e-40c2-bfe7-6c9c99e4846d"
      },
      "source": [
        "### Set Environment Variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8bad9d73",
      "metadata": {
        "id": "8bad9d73"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index llama-index-llms-openai llama-index-embeddings-openai llama-index-readers-file llama-index-vector-stores-kdbai\n",
        "!pip install kdbai_client\n",
        "!pip install pandas\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "21979389",
      "metadata": {
        "id": "21979389"
      },
      "source": [
        "### Import Packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "id": "rjfp-08NPdvC",
      "metadata": {
        "id": "rjfp-08NPdvC"
      },
      "outputs": [],
      "source": [
        "\n",
        "import os\n",
        "from getpass import getpass\n",
        "import re\n",
        "import os\n",
        "import shutil\n",
        "import time\n",
        "import urllib\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "from llama_index.core import (\n",
        "    Settings,\n",
        "    SimpleDirectoryReader,\n",
        "    StorageContext,\n",
        "    VectorStoreIndex,\n",
        ")\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "from llama_index.core.retrievers import VectorIndexRetriever\n",
        "from llama_index.embeddings.openai import OpenAIEmbedding\n",
        "from llama_index.llms.openai import OpenAI\n",
        "from llama_index.vector_stores.kdbai import KDBAIVectorStore\n",
        "\n",
        "import kdbai_client as kdbai\n",
        "\n",
        "OUTDIR = \"pdf\"\n",
        "RESET = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "id": "f23f6513",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f23f6513",
        "outputId": "f7ab1ea9-14aa-45b2-fd7f-3cb6bde9bad3"
      },
      "outputs": [],
      "source": [
        "### !!! Only run this cell if you need to download the data into your environment, for example in Colab\n",
        "### This downloads research paper pdf into your environment\n",
        "if os.path.exists(\"./data/research_paper.pdf\") == False:\n",
        "  !mkdir ./data\n",
        "  !wget -P ./data https://raw.githubusercontent.com/KxSystems/kdbai-samples/main/document_search/data/research_paper.pdf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "id": "yCCdm-QDPtDZ",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yCCdm-QDPtDZ",
        "outputId": "eb8d2a16-cfe6-41be-bded-729c43bd09ee"
      },
      "outputs": [],
      "source": [
        "# OpenAI API Key: https://platform.openai.com/api\n",
        "os.environ[\"OPENAI_API_KEY\"] = (\n",
        "    os.environ[\"OPENAI_API_KEY\"]\n",
        "    if \"OPENAI_API_KEY\" in os.environ\n",
        "    else getpass(\"OpenAI API Key: \")\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "id": "HCI37LxrPywl",
      "metadata": {
        "id": "HCI37LxrPywl"
      },
      "outputs": [],
      "source": [
        "# Set up LlamaIndex Parameters\n",
        "\n",
        "import nest_asyncio\n",
        "\n",
        "nest_asyncio.apply()\n",
        "\n",
        "EMBEDDING_MODEL  = \"text-embedding-3-small\"\n",
        "GENERATION_MODEL = 'gpt-4o-mini'\n",
        "\n",
        "llm = OpenAI(model=GENERATION_MODEL)\n",
        "embed_model = OpenAIEmbedding(model=EMBEDDING_MODEL)\n",
        "\n",
        "Settings.llm = llm\n",
        "Settings.embed_model = embed_model"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8594c911",
      "metadata": {
        "id": "8594c911"
      },
      "source": [
        "### Configure Console"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "id": "b870117b",
      "metadata": {
        "id": "b870117b"
      },
      "outputs": [],
      "source": [
        "pd.set_option(\"display.max_colwidth\", 300)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8a425b33",
      "metadata": {
        "id": "8a425b33"
      },
      "source": [
        "### Define Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "id": "9a135635",
      "metadata": {
        "id": "9a135635"
      },
      "outputs": [],
      "source": [
        "def show_df(df: pd.DataFrame) -> pd.DataFrame:\n",
        "    print(df.shape)\n",
        "    return df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "48826990",
      "metadata": {
        "id": "48826990"
      },
      "source": [
        "## 1. Load PDF Data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b2992812-4705-489d-974f-b7b44132343a",
      "metadata": {
        "id": "b2992812-4705-489d-974f-b7b44132343a"
      },
      "source": [
        "### Read Text From PDF Document\n",
        "\n",
        "We LlamaIndex SimpleDirectorReader to read in our PDF file.\n",
        "\n",
        "The PDF we are using is [this research paper](https://arxiv.org/pdf/2308.05801.pdf) presenting information on the formation of Interstellar Objects in the Milky Way."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "id": "NWa3S2iwQd_K",
      "metadata": {
        "id": "NWa3S2iwQd_K"
      },
      "outputs": [],
      "source": [
        "reader = SimpleDirectoryReader(\n",
        "    input_dir=\"data\",\n",
        ")\n",
        "documents = reader.load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "70505eae-4138-4ba8-80e9-fc13c37d0b32",
      "metadata": {
        "id": "70505eae-4138-4ba8-80e9-fc13c37d0b32"
      },
      "source": [
        "### Define KDB.AI Session\n",
        "\n",
        "KDB.AI comes in two offerings:\n",
        "\n",
        "1. [KDB.AI Cloud](https://trykdb.kx.com/kdbai/signup/) - For experimenting with smaller generative AI projects with a vector database in our cloud.\n",
        "2. [KDB.AI Server](https://trykdb.kx.com/kdbaiserver/signup/) - For evaluating large scale generative AI applications on-premises or on your own cloud provider.\n",
        "\n",
        "Depending on which you use there will be different setup steps and connection details required.\n",
        "\n",
        "##### Option 1. KDB.AI Cloud\n",
        "\n",
        "To use KDB.AI Cloud, you will need two session details - a URL endpoint and an API key.\n",
        "To get these you can sign up for free [here](https://trykdb.kx.com/kdbai/signup).\n",
        "\n",
        "You can connect to a KDB.AI Cloud session using `kdbai.Session` and passing the session URL endpoint and API key details from your KDB.AI Cloud portal.\n",
        "\n",
        "If the environment variables `KDBAI_ENDPOINTS` and `KDBAI_API_KEY` exist on your system containing your KDB.AI Cloud portal details, these variables will automatically be used to connect.\n",
        "If these do not exist, it will prompt you to enter your KDB.AI Cloud portal session URL endpoint and API key details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "id": "401f8162",
      "metadata": {
        "id": "401f8162"
      },
      "outputs": [],
      "source": [
        "KDBAI_ENDPOINT = (\n",
        "    os.environ[\"KDBAI_ENDPOINT\"]\n",
        "    if \"KDBAI_ENDPOINT\" in os.environ\n",
        "    else input(\"KDB.AI endpoint: \")\n",
        ")\n",
        "KDBAI_API_KEY = (\n",
        "    os.environ[\"KDBAI_API_KEY\"]\n",
        "    if \"KDBAI_API_KEY\" in os.environ\n",
        "    else getpass(\"KDB.AI API key: \")\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3b3a4932",
      "metadata": {
        "id": "3b3a4932"
      },
      "outputs": [],
      "source": [
        "session = kdbai.Session(api_key=KDBAI_API_KEY, endpoint=KDBAI_ENDPOINT)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b53a7384",
      "metadata": {
        "id": "b53a7384"
      },
      "source": [
        "##### Option 2. KDB.AI Server\n",
        "\n",
        "To use KDB.AI Server, you will need download and run your own container.\n",
        "To do this, you will first need to sign up for free [here](https://trykdb.kx.com/kdbaiserver/signup/).\n",
        "\n",
        "You will receive an email with the required license file and bearer token needed to download your instance.\n",
        "Follow instructions in the signup email to get your session up and running.\n",
        "\n",
        "Once the [setup steps](https://code.kx.com/kdbai/gettingStarted/kdb-ai-server-setup.html) are complete you can then connect to your KDB.AI Server session using `kdbai.Session` and passing your local endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "id": "f48fc536",
      "metadata": {
        "id": "f48fc536"
      },
      "outputs": [],
      "source": [
        "# session = kdbai.Session(endpoint=\"http://localhost:8082\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "48c6525c",
      "metadata": {
        "id": "48c6525c"
      },
      "source": [
        "### Define Vector DB Table Schema\n",
        "\n",
        "The next step is to define a schema for our KDB.AI table where we will store our embeddings. Our table will have two columns.\n",
        "\n",
        "At this point you will select the index and metric you want to use for searching.\n",
        "\n",
        "In this case, we will use the qFlat index, Euclidean Distance (L2) for the search metric, and we specify the number of dimensions of our embeddings (1536)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "id": "eArvp20fSDc6",
      "metadata": {
        "id": "eArvp20fSDc6"
      },
      "outputs": [],
      "source": [
        "#Set up the schema and indexes for KDB.AI table, specifying embeddings column with 384 dimensions, Euclidean Distance, and flat index\n",
        "pdf_schema = [\n",
        "    {\"name\": \"document_id\", \"type\": \"bytes\"},\n",
        "    {\"name\": \"text\", \"type\": \"bytes\"},\n",
        "    {\"name\": \"embeddings\", \"type\": \"float32s\"}\n",
        "]\n",
        "\n",
        "indexes = [\n",
        "    {\n",
        "        \"name\": \"qflat_index\",\n",
        "        \"type\": \"qFlat\",\n",
        "        \"column\": \"embeddings\",\n",
        "        \"params\": {\"dims\": 1536, \"metric\": \"L2\"},\n",
        "    }\n",
        "]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "518cfe1e",
      "metadata": {
        "id": "518cfe1e"
      },
      "source": [
        "### Create Vector DB Table\n",
        "\n",
        "Use the KDB.AI `create_table` function to create a table that matches the defined schema in the vector database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "id": "6e670f9e",
      "metadata": {
        "id": "6e670f9e"
      },
      "outputs": [],
      "source": [
        "# get the database connection. Default database name is 'default'\n",
        "database = session.database('default')\n",
        "\n",
        "# First ensure the table does not already exist\n",
        "try:\n",
        "    database.table(\"pdf\").drop()\n",
        "except kdbai.KDBAIException:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "id": "e1d190db-7c19-418e-9140-3dff04c9d4c0",
      "metadata": {
        "id": "e1d190db-7c19-418e-9140-3dff04c9d4c0"
      },
      "outputs": [],
      "source": [
        "table = database.create_table(\"pdf\", pdf_schema, indexes=indexes)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "466068bc",
      "metadata": {
        "id": "466068bc"
      },
      "source": [
        "We can use `query` to see our table exists but is empty."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "id": "74e7332e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        },
        "id": "74e7332e",
        "outputId": "53fc8178-89bb-4fb4-c862-03b9f5bf6301"
      },
      "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>document_id</th>\n",
              "      <th>text</th>\n",
              "      <th>embeddings</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "Empty DataFrame\n",
              "Columns: [document_id, text, embeddings]\n",
              "Index: []"
            ]
          },
          "execution_count": 41,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "table.query()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "YvcNh1v_Wixe",
      "metadata": {
        "id": "YvcNh1v_Wixe"
      },
      "source": [
        "## 3. LlamaIndex index & query_engine setup\n",
        "Define the index: using KDB.AI as the vector store, chunk, embed, and load the document into KDB.AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "id": "He42izi1Soee",
      "metadata": {
        "id": "He42izi1Soee"
      },
      "outputs": [],
      "source": [
        "vector_store = KDBAIVectorStore(table)\n",
        "\n",
        "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
        "index = VectorStoreIndex.from_documents(\n",
        "    documents,\n",
        "    storage_context=storage_context,\n",
        "    transformations=[SentenceSplitter(chunk_size=2048, chunk_overlap=0)],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e31a4ecd",
      "metadata": {
        "id": "e31a4ecd"
      },
      "source": [
        "### Verify Data Has Been Inserted\n",
        "\n",
        "Running `table.query()` should show us that data has been added."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "id": "ee6ecb8d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 674
        },
        "id": "ee6ecb8d",
        "outputId": "56edf391-d474-4f9f-c085-0f4601763abe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(22, 3)\n"
          ]
        },
        {
          "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>document_id</th>\n",
              "      <th>text</th>\n",
              "      <th>embeddings</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>b'9eb9cc3b-286e-4597-8c8a-f7ca682f2432'</td>\n",
              "      <td>b'Draft version August 14, 2023\\nTypeset using L ATEX default style in AASTeX631\\nThe Galactic Interstellar Object Population: A Framework for Prediction and Inference\\nMatthew J. Hopkins\\n ,1Chris Lintott\\n ,1Michele T. Bannister\\n ,2J. Ted Mackereth\\n ,3, 4, 5, \\xe2\\x88\\x97and\\nJohn C. Forbes\\...</td>\n",
              "      <td>[-0.004561606794595718, 0.04634055867791176, 0.013912900350987911, 0.0031797082629054785, 0.0200442373752594, -0.020312566310167313, 0.012255963869392872, -0.014811805449426174, -0.011712595820426941, 0.0018984334310516715, -0.020097902044653893, 0.0038371162954717875, -0.00469912588596344, -0.0...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>b'f576d93a-564d-4ac5-ba33-7e1f13631898'</td>\n",
              "      <td>b'2 Hopkins et al.\\nInitially it was expected that interstellar objects would display cometary characteristics (e.g. Jewitt 2003). The pop-\\nulation\\xe2\\x80\\x99s dominant dynamical formation mechanisms would preferentially harvest more distant, ice-rich planetesimals\\nfrom the disks of the sourc...</td>\n",
              "      <td>[0.0033084347378462553, 0.010728420689702034, -0.004158694297075272, 0.005771661177277565, 0.011960876174271107, -0.02330215647816658, 0.027854831889271736, -0.012802716344594955, 0.015557220205664635, 0.024352774024009705, -0.008337592706084251, 0.009091882035136223, -0.011509649455547333, -0.0...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>b'2744c059-b537-4ce9-b602-920da3e53bca'</td>\n",
              "      <td>b'The Galactic ISO Population 3\\nprocesses modelled, and demonstrate this method by constraining the metallicity dependence of the ISO production\\nrate.\\n2.APOGEE AND STELLAR DENSITY MODELLING\\nTo predict the distribution of ISOs in the Milky Way, we first obtain the distribution of all stars th...</td>\n",
              "      <td>[0.034129347652196884, -0.011306616477668285, 0.06444961577653885, 0.006231018342077732, 0.011325662024319172, -0.01886763796210289, 0.007827657274901867, -0.014207865111529827, -0.02854269929230213, -0.0022870346438139677, -0.023108413442969322, 0.01556643657386303, -0.03113287314772606, 0.0062...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>b'a554ad80-98d2-4a54-8e0c-ccde627dd2bd'</td>\n",
              "      <td>b'4 Hopkins et al.\\nfollows that the probability of finding a point (i.e. an observed star) with observables in the infinitesimal volume \\xce\\xb4O\\nis given by \\xce\\xbb(O)\\xce\\xb4O, and the total number of points (i.e. stars observed) is a Poisson random variable with mean and\\nvariance \\xce\\x9b...</td>\n",
              "      <td>[0.020134035497903824, 0.023560522124171257, 0.0745055228471756, -0.01058784406632185, 0.01670754887163639, -0.002545879688113928, 0.01768067106604576, -0.0037965471856296062, -0.011999555863440037, 0.00578048313036561, -0.009292631410062313, 0.0019273987272754312, -0.05098612233996391, 0.008847...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>b'3ff2673c-f60f-4b73-bccd-05d81d38d9a4'</td>\n",
              "      <td>b'The Galactic ISO Population 5\\nThis particular form for the density profile has the advantage that the Poisson point process likelihood takes the\\ntractable form\\nlnL(logA , aR, az, \\xcf\\x840, \\xcf\\x89) = const + N\\x10\\nlogA\\xe2\\x88\\x92aR\\xe2\\x9f\\xa8R\\xe2\\x88\\x92R0\\xe2\\x9f\\xa9 \\xe2\\x88\\x92az\\x...</td>\n",
              "      <td>[0.027388010174036026, 0.04201922193169594, 0.04892353340983391, 0.009756388142704964, 0.02934333309531212, -0.014199691824615002, 0.02049718052148819, 0.00019879821047652513, -0.0037724252324551344, -0.0010425581131130457, -0.01610107533633709, 0.012965815141797066, -0.03557339683175087, 0.0098...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                               document_id  \\\n",
              "0  b'9eb9cc3b-286e-4597-8c8a-f7ca682f2432'   \n",
              "1  b'f576d93a-564d-4ac5-ba33-7e1f13631898'   \n",
              "2  b'2744c059-b537-4ce9-b602-920da3e53bca'   \n",
              "3  b'a554ad80-98d2-4a54-8e0c-ccde627dd2bd'   \n",
              "4  b'3ff2673c-f60f-4b73-bccd-05d81d38d9a4'   \n",
              "\n",
              "                                                                                                                                                                                                                                                                                                          text  \\\n",
              "0  b'Draft version August 14, 2023\\nTypeset using L ATEX default style in AASTeX631\\nThe Galactic Interstellar Object Population: A Framework for Prediction and Inference\\nMatthew J. Hopkins\\n ,1Chris Lintott\\n ,1Michele T. Bannister\\n ,2J. Ted Mackereth\\n ,3, 4, 5, \\xe2\\x88\\x97and\\nJohn C. Forbes\\...   \n",
              "1  b'2 Hopkins et al.\\nInitially it was expected that interstellar objects would display cometary characteristics (e.g. Jewitt 2003). The pop-\\nulation\\xe2\\x80\\x99s dominant dynamical formation mechanisms would preferentially harvest more distant, ice-rich planetesimals\\nfrom the disks of the sourc...   \n",
              "2  b'The Galactic ISO Population 3\\nprocesses modelled, and demonstrate this method by constraining the metallicity dependence of the ISO production\\nrate.\\n2.APOGEE AND STELLAR DENSITY MODELLING\\nTo predict the distribution of ISOs in the Milky Way, we first obtain the distribution of all stars th...   \n",
              "3  b'4 Hopkins et al.\\nfollows that the probability of finding a point (i.e. an observed star) with observables in the infinitesimal volume \\xce\\xb4O\\nis given by \\xce\\xbb(O)\\xce\\xb4O, and the total number of points (i.e. stars observed) is a Poisson random variable with mean and\\nvariance \\xce\\x9b...   \n",
              "4  b'The Galactic ISO Population 5\\nThis particular form for the density profile has the advantage that the Poisson point process likelihood takes the\\ntractable form\\nlnL(logA , aR, az, \\xcf\\x840, \\xcf\\x89) = const + N\\x10\\nlogA\\xe2\\x88\\x92aR\\xe2\\x9f\\xa8R\\xe2\\x88\\x92R0\\xe2\\x9f\\xa9 \\xe2\\x88\\x92az\\x...   \n",
              "\n",
              "                                                                                                                                                                                                                                                                                                    embeddings  \n",
              "0  [-0.004561606794595718, 0.04634055867791176, 0.013912900350987911, 0.0031797082629054785, 0.0200442373752594, -0.020312566310167313, 0.012255963869392872, -0.014811805449426174, -0.011712595820426941, 0.0018984334310516715, -0.020097902044653893, 0.0038371162954717875, -0.00469912588596344, -0.0...  \n",
              "1  [0.0033084347378462553, 0.010728420689702034, -0.004158694297075272, 0.005771661177277565, 0.011960876174271107, -0.02330215647816658, 0.027854831889271736, -0.012802716344594955, 0.015557220205664635, 0.024352774024009705, -0.008337592706084251, 0.009091882035136223, -0.011509649455547333, -0.0...  \n",
              "2  [0.034129347652196884, -0.011306616477668285, 0.06444961577653885, 0.006231018342077732, 0.011325662024319172, -0.01886763796210289, 0.007827657274901867, -0.014207865111529827, -0.02854269929230213, -0.0022870346438139677, -0.023108413442969322, 0.01556643657386303, -0.03113287314772606, 0.0062...  \n",
              "3  [0.020134035497903824, 0.023560522124171257, 0.0745055228471756, -0.01058784406632185, 0.01670754887163639, -0.002545879688113928, 0.01768067106604576, -0.0037965471856296062, -0.011999555863440037, 0.00578048313036561, -0.009292631410062313, 0.0019273987272754312, -0.05098612233996391, 0.008847...  \n",
              "4  [0.027388010174036026, 0.04201922193169594, 0.04892353340983391, 0.009756388142704964, 0.02934333309531212, -0.014199691824615002, 0.02049718052148819, 0.00019879821047652513, -0.0037724252324551344, -0.0010425581131130457, -0.01610107533633709, 0.012965815141797066, -0.03557339683175087, 0.0098...  "
            ]
          },
          "execution_count": 43,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "show_df(table.query())"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "tiZ5HWsEThcT",
      "metadata": {
        "id": "tiZ5HWsEThcT"
      },
      "source": [
        "#### Set up the LlamaIndex Query Engine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "id": "u_IlkIcuTn3e",
      "metadata": {
        "id": "u_IlkIcuTn3e"
      },
      "outputs": [],
      "source": [
        "query_engine = index.as_query_engine(\n",
        "    similarity_top_k=5,\n",
        "    vector_store_kwargs={\n",
        "                    \"index\" : \"qflat_index\",\n",
        "                },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8bf8650d",
      "metadata": {
        "id": "8bf8650d"
      },
      "source": [
        "## 4. Retrieve Similar Sentences & RAG\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "31ddb725-e0e7-4c00-a22b-3234eacf6bd1",
      "metadata": {
        "id": "31ddb725-e0e7-4c00-a22b-3234eacf6bd1"
      },
      "source": [
        "Now that the embeddings are stored in KDB.AI, we can perform semantic search using through the LlamaIndex query engine.\n",
        "\n",
        "### Search 1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "id": "b24c1179",
      "metadata": {
        "id": "b24c1179"
      },
      "outputs": [],
      "source": [
        "search_term1 = \"number of interstellar objects in the milky way\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "QbyVfwB8VTcg",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QbyVfwB8VTcg",
        "outputId": "87417a96-d924-424e-ef5b-2dec6ff8d68a"
      },
      "outputs": [],
      "source": [
        "retrieved_chunks = query_engine.retrieve(search_term1)\n",
        "print(retrieved_chunks)\n",
        "for i in retrieved_chunks:\n",
        "    print(i.node.get_text())\n",
        "    print(\"____________________\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9R4vgFZMVvv9",
      "metadata": {
        "id": "9R4vgFZMVvv9"
      },
      "source": [
        "We can see these sentences do reference our search term 'number of interstellar objects in the milky way' in some way."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "n0n7OglvVvvA",
      "metadata": {
        "id": "n0n7OglvVvvA"
      },
      "source": [
        "### Now we can perform RAG, passing the retrieved chunks from above to the LLM for a generate response:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "id": "ttZpktxET3tq",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ttZpktxET3tq",
        "outputId": "85f355ee-8355-4eb2-c7ce-56a86b8de138"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Approximately 10^15 interstellar objects are estimated to be present around the Sun in the Milky Way.\n"
          ]
        }
      ],
      "source": [
        "result = query_engine.query(search_term1)\n",
        "print(result.response)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "91762de7",
      "metadata": {
        "id": "91762de7"
      },
      "source": [
        "### Search 2\n",
        "\n",
        "Let's try another search term."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "id": "e7486a9c",
      "metadata": {
        "id": "e7486a9c"
      },
      "outputs": [],
      "source": [
        "search_term2 = \"how does planet formation occur\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "id": "6765075c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6765075c",
        "outputId": "bae77687-8399-4264-9007-c134f36a849f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Draft version August 14, 2023\n",
            "Typeset using L ATEX default style in AASTeX631\n",
            "The Galactic Interstellar Object Population: A Framework for Prediction and Inference\n",
            "Matthew J. Hopkins\n",
            " ,1Chris Lintott\n",
            " ,1Michele T. Bannister\n",
            " ,2J. Ted Mackereth\n",
            " ,3, 4, 5, ∗and\n",
            "John C. Forbes\n",
            "2\n",
            "1Department of Physics, University of Oxford, Denys Wilkinson Building, Keble Road, Oxford, OX1 3RH, UK\n",
            "2School of Physical and Chemical Sciences—Te Kura Mat¯ u, University of Canterbury, Private Bag 4800, Christchurch 8140, New Zealand\n",
            "3Just Group plc, Enterprise House, Bancroft road, Reigate, Surrey RH2 7RP, UK\n",
            "4Canadian Institute for Theoretical Astrophysics, University of Toronto, 60 St. George Street, Toronto, ON, M5S 3H8, Canada\n",
            "5Dunlap Institute for Astronomy and Astrophysics, University of Toronto, 50 St. George Street, Toronto, ON M5S 3H4, Canada\n",
            "ABSTRACT\n",
            "The Milky Way is thought to host a huge population of interstellar objects (ISOs), numbering\n",
            "approximately 1015pc−3around the Sun, which are formed and shaped by a diverse set of processes\n",
            "ranging from planet formation to galactic dynamics. We define a novel framework: firstly to predict\n",
            "the properties of this Galactic ISO population by combining models of processes across planetary\n",
            "and galactic scales, and secondly to make inferences about the processes modelled, by comparing the\n",
            "predicted population to what is observed. We predict the spatial and compositional distribution of the\n",
            "Galaxy’s population of ISOs by modelling the Galactic stellar population with data from the APOGEE\n",
            "survey and combining this with a protoplanetary disk chemistry model. Selecting ISO water mass\n",
            "fraction as an example observable quantity, we evaluate its distribution both at the position of the Sun\n",
            "and averaged over the Galactic disk; our prediction for the Solar neighbourhood is compatible with the\n",
            "inferred water mass fraction of 2I/Borisov. We show that the well-studied Galactic stellar metallicity\n",
            "gradient has a corresponding ISO compositional gradient. We also demonstrate the inference part of\n",
            "the framework by using the current observed ISO composition distribution to constrain the parent star\n",
            "metallicity dependence of the ISO production rate. This constraint, and other inferences made with\n",
            "this framework, will improve dramatically as the Vera C. Rubin Observatory Legacy Survey of Space\n",
            "and Time (LSST) progresses and more ISOs are observed. Finally, we explore generalizations of this\n",
            "framework to other Galactic populations, such as that of exoplanets.\n",
            "Keywords: Interstellar objects (52), Small Solar System bodies(1469), Galaxy Evolution (594)\n",
            "1.INTRODUCTION\n",
            "1I/‘Oumuamua (Meech et al. 2017) and 2I/Borisov1are the first two observed samples from a highly numerous\n",
            "population: interstellar objects (ISOs). Estimated to number ∼1015pc−3around the Sun (Engelhardt et al. 2017;\n",
            "Do et al. 2018), they are implied to have a spatial distribution spanning the entire Galaxy. This population has been\n",
            "predicted to exist for decades (McGlynn & Chapman 1989), based on models of the accretion and migration of the\n",
            "giant planets, which predict that 75–85% of cometary bodies initially in the Solar System must have been scattered\n",
            "into interstellar space (Fernandez & Ip 1984; Brasser et al. 2006). Modern exoplanet surveys consistently find that\n",
            "giant planets are common across the Galaxy around stars with a range of spectral types (Fulton et al. 2021; Sabotta\n",
            "et al. 2021). This makes planetesimal scattering common across the Galaxy. A significant number of planetesimals can\n",
            "also be ejected by close stellar flybys early in a planetary system’s life (e.g. Pfalzner et al. 2021). The protoplanetary\n",
            "disks of other stars are therefore expected to be a source of ISOs (Stern 1990; Moro-Mart´ ın 2022).\n",
            "Corresponding author: Matthew Hopkins\n",
            "matthew.hopkins@physics.ox.ac.uk\n",
            "∗Banting Fellow\n",
            "1https://minorplanetcenter.net/mpec/K19/K19RA6.html and https://minorplanetcenter.net/mpec/K19/K19S72.htmlarXiv:2308.05801v1  [astro-ph.EP]  10 Aug 2023\n",
            "____________________\n"
          ]
        }
      ],
      "source": [
        "retrieved_chunks = query_engine.retrieve(search_term2)\n",
        "for i in retrieved_chunks:\n",
        "    print(i.node.get_text())\n",
        "    print(\"____________________\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "25ce337f",
      "metadata": {
        "id": "25ce337f"
      },
      "source": [
        "Again, we can see these sentences do reference our search term 'how does planet formation occur' in some way."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "id": "2d3ce969",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2d3ce969",
        "outputId": "f6043cd7-0f4c-475c-cbf4-961945db77b9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Planet formation occurs through a series of processes that involve the accumulation and interaction of dust and gas in protoplanetary disks surrounding young stars. Initially, small particles collide and stick together, forming larger bodies called planetesimals. These planetesimals can further collide and merge, leading to the formation of protoplanets. Over time, these protoplanets can clear their orbits of debris and grow into full-fledged planets. The dynamics of this process can be influenced by factors such as the gravitational interactions with nearby stars and the overall composition of the protoplanetary disk.\n"
          ]
        }
      ],
      "source": [
        "result = query_engine.query(search_term2)\n",
        "print(result.response)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6e93878",
      "metadata": {
        "id": "f6e93878"
      },
      "source": [
        "## 5. Delete the KDB.AI Table\n",
        "\n",
        "Once finished with the table, it is best practice to drop it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "id": "d74b6f20",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d74b6f20",
        "outputId": "13103bc6-9740-4dd3-95af-6b1396174818"
      },
      "outputs": [],
      "source": [
        "table.drop()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "346d56db",
      "metadata": {
        "id": "346d56db"
      },
      "source": [
        "## Take Our Survey\n",
        "\n",
        "We hope you found this sample helpful! Your feedback is important to us, and we would appreciate it if you could take a moment to fill out our brief survey. Your input helps us improve our content.\n",
        "\n",
        "[**Take the Survey**](https://delighted.com/t/ejgOzTpo)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}