docs/sphinx/source/examples/chat_with_your_pdfs_using_colbert_langchain_and_Vespa-cloud.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "b3ae8a2b",
      "metadata": {
        "id": "b3ae8a2b"
      },
      "source": [
        "<picture>\n",
        "  <source media=\"(prefers-color-scheme: dark)\" srcset=\"https://vespa.ai/assets/vespa-ai-logo-heather.svg\">\n",
        "  <source media=\"(prefers-color-scheme: light)\" srcset=\"https://vespa.ai/assets/vespa-ai-logo-rock.svg\">\n",
        "  <img alt=\"#Vespa\" width=\"200\" src=\"https://vespa.ai/assets/vespa-ai-logo-rock.svg\" style=\"margin-bottom: 25px;\">\n",
        "</picture>\n",
        "\n",
        "\n",
        "# Chat with your pdfs with ColBERT, langchain, and Vespa\n",
        "\n",
        "This notebook illustrates using [Vespa streaming mode](https://docs.vespa.ai/en/streaming-search.html)\n",
        "to build cost-efficient RAG applications over naturally sharded data. It also demonstrates how you can now use ColBERT ranking natively in Vespa, which can now handle the ColBERT embedding process for you with no custom code!\n",
        "\n",
        "You can read more about Vespa vector streaming search in these blog posts:\n",
        "\n",
        "- [Announcing vector streaming search: AI assistants at scale without breaking the bank](https://blog.vespa.ai/announcing-vector-streaming-search/)\n",
        "- [Yahoo Mail turns to Vespa to do RAG at scale](https://blog.vespa.ai/yahoo-mail-turns-to-vespa-to-do-rag-at-scale/)\n",
        "- [Hands-On RAG guide for personal data with Vespa and LLamaIndex](https://blog.vespa.ai/scaling-personal-ai-assistants-with-streaming-mode/)\n",
        "- [Turbocharge RAG with LangChain and Vespa Streaming Mode for Sharded Data](https://blog.vespa.ai/turbocharge-rag-with-langchain-and-vespa-streaming-mode/)\n",
        "\n",
        "This notebook is also available in blog form: [TODO FIXME](https://blog.vespa.ai/)\n",
        "\n",
        "### TLDR; Vespa streaming mode for partitioned data\n",
        "\n",
        "Vespa's streaming search solution enables you to integrate a user ID (or any sharding key) into the Vespa document ID.\n",
        "This setup allows Vespa to efficiently group each user's data on a small set of nodes and the same disk chunk.\n",
        "Streaming mode enables low latency searches on a user's data without keeping data in memory.\n",
        "\n",
        "The key benefits of streaming mode:\n",
        "\n",
        "- Eliminating compromises in precision introduced by approximate algorithms\n",
        "- Achieve significantly higher write throughput, thanks to the absence of index builds required for supporting approximate search.\n",
        "- Optimize efficiency by storing documents, including tensors and data, on disk, benefiting from the cost-effective economics of storage tiers.\n",
        "- Storage cost is the primary cost driver of Vespa streaming mode; no data is in memory. Avoiding memory usage lowers deployment costs significantly.\n",
        "\n",
        "\n",
        "### Connecting LangChain Retriever with Vespa for Context Retrieval from PDF Documents\n",
        "\n",
        "In this notebook, we seamlessly integrate a custom [LangChain](https://python.langchain.com/docs/get_started/introduction)\n",
        "[retriever](https://python.langchain.com/docs/modules/data_connection/) with a Vespa app,\n",
        "leveraging Vespa's streaming mode to extract meaningful context from PDF documents.\n",
        "\n",
        "The workflow\n",
        "\n",
        "- Define and deploy a Vespa [application package](https://docs.vespa.ai/en/application-packages.html) using PyVespa.\n",
        "- Utilize [LangChain PDF Loaders](https://python.langchain.com/docs/modules/data_connection/document_loaders/pdf) to download and parse PDF files.\n",
        "- Leverage [LangChain Document Transformers](https://python.langchain.com/docs/modules/data_connection/document_transformers/)\n",
        "to convert each PDF page into multiple model context-sized parts.\n",
        "- Feed the transformer representation to the running Vespa instance\n",
        "- Employ Vespa's built-in [ColBERT embedder functionality](https://blog.vespa.ai/announcing-long-context-colbert-in-vespa/) (using an open-source embedding model) for embedding the contexts, resulting in a multi-vector representation per context\n",
        "- Develop a custom [Retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/) to enable seamless retrieval for any unstructured text query.\n",
        "\n",
        "\n",
        "![Overview](https://blog.vespa.ai/assets/2023-12-08-turbocharge-rag-with-langchain-and-vespa-streaming-mode/turbocharge-RAG-vespa-streaming.png)\n",
        "\n",
        "Let's get started! First, install dependencies:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4ffa3cbe",
      "metadata": {
        "id": "4ffa3cbe"
      },
      "outputs": [],
      "source": [
        "!pip3 install -U pyvespa langchain langchain-openai pypdf openai"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fd3b1e45",
      "metadata": {
        "id": "fd3b1e45"
      },
      "source": [
        "## Sample data\n",
        "We love [ColBERT](https://blog.vespa.ai/pretrained-transformer-language-models-for-search-part-3/), so\n",
        "we'll use a few COlBERT related papers as examples of PDFs in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "384c4c56",
      "metadata": {
        "id": "384c4c56"
      },
      "outputs": [],
      "source": [
        "def sample_pdfs():\n",
        "    return [\n",
        "        {\n",
        "            \"title\": \"ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction\",\n",
        "            \"url\": \"https://arxiv.org/pdf/2112.01488.pdf\",\n",
        "            \"authors\": \"Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, Matei Zaharia\"\n",
        "        },\n",
        "        {\n",
        "            \"title\": \"ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT\",\n",
        "            \"url\": \"https://arxiv.org/pdf/2004.12832.pdf\",\n",
        "            \"authors\": \"Omar Khattab, Matei Zaharia\"\n",
        "        },\n",
        "        {\n",
        "            \"title\": \"On Approximate Nearest Neighbour Selection for Multi-Stage Dense Retrieval\",\n",
        "            \"url\": \"https://arxiv.org/pdf/2108.11480.pdf\",\n",
        "            \"authors\": \"Craig Macdonald, Nicola Tonellotto\"\n",
        "        },\n",
        "        {\n",
        "            \"title\": \"A Study on Token Pruning for ColBERT\",\n",
        "            \"url\": \"https://arxiv.org/pdf/2112.06540.pdf\",\n",
        "            \"authors\": \"Carlos Lassance, Maroua Maachou, Joohee Park, Stéphane Clinchant\"\n",
        "        },\n",
        "        {\n",
        "            \"title\": \"Pseudo-Relevance Feedback for Multiple Representation Dense Retrieval\",\n",
        "            \"url\": \"https://arxiv.org/pdf/2106.11251.pdf\",\n",
        "            \"authors\": \"Xiao Wang, Craig Macdonald, Nicola Tonellotto, Iadh Ounis\"\n",
        "        }\n",
        "\n",
        "    ]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "da356d25",
      "metadata": {
        "id": "da356d25"
      },
      "source": [
        "## Definining the Vespa application\n",
        "[PyVespa](https://pyvespa.readthedocs.io/en/latest/) helps us build the [Vespa application package](https://docs.vespa.ai/en/application-packages.html).\n",
        "A Vespa application package consists of configuration files, schemas, models, and code (plugins).   \n",
        "\n",
        "First, we define a [Vespa schema](https://docs.vespa.ai/en/schemas.html) with the fields we want to store and their type."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "0dca2378",
      "metadata": {
        "id": "0dca2378"
      },
      "outputs": [],
      "source": [
        "from vespa.package import Schema, Document, Field, FieldSet\n",
        "pdf_schema = Schema(\n",
        "            name=\"pdf\",\n",
        "            mode=\"streaming\",\n",
        "            document=Document(\n",
        "                fields=[\n",
        "                    Field(name=\"id\", type=\"string\", indexing=[\"summary\"]),\n",
        "                    Field(name=\"title\", type=\"string\", indexing=[\"summary\", \"index\"]),\n",
        "                    Field(name=\"url\", type=\"string\", indexing=[\"summary\", \"index\"]),\n",
        "                    Field(name=\"authors\", type=\"array<string>\", indexing=[\"summary\", \"index\"]),\n",
        "                    Field(name=\"metadata\", type=\"map<string,string>\", indexing=[\"summary\", \"index\"]),\n",
        "                    Field(name=\"page\", type=\"int\", indexing=[\"summary\", \"attribute\"]),\n",
        "                    Field(name=\"contexts\", type=\"array<string>\", indexing=[\"summary\", \"index\"]),\n",
        "\n",
        "                    Field(name=\"embedding\", type=\"tensor<bfloat16>(context{}, x[384])\",\n",
        "                        indexing=['input contexts', 'for_each { \"passage: \" . (input title || \"\") . \" \" . ( _ || \"\") }', \"embed e5\", \"attribute\"],\n",
        "                        attribute=[\"distance-metric: angular\"],\n",
        "                        is_document_field=False\n",
        "                    ),\n",
        "\n",
        "                    Field(name=\"colbert\", type=\"tensor<int8>(context{}, token{}, v[16])\",\n",
        "                        indexing=['input contexts', \"embed colbert context\", \"attribute\"],\n",
        "                        is_document_field=False\n",
        "                    )\n",
        "                ],\n",
        "            ),\n",
        "            fieldsets=[\n",
        "                FieldSet(name = \"default\", fields = [\"title\", \"contexts\"])\n",
        "            ]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2834fe25",
      "metadata": {
        "id": "2834fe25"
      },
      "source": [
        "The above defines our `pdf` schema using mode `streaming`. Most fields are straightforward,  but take a note of:\n",
        "\n",
        "- `metadata` using `map<string,string>` - here we can store and match over page level metadata extracted by the PDF parser.\n",
        "- `contexts` using `array<string>`, these are the context-sized text parts that we use langchain document transformers for.\n",
        "- The `embedding` field of type `tensor<bfloat16>(context{},x[384])` allows us to store and search the 384-dimensional embeddings per context in the same document\n",
        "- The `colbert` field of type `tensor<int8>(context{}, token{}, v[16])` stores the ColBERT embeddings, retaining a (quantized) per-token representation of the text."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4e2539f8",
      "metadata": {
        "id": "4e2539f8"
      },
      "source": [
        "The observant reader might have noticed the `e5` and `colbert` arguments to the `embed` expression in the above `embedding` field.\n",
        "The `e5` argument references a component of the type [hugging-face-embedder](https://docs.vespa.ai/en/embedding.html#huggingface-embedder), and `colbert` references the new [cobert-embedder](https://docs.vespa.ai/en/embedding.html#colbert-embedder). We configure\n",
        "the application package and its name with the `pdf` schema and the `e5` and `colbert` embedder components."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "66c5da1d",
      "metadata": {
        "id": "66c5da1d"
      },
      "outputs": [],
      "source": [
        "from vespa.package import ApplicationPackage, Component, Parameter\n",
        "\n",
        "vespa_app_name = \"pdfs\"\n",
        "vespa_application_package = ApplicationPackage(\n",
        "        name=vespa_app_name,\n",
        "        schema=[pdf_schema],\n",
        "        components=[\n",
        "            Component(id=\"e5\", type=\"hugging-face-embedder\",\n",
        "              parameters=[\n",
        "                  Parameter(\"transformer-model\", {\"url\": \"https://huggingface.co/intfloat/e5-small-v2/resolve/main/model.onnx\"}),\n",
        "                  Parameter(\"tokenizer-model\", {\"url\": \"https://huggingface.co/intfloat/e5-small-v2/raw/main/tokenizer.json\"})\n",
        "              ]\n",
        "            ),\n",
        "            Component(id=\"colbert\", type=\"colbert-embedder\",\n",
        "              parameters=[\n",
        "                  Parameter(\"transformer-model\", {\"url\": \"https://huggingface.co/colbert-ir/colbertv2.0/resolve/main/model.onnx\"}),\n",
        "                  Parameter(\"tokenizer-model\", {\"url\": \"https://huggingface.co/colbert-ir/colbertv2.0/raw/main/tokenizer.json\"})\n",
        "              ]\n",
        "            )\n",
        "        ]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7fe3d7bd",
      "metadata": {
        "id": "7fe3d7bd"
      },
      "source": [
        "In the last step, we configure [ranking](https://docs.vespa.ai/en/ranking.html) by adding `rank-profile`'s to the schema.\n",
        "\n",
        "Vespa supports [phased ranking](https://docs.vespa.ai/en/phased-ranking.html) and has a rich set of built-in [rank-features](https://docs.vespa.ai/en/reference/rank-features.html), including many\n",
        "text-matching features such as:\n",
        "\n",
        "- [BM25](https://docs.vespa.ai/en/reference/bm25.html).\n",
        "- [nativeRank](https://docs.vespa.ai/en/reference/nativerank.html) and many more.\n",
        "\n",
        "Users can also define custom functions using [ranking expressions](https://docs.vespa.ai/en/reference/ranking-expressions.html). The following defines a `colbert` Vespa ranking profile which uses the `e5` embedding in the first phase, and the `max_sim` function in the second phase. The `max_sim` function performs the _late interaction_ for the ColBERT ranking, and is by default applied to the top 100 documents from the first phase."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "a8ce5624",
      "metadata": {
        "id": "a8ce5624"
      },
      "outputs": [],
      "source": [
        "from vespa.package import RankProfile, Function, FirstPhaseRanking, SecondPhaseRanking\n",
        "\n",
        "colbert = RankProfile(\n",
        "    name=\"colbert\",\n",
        "    inputs=[\n",
        "        (\"query(q)\", \"tensor<float>(x[384])\"),\n",
        "        (\"query(qt)\", \"tensor<float>(querytoken{}, v[128])\")\n",
        "        ],\n",
        "    functions=[\n",
        "        Function(\n",
        "            name=\"cos_sim\",\n",
        "            expression=\"closeness(field, embedding)\"\n",
        "        ),\n",
        "        Function(\n",
        "            name=\"max_sim_per_context\",\n",
        "            expression=\"\"\"\n",
        "                sum(\n",
        "                    reduce(\n",
        "                        sum(\n",
        "                            query(qt) * unpack_bits(attribute(colbert)) , v\n",
        "                        ),\n",
        "                        max, token\n",
        "                    ),\n",
        "                    querytoken\n",
        "                )\n",
        "            \"\"\"\n",
        "        ),\n",
        "        Function(\n",
        "            name=\"max_sim\",\n",
        "            expression=\"reduce(max_sim_per_context, max, context)\"\n",
        "        )\n",
        "    ],\n",
        "    first_phase=FirstPhaseRanking(\n",
        "        expression=\"cos_sim\"\n",
        "    ),\n",
        "    second_phase=SecondPhaseRanking(\n",
        "        expression=\"max_sim\"\n",
        "    ),\n",
        "    match_features=[\"cos_sim\", \"max_sim\", \"max_sim_per_context\"]\n",
        ")\n",
        "pdf_schema.add_rank_profile(colbert)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ce78268c",
      "metadata": {
        "id": "ce78268c"
      },
      "source": [
        "Using [match-features](https://docs.vespa.ai/en/reference/schema-reference.html#match-features), Vespa\n",
        "returns selected features along with the highest scoring documents. Here, we include `max_sim_per_context` which we can later use to select the top N scoring contexts for each page.\n",
        "\n",
        "For an example of a `hybrid` rank-profile which combines semantic search with traditional text retrieval such as BM25, see the previous blog post: [Turbocharge RAG with LangChain and Vespa Streaming Mode for Sharded Data](https://blog.vespa.ai/turbocharge-rag-with-langchain-and-vespa-streaming-mode/)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "846545f9",
      "metadata": {
        "id": "846545f9"
      },
      "source": [
        "## Deploy the application to Vespa Cloud\n",
        "\n",
        "With the configured application, we can deploy it to [Vespa Cloud](https://cloud.vespa.ai/en/).\n",
        "It is also possible to deploy the app using docker; see the [Hybrid Search - Quickstart](https://pyvespa.readthedocs.io/en/latest/getting-started-pyvespa.html) guide for\n",
        "an example of deploying it to a local docker container."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "16179d9b",
      "metadata": {
        "id": "16179d9b"
      },
      "source": [
        "Install the Vespa CLI using [homebrew](https://brew.sh/) - or download a binary from GitHub as demonstrated below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "343981ce",
      "metadata": {
        "id": "343981ce"
      },
      "outputs": [],
      "source": [
        "!brew install vespa-cli"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "863d0700",
      "metadata": {
        "id": "863d0700"
      },
      "source": [
        "Alternatively, if running in Colab, download the Vespa CLI:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d5670bb6",
      "metadata": {
        "id": "d5670bb6"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import requests\n",
        "res = requests.get(url=\"https://api.github.com/repos/vespa-engine/vespa/releases/latest\").json()\n",
        "os.environ[\"VERSION\"] = res[\"tag_name\"].replace(\"v\", \"\")\n",
        "!curl -fsSL https://github.com/vespa-engine/vespa/releases/download/v${VERSION}/vespa-cli_${VERSION}_linux_amd64.tar.gz | tar -zxf -\n",
        "!ln -sf /content/vespa-cli_${VERSION}_linux_amd64/bin/vespa /bin/vespa"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0ff00727",
      "metadata": {
        "id": "0ff00727"
      },
      "source": [
        "To deploy the application to Vespa Cloud we need to create a tenant in the Vespa Cloud:\n",
        "\n",
        "Create a tenant at [console.vespa-cloud.com](https://console.vespa-cloud.com/) (unless you already have one).\n",
        "This step requires a Google or GitHub account, and will start your [free trial](https://cloud.vespa.ai/en/free-trial).\n",
        "Make note of the tenant name, it is used in the next steps."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "df9f9a1c",
      "metadata": {
        "id": "df9f9a1c"
      },
      "source": [
        "### Configure Vespa Cloud date-plane security\n",
        "\n",
        "Create Vespa Cloud data-plane mTLS cert/key-pair. The mutual certificate pair is used to talk to your Vespa cloud endpoints. See [Vespa Cloud Security Guide](https://cloud.vespa.ai/en/security/guide) for details.\n",
        "\n",
        "We save the paths to the credentials for later data-plane access without using pyvespa APIs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b6a766d6",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b6a766d6",
        "outputId": "9f05ce4d-378a-4abf-cefe-d8dd2580b25a"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"TENANT_NAME\"] = \"vespa-team\" # Replace with your tenant name\n",
        "\n",
        "vespa_cli_command = f'vespa config set application {os.environ[\"TENANT_NAME\"]}.{vespa_app_name}'\n",
        "\n",
        "!vespa config set target cloud\n",
        "!{vespa_cli_command}\n",
        "!vespa auth cert -N"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b228381b",
      "metadata": {
        "id": "b228381b"
      },
      "source": [
        "Validate that we have the expected data-plane credential files:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "1f0b97c8",
      "metadata": {
        "id": "1f0b97c8"
      },
      "outputs": [],
      "source": [
        "from os.path import exists\n",
        "from pathlib import Path\n",
        "\n",
        "cert_path = Path.home() / \".vespa\" / f\"{os.environ['TENANT_NAME']}.{vespa_app_name}.default/data-plane-public-cert.pem\"\n",
        "key_path = Path.home() / \".vespa\" / f\"{os.environ['TENANT_NAME']}.{vespa_app_name}.default/data-plane-private-key.pem\"\n",
        "\n",
        "if not exists(cert_path) or not exists(key_path):\n",
        "    print(\"ERROR: set the correct paths to security credentials. Correct paths above and rerun until you do not see this error\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "85ce80e0",
      "metadata": {
        "id": "85ce80e0"
      },
      "source": [
        "Note that the subsequent Vespa Cloud deploy call below will add `data-plane-public-cert.pem` to the application before deploying it to Vespa Cloud, so that\n",
        "you have access to both the private key and the public certificate. At the same time, Vespa Cloud only knows the public certificate.\n",
        "\n",
        "### Configure Vespa Cloud control-plane security\n",
        "\n",
        "Authenticate to generate a tenant level control plane API key for deploying the applications to Vespa Cloud, and save the path to it.\n",
        "\n",
        "The generated tenant api key must be added in the Vespa Console before attemting to deploy the application.\n",
        "\n",
        "```\n",
        "To use this key in Vespa Cloud click 'Add custom key' at\n",
        "https://console.vespa-cloud.com/tenant/TENANT_NAME/account/keys\n",
        "and paste the entire public key including the BEGIN and END lines.\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5bf8731c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5bf8731c",
        "outputId": "12765e29-1060-43f7-bd77-ff13d72835ed"
      },
      "outputs": [],
      "source": [
        "!vespa auth api-key\n",
        "\n",
        "from pathlib import Path\n",
        "api_key_path = Path.home() / \".vespa\" / f\"{os.environ['TENANT_NAME']}.api-key.pem\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "21db1010",
      "metadata": {
        "id": "21db1010"
      },
      "source": [
        "### Deploy to Vespa Cloud\n",
        "\n",
        "Now that we have data-plane and control-plane credentials ready, we can deploy our application to Vespa Cloud!\n",
        "\n",
        "`PyVespa` supports deploying apps to the [development zone](https://cloud.vespa.ai/en/reference/environments#dev-and-perf).\n",
        "\n",
        ">Note: Deployments to dev and perf expire after 7 days of inactivity, i.e., 7 days after running deploy. This applies to all plans, not only the Free Trial. Use the Vespa Console to extend the expiry period, or redeploy the application to add 7 more days."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "b5fddf9f",
      "metadata": {
        "id": "b5fddf9f"
      },
      "outputs": [],
      "source": [
        "from vespa.deployment import VespaCloud\n",
        "\n",
        "def read_secret():\n",
        "    \"\"\"Read the API key from the environment variable. This is\n",
        "    only used for CI/CD purposes.\"\"\"\n",
        "    t = os.getenv(\"VESPA_TEAM_API_KEY\")\n",
        "    if t:\n",
        "        return t.replace(r\"\\n\", \"\\n\")\n",
        "    else:\n",
        "        return t\n",
        "\n",
        "vespa_cloud = VespaCloud(\n",
        "    tenant=os.environ[\"TENANT_NAME\"],\n",
        "    application=vespa_app_name,\n",
        "    key_content=read_secret() if read_secret() else None,\n",
        "    key_location=api_key_path,\n",
        "    application_package=vespa_application_package)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fa9baa5a",
      "metadata": {
        "id": "fa9baa5a"
      },
      "source": [
        "Now deploy the app to Vespa Cloud dev zone.\n",
        "\n",
        "The first deployment typically takes 2 minutes until the endpoint is up."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "fe954dc4",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fe954dc4",
        "outputId": "a0764bd3-98c2-492a-b8d9-b99ecacf4bdb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Deployment started in run 1 of dev-aws-us-east-1c for vespa-team.pdfs. This may take a few minutes the first time.\n",
            "INFO    [19:04:30]  Deploying platform version 8.314.57 and application dev build 1 for dev-aws-us-east-1c of default ...\n",
            "INFO    [19:04:30]  Using CA signed certificate version 1\n",
            "INFO    [19:04:30]  Using 1 nodes in container cluster 'pdfs_container'\n",
            "INFO    [19:04:35]  Session 285265 for tenant 'vespa-team' prepared and activated.\n",
            "INFO    [19:04:39]  ######## Details for all nodes ########\n",
            "INFO    [19:04:44]  h88969d.dev.aws-us-east-1c.vespa-external.aws.oath.cloud: expected to be UP\n",
            "INFO    [19:04:44]  --- platform vespa/cloud-tenant-rhel8:8.314.57 <-- :\n",
            "INFO    [19:04:44]  --- container-clustercontroller on port 19050 has not started \n",
            "INFO    [19:04:44]  --- metricsproxy-container on port 19092 has not started \n",
            "INFO    [19:04:44]  h88978a.dev.aws-us-east-1c.vespa-external.aws.oath.cloud: expected to be UP\n",
            "INFO    [19:04:44]  --- platform vespa/cloud-tenant-rhel8:8.314.57 <-- :\n",
            "INFO    [19:04:44]  --- logserver-container on port 4080 has not started \n",
            "INFO    [19:04:44]  --- metricsproxy-container on port 19092 has not started \n",
            "INFO    [19:04:44]  h90615b.dev.aws-us-east-1c.vespa-external.aws.oath.cloud: expected to be UP\n",
            "INFO    [19:04:44]  --- platform vespa/cloud-tenant-rhel8:8.314.57 <-- :\n",
            "INFO    [19:04:44]  --- storagenode on port 19102 has not started \n",
            "INFO    [19:04:44]  --- searchnode on port 19107 has not started \n",
            "INFO    [19:04:44]  --- distributor on port 19111 has not started \n",
            "INFO    [19:04:44]  --- metricsproxy-container on port 19092 has not started \n",
            "INFO    [19:04:44]  h91135a.dev.aws-us-east-1c.vespa-external.aws.oath.cloud: expected to be UP\n",
            "INFO    [19:04:44]  --- platform vespa/cloud-tenant-rhel8:8.314.57 <-- :\n",
            "INFO    [19:04:44]  --- container on port 4080 has not started \n",
            "INFO    [19:04:44]  --- metricsproxy-container on port 19092 has not started \n",
            "INFO    [19:05:52]  Waiting for convergence of 10 services across 4 nodes\n",
            "INFO    [19:05:52]  1/1 nodes upgrading platform\n",
            "INFO    [19:05:52]  1 application services still deploying\n",
            "DEBUG   [19:05:52]  h91135a.dev.aws-us-east-1c.vespa-external.aws.oath.cloud: expected to be UP\n",
            "DEBUG   [19:05:52]  --- platform vespa/cloud-tenant-rhel8:8.314.57 <-- :\n",
            "DEBUG   [19:05:52]  --- container on port 4080 has not started \n",
            "DEBUG   [19:05:52]  --- metricsproxy-container on port 19092 has config generation 285265, wanted is 285265\n",
            "INFO    [19:06:21]  Found endpoints:\n",
            "INFO    [19:06:21]  - dev.aws-us-east-1c\n",
            "INFO    [19:06:21]   |-- https://bac3e5ad.c81e7b13.z.vespa-app.cloud/ (cluster 'pdfs_container')\n",
            "INFO    [19:06:22]  Installation succeeded!\n",
            "Using mTLS (key,cert) Authentication against endpoint https://bac3e5ad.c81e7b13.z.vespa-app.cloud//ApplicationStatus\n",
            "Application is up!\n",
            "Finished deployment.\n"
          ]
        }
      ],
      "source": [
        "from vespa.application import Vespa\n",
        "app:Vespa = vespa_cloud.deploy()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4cde8f22",
      "metadata": {
        "id": "4cde8f22"
      },
      "source": [
        "### Processing PDFs with LangChain\n",
        "\n",
        "[LangChain](https://python.langchain.com/) has a rich set of [document loaders](https://python.langchain.com/docs/modules/data_connection/document_loaders/) that can be used to load and process various file formats. In this notebook, we use the [PyPDFLoader](https://python.langchain.com/docs/modules/data_connection/document_loaders/pdf#using-pypdf).\n",
        "\n",
        "We also want to split the extracted text into *contexts* using a [text splitter](https://python.langchain.com/docs/modules/data_connection/document_transformers/). Most text embedding models have limited input lengths (typically less than 512 language model tokens, so splitting the text\n",
        "into multiple contexts that each fits into the context limit of the embedding model is a common strategy.\n",
        "\n",
        "For embedding text data, models based on the Transformer architecture have become the de facto standard. A challenge with Transformer-based models is their input length limitation due to the quadratic self-attention computational complexity. For example, a popular open-source text embedding model like\n",
        "[e5](https://huggingface.co/intfloat/e5-small) has an absolute maximum input length of 512 wordpiece tokens. In addition to\n",
        "the technical limitation, trying to fit more tokens than used during fine-tuning of the model will impact the quality of the vector representation.\n",
        "\n",
        "One can view this text embedding encoding as a lossy compression technique, where variable-length texts are compressed\n",
        "into a fixed dimensional vector representation.\n",
        "\n",
        "Although this compressed representation is very useful, it can be imprecise especially as the size of the text increases. By adding the ColBERT embedding, we also retain token-level information which retains more of the original meaning of the text and allows the richer _late interaction_ between the query and the document text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "d9e42b0f",
      "metadata": {
        "id": "d9e42b0f"
      },
      "outputs": [],
      "source": [
        "from langchain.document_loaders import PyPDFLoader\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "text_splitter = RecursiveCharacterTextSplitter(\n",
        "    chunk_size = 1024, #chars, not llm tokens\n",
        "    chunk_overlap  = 0,\n",
        "    length_function = len,\n",
        "    is_separator_regex = False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "adaccdfc",
      "metadata": {
        "id": "adaccdfc"
      },
      "source": [
        "The following iterates over the `sample_pdfs` and performs the following:\n",
        "\n",
        "- Load the URL and extract the text into pages. A page is the retrievable unit we will use in Vespa\n",
        "- For each page, use the text splitter to split the text into contexts. The contexts are represented as an `array<string>` in the Vespa schema\n",
        "- Create the page level Vespa `fields`, note that we duplicate some content like the title and URL into the page level representation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "bf8ac8c7",
      "metadata": {
        "id": "bf8ac8c7"
      },
      "outputs": [],
      "source": [
        "import hashlib\n",
        "import unicodedata\n",
        "def remove_control_characters(s):\n",
        "    return \"\".join(ch for ch in s if unicodedata.category(ch)[0]!=\"C\")\n",
        "\n",
        "my_docs_to_feed = []\n",
        "for pdf in sample_pdfs():\n",
        "    url = pdf['url']\n",
        "    loader = PyPDFLoader(url)\n",
        "    pages = loader.load_and_split()\n",
        "    for index, page in enumerate(pages):\n",
        "        source = page.metadata['source']\n",
        "        chunks = text_splitter.transform_documents([page])\n",
        "        text_chunks = [chunk.page_content for chunk in chunks]\n",
        "        text_chunks = [remove_control_characters(chunk) for chunk in text_chunks]\n",
        "        page_number = index + 1\n",
        "        vespa_id = f\"{url}#{page_number}\"\n",
        "        hash_value = hashlib.sha1(vespa_id.encode()).hexdigest()\n",
        "        fields = {\n",
        "            \"title\" : pdf['title'],\n",
        "            \"url\" : url,\n",
        "            \"page\": page_number,\n",
        "            \"id\": hash_value,\n",
        "            \"authors\": [a.strip() for a in pdf['authors'].split(\",\")],\n",
        "            \"contexts\": text_chunks,\n",
        "            \"metadata\": page.metadata\n",
        "        }\n",
        "        my_docs_to_feed.append(fields)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "54db44b1",
      "metadata": {
        "id": "54db44b1"
      },
      "source": [
        "Now that we have parsed the input PDFs and created a list of pages that we want to add to Vespa, we must format the\n",
        "list into the format that PyVespa accepts. Notice the `fields`, `id` and `groupname` keys. The `groupname` is the\n",
        "key that is used to shard and co-locate the data and is only relevant when using Vespa with streaming mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "bcbfa981",
      "metadata": {
        "id": "bcbfa981"
      },
      "outputs": [],
      "source": [
        "from typing import Iterable\n",
        "def vespa_feed(user:str) -> Iterable[dict]:\n",
        "    for doc in my_docs_to_feed:\n",
        "        yield {\n",
        "            \"fields\": doc,\n",
        "            \"id\": doc[\"id\"],\n",
        "            \"groupname\": user\n",
        "        }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "MMvbUZ1Vpuup",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MMvbUZ1Vpuup",
        "outputId": "3dd357d7-1a2f-4d51-abe2-efdd2dd6a829"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'title': 'ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction',\n",
              " 'url': 'https://arxiv.org/pdf/2112.01488.pdf',\n",
              " 'page': 1,\n",
              " 'id': 'a731a839198de04fa3d1a3cee6890d0d170ab025',\n",
              " 'authors': ['Keshav Santhanam',\n",
              "  'Omar Khattab',\n",
              "  'Jon Saad-Falcon',\n",
              "  'Christopher Potts',\n",
              "  'Matei Zaharia'],\n",
              " 'contexts': ['ColBERTv2:Effective and Efﬁcient Retrieval via Lightweight Late InteractionKeshav Santhanam∗Stanford UniversityOmar Khattab∗Stanford UniversityJon Saad-FalconGeorgia Institute of TechnologyChristopher PottsStanford UniversityMatei ZahariaStanford UniversityAbstractNeural information retrieval (IR) has greatlyadvanced search and other knowledge-intensive language tasks. While many neuralIR methods encode queries and documentsinto single-vector representations, lateinteraction models produce multi-vector repre-sentations at the granularity of each token anddecompose relevance modeling into scalabletoken-level computations. This decompositionhas been shown to make late interaction moreeffective, but it inﬂates the space footprint ofthese models by an order of magnitude. In thiswork, we introduce ColBERTv2, a retrieverthat couples an aggressive residual compres-sion mechanism with a denoised supervisionstrategy to simultaneously improve the quality',\n",
              "  'and space footprint of late interaction. Weevaluate ColBERTv2 across a wide rangeof benchmarks, establishing state-of-the-artquality within and outside the training domainwhile reducing the space footprint of lateinteraction models by 6–10 ×.1 IntroductionNeural information retrieval (IR) has quickly domi-nated the search landscape over the past 2–3 years,dramatically advancing not only passage and doc-ument search (Nogueira and Cho, 2019) but alsomany knowledge-intensive NLP tasks like open-domain question answering (Guu et al., 2020),multi-hop claim veriﬁcation (Khattab et al., 2021a),and open-ended generation (Paranjape et al., 2022).Many neural IR methods follow a single-vectorsimilarity paradigm: a pretrained language modelis used to encode each query and each documentinto a single high-dimensional vector, and rele-vance is modeled as a simple dot product betweenboth vectors. An alternative is late interaction , in-troduced in ColBERT (Khattab and Zaharia, 2020),',\n",
              "  'where queries and documents are encoded at a ﬁner-granularity into multi-vector representations, and∗Equal contribution.relevance is estimated using rich yet scalable in-teractions between these two sets of vectors. Col-BERT produces an embedding for every token inthe query (and document) and models relevanceas the sum of maximum similarities between eachquery vector and all vectors in the document.By decomposing relevance modeling into token-level computations, late interaction aims to reducethe burden on the encoder: whereas single-vectormodels must capture complex query–document re-lationships within one dot product, late interactionencodes meaning at the level of tokens and del-egates query–document matching to the interac-tion mechanism. This added expressivity comesat a cost: existing late interaction systems imposean order-of-magnitude larger space footprint thansingle-vector models, as they must store billionsof small vectors for Web-scale collections. Con-',\n",
              "  'sidering this challenge, it might seem more fruit-ful to focus instead on addressing the fragility ofsingle-vector models (Menon et al., 2022) by in-troducing new supervision paradigms for negativemining (Xiong et al., 2020), pretraining (Gao andCallan, 2021), and distillation (Qu et al., 2021).Indeed, recent single-vector models with highly-tuned supervision strategies (Ren et al., 2021b; For-mal et al., 2021a) sometimes perform on-par oreven better than “vanilla” late interaction models,and it is not necessarily clear whether late inter-action architectures—with their ﬁxed token-levelinductive biases—admit similarly large gains fromimproved supervision.In this work, we show that late interaction re-trievers naturally produce lightweight token rep-resentations that are amenable to efﬁcient storageoff-the-shelf and that they can beneﬁt drasticallyfrom denoised supervision. We couple those inColBERTv2 ,1a new late-interaction retriever that'],\n",
              " 'metadata': {'source': 'https://arxiv.org/pdf/2112.01488.pdf', 'page': 0}}"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "my_docs_to_feed[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2ff628ac",
      "metadata": {
        "id": "2ff628ac"
      },
      "source": [
        "Now, we can feed to the Vespa instance (`app`), using the `feed_iterable` API, using the generator function above as input\n",
        "with a custom `callback` function. Vespa also performs embedding inference during this step using the built-in Vespa [embedding](https://docs.vespa.ai/en/embedding.html#huggingface-embedder) functionality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "dc1b3029",
      "metadata": {
        "id": "dc1b3029"
      },
      "outputs": [],
      "source": [
        "from vespa.io import VespaResponse\n",
        "\n",
        "def callback(response:VespaResponse, id:str):\n",
        "    if not response.is_successful():\n",
        "        print(f\"Document {id} failed to feed with status code {response.status_code}, url={response.url} response={response.json}\")\n",
        "\n",
        "app.feed_iterable(schema=\"pdf\", iter=vespa_feed(\"jo-bergum\"), namespace=\"personal\", callback=callback)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "431dc2f9",
      "metadata": {
        "id": "431dc2f9"
      },
      "source": [
        "Notice the `schema` and `namespace` arguments. PyVespa transforms the input operations to Vespa [document v1](https://docs.vespa.ai/en/document-v1-api-guide.html)\n",
        "requests.\n",
        "\n",
        "![Document id](https://blog.vespa.ai/assets/2023-12-08-turbocharge-rag-with-langchain-and-vespa-streaming-mode/docid.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "20b007ec",
      "metadata": {
        "id": "20b007ec"
      },
      "source": [
        "### Querying data\n",
        "\n",
        "Now, we can also query our data. With [streaming mode](https://docs.vespa.ai/en/reference/query-api-reference.html#streaming),\n",
        "we must pass the `groupname` parameter, or the request will fail with an error.\n",
        "\n",
        "The query request uses the Vespa Query API  and the `Vespa.query()` function\n",
        "supports passing any of the Vespa query API parameters.\n",
        "\n",
        "Read more about querying Vespa in:\n",
        "\n",
        "- [Vespa Query API](https://docs.vespa.ai/en/query-api.html)\n",
        "- [Vespa Query API reference](https://docs.vespa.ai/en/reference/query-api-reference.html)\n",
        "- [Vespa Query Language API (YQL)](https://docs.vespa.ai/en/query-language.html)\n",
        "\n",
        "Sample query request for `why is colbert effective?` for the user `jo-bergum`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "id": "b9349fb4",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b9349fb4",
        "outputId": "08eafc2f-0856-4c6b-9f13-2decf754228c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"id\": \"id:personal:pdf:g=jo-bergum:55ea3f735cb6748a2eddb9f76d3f0e7fff0c31a8\",\n",
            "  \"relevance\": 103.17699432373047,\n",
            "  \"source\": \"pdfs_content.pdf\",\n",
            "  \"fields\": {\n",
            "    \"matchfeatures\": {\n",
            "      \"cos_sim\": 0.6534222205340683,\n",
            "      \"max_sim\": 103.17699432373047,\n",
            "      \"max_sim_per_context\": {\n",
            "        \"0\": 74.16375732421875,\n",
            "        \"1\": 103.17699432373047\n",
            "      }\n",
            "    },\n",
            "    \"id\": \"55ea3f735cb6748a2eddb9f76d3f0e7fff0c31a8\",\n",
            "    \"title\": \"ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT\",\n",
            "    \"page\": 18,\n",
            "    \"contexts\": [\n",
            "      \"at least once. While ColBERT encodes each document with BERTexactly once, existing BERT-based rankers would repeat similarcomputations on possibly hundreds of documents for each query.Se/t_ting Dimension( m) Bytes/Dim Space(GiBs) MRR@10Re-rank Cosine 128 4 286 34.9End-to-end L2 128 2 154 36.0Re-rank L2 128 2 143 34.8Re-rank Cosine 48 4 54 34.4Re-rank Cosine 24 2 27 33.9Table 4: Space Footprint vs MRR@10 (Dev) on MS MARCO.Table 4 reports the space footprint of ColBERT under variousse/t_tings as we reduce the embeddings dimension and/or the bytesper dimension. Interestingly, the most space-e\\ufb03cient se/t_ting, thatis, re-ranking with cosine similarity with 24-dimensional vectorsstored as 2-byte /f_loats, is only 1% worse in MRR@10 than the mostspace-consuming one, while the former requires only 27 GiBs torepresent the MS MARCO collection.5 CONCLUSIONSIn this paper, we introduced ColBERT, a novel ranking model thatemploys contextualized late interaction over deep LMs (in particular,\",\n",
            "      \"BERT) for e\\ufb03cient retrieval. By independently encoding queriesand documents into /f_ine-grained representations that interact viacheap and pruning-friendly computations, ColBERT can leveragethe expressiveness of deep LMs while greatly speeding up queryprocessing. In addition, doing so allows using ColBERT for end-to-end neural retrieval directly from a large document collection. Ourresults show that ColBERT is more than 170 \\u00d7faster and requires14,000\\u00d7fewer FLOPs/query than existing BERT-based models, allwhile only minimally impacting quality and while outperformingevery non-BERT baseline.Acknowledgments. OK was supported by the Eltoukhy FamilyGraduate Fellowship at the Stanford School of Engineering. /T_hisresearch was supported in part by a\\ufb03liate members and othersupporters of the Stanford DAWN project\\u2014Ant Financial, Facebook,Google, Infosys, NEC, and VMware\\u2014as well as Cisco, SAP, and the\"\n",
            "    ]\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "from vespa.io import VespaQueryResponse\n",
        "import json\n",
        "\n",
        "response:VespaQueryResponse = app.query(\n",
        "    yql=\"select id,title,page,contexts from pdf where ({targetHits:10}nearestNeighbor(embedding,q))\",\n",
        "    groupname=\"jo-bergum\",\n",
        "    ranking=\"colbert\",\n",
        "    query=\"why is colbert effective?\",\n",
        "    body={\n",
        "        \"presentation.format.tensors\": \"short-value\",\n",
        "        \"input.query(q)\": \"embed(e5, \\\"why is colbert effective?\\\")\",\n",
        "        \"input.query(qt)\": \"embed(colbert, \\\"why is colbert effective?\\\")\",\n",
        "    }\n",
        ")\n",
        "assert(response.is_successful())\n",
        "print(json.dumps(response.hits[0], indent=2))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4d3ca1da",
      "metadata": {
        "id": "4d3ca1da"
      },
      "source": [
        "Notice the `matchfeatures` that returns the configured match-features from the rank-profile, including all the context similarities."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "57f323df",
      "metadata": {
        "id": "57f323df"
      },
      "source": [
        "## LangChain Retriever\n",
        "\n",
        "We use the [LangChain Retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/) interface so that\n",
        "we can connect our Vespa app with the flexibility and power of the [LangChain](https://python.langchain.com/docs/get_started/introduction) LLM framework.\n",
        "\n",
        ">A retriever is an interface that returns documents given an unstructured query. It is more general than a vector store. A retriever does not need to be able to store documents, only to return (or retrieve) them. Vector stores can be used as the backbone of a retriever, but there are other types of retrievers as well.\n",
        "\n",
        "The retriever interface fits perfectly with Vespa, as Vespa can support a wide range of features and ways to retrieve and\n",
        "rank content. The following implements a custom retriever `VespaStreamingColBERTRetriever` that takes the following arguments:\n",
        "\n",
        "- `app:Vespa` The Vespa application we retrieve from. This could be a Vespa Cloud instance or a local instance, for example running on a laptop. \n",
        "-  `user:str` The user that that we want to retrieve for, this argument maps to the [Vespa streaming mode groupname parameter](https://docs.vespa.ai/en/reference/query-api-reference.html#streaming.groupname)\n",
        "- `pages:int` The  target number of PDF pages we want to retrieve for a given query \n",
        "- `chunks_per_page` The is the target number of relevant text chunks that are associated with the page\n",
        "- `chunk_similarity_threshold` - The chunk similarity threshold, only chunks with a similarity above this threshold \n",
        "\n",
        "The core idea is to _retrieve_ pages using max context similarity as the initial scoring function, then re-rank the top-K pages using the ColBERT embeddings. This re-ranking is handled by the second phase of the Vespa ranking expression defined above, and is transparent to the retriever code below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "id": "66756a7f",
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_core.documents import Document\n",
        "from langchain_core.retrievers import BaseRetriever\n",
        "from typing import List\n",
        "\n",
        "\n",
        "class VespaStreamingColBERTRetriever(BaseRetriever):\n",
        "\n",
        "    app: Vespa\n",
        "    user:str\n",
        "    pages: int = 5\n",
        "    chunks_per_page: int = 3\n",
        "    chunk_similarity_threshold: float = 0.8\n",
        "\n",
        "    def _get_relevant_documents(self, query: str) -> List[Document]:\n",
        "        response:VespaQueryResponse = self.app.query(\n",
        "            yql=\"select id, url, title, page, authors, contexts from pdf where userQuery() or ({targetHits:20}nearestNeighbor(embedding,q))\",\n",
        "            groupname=self.user,\n",
        "            ranking=\"colbert\",\n",
        "            query=query,\n",
        "            hits = self.pages,\n",
        "            body={\n",
        "                \"presentation.format.tensors\": \"short-value\",\n",
        "                \"input.query(q)\": f\"embed(e5, \\\"query: {query} \\\")\",\n",
        "                \"input.query(qt)\": f\"embed(colbert, \\\"{query}\\\")\"\n",
        "            }\n",
        "        )\n",
        "        if not response.is_successful():\n",
        "            raise ValueError(f\"Query failed with status code {response.status_code}, url={response.url} response={response.json}\")\n",
        "        return self._parse_response(response)\n",
        "\n",
        "\n",
        "    def _parse_response(self, response: VespaQueryResponse) -> List[Document]:\n",
        "        documents: List[Document] = []\n",
        "        for hit in response.hits:\n",
        "            fields = hit['fields']\n",
        "            chunks_with_scores = self._get_chunk_similarities(fields)\n",
        "            ## Best k chunks from each page\n",
        "            best_chunks_on_page = \" ### \".join(\n",
        "                [chunk for chunk, score in\n",
        "                    chunks_with_scores[0:self.chunks_per_page] if score > self.chunk_similarity_threshold])\n",
        "            documents.append(\n",
        "                Document(\n",
        "                    id=fields['id'],\n",
        "                    page_content=best_chunks_on_page,\n",
        "                    title=fields['title'],\n",
        "                    metadata={\n",
        "                        \"title\": fields['title'],\n",
        "                        \"url\": fields['url'],\n",
        "                        \"page\": fields['page'],\n",
        "                        \"authors\": fields['authors'],\n",
        "                        \"features\": fields['matchfeatures']\n",
        "                    }\n",
        "                )\n",
        "            )\n",
        "        return documents\n",
        "\n",
        "    def _get_chunk_similarities(self, hit_fields: dict) -> List[tuple]:\n",
        "        match_features = hit_fields['matchfeatures']\n",
        "        similarities = match_features['max_sim_per_context']\n",
        "        chunk_scores = []\n",
        "        for i in range(0,len(similarities)):\n",
        "            chunk_scores.append(similarities.get(str(i), 0))\n",
        "        chunks = hit_fields['contexts']\n",
        "        chunks_with_scores = list(zip(chunks, chunk_scores))\n",
        "        return sorted(chunks_with_scores, key=lambda x: x[1], reverse=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "341dd861",
      "metadata": {
        "id": "341dd861"
      },
      "source": [
        "That's it! We can give our newborn retriever a spin for the user `jo-bergum` by"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "id": "ac9088a4",
      "metadata": {
        "id": "ac9088a4"
      },
      "outputs": [],
      "source": [
        "vespa_hybrid_retriever = VespaStreamingColBERTRetriever(app=app, user=\"jo-bergum\", pages=1, chunks_per_page=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "id": "3198db04",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3198db04",
        "outputId": "8d25439c-e8a2-4c2e-9d70-8f1bd5f57ae4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[Document(page_content='ture that precisely does so. As illustrated, every query embeddinginteracts with all document embeddings via a MaxSim operator,which computes maximum similarity (e.g., cosine similarity), andthe scalar outputs of these operators are summed across queryterms. /T_his paradigm allows ColBERT to exploit deep LM-basedrepresentations while shi/f_ting the cost of encoding documents of-/f_line and amortizing the cost of encoding the query once acrossall ranked documents. Additionally, it enables ColBERT to lever-age vector-similarity search indexes (e.g., [ 1,15]) to retrieve thetop-kresults directly from a large document collection, substan-tially improving recall over models that only re-rank the output ofterm-based retrieval.As Figure 1 illustrates, ColBERT can serve queries in tens orfew hundreds of milliseconds. For instance, when used for re-ranking as in “ColBERT (re-rank)”, it delivers over 170 ×speedup(and requires 14,000 ×fewer FLOPs) relative to existing BERT-based ### models, while being more eﬀective than every non-BERT baseline(§4.2 & 4.3). ColBERT’s indexing—the only time it needs to feeddocuments through BERT—is also practical: it can index the MSMARCO collection of 9M passages in about 3 hours using a singleserver with four GPUs ( §4.5), retaining its eﬀectiveness with a spacefootprint of as li/t_tle as few tens of GiBs. Our extensive ablationstudy ( §4.4) shows that late interaction, its implementation viaMaxSim operations, and crucial design choices within our BERT-based encoders are all essential to ColBERT’s eﬀectiveness.Our main contributions are as follows.(1)We propose late interaction (§3.1) as a paradigm for eﬃcientand eﬀective neural ranking.(2)We present ColBERT ( §3.2 & 3.3), a highly-eﬀective modelthat employs novel BERT-based query and document en-coders within the late interaction paradigm.', metadata={'title': 'ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT', 'url': 'https://arxiv.org/pdf/2004.12832.pdf', 'page': 4, 'authors': ['Omar Khattab', 'Matei Zaharia'], 'features': {'cos_sim': 0.6664045997289173, 'max_sim': 124.19231414794922, 'max_sim_per_context': {'0': 124.19231414794922, '1': 92.21265411376953}}})]"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vespa_hybrid_retriever.get_relevant_documents(\"what is the maxsim operator in colbert?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fcca4fc7",
      "metadata": {
        "id": "fcca4fc7"
      },
      "source": [
        "## RAG\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a84b98db",
      "metadata": {
        "id": "a84b98db"
      },
      "source": [
        "Finally, we can connect our custom retriever with the complete flexibility and power of the [LangChain] LLM framework.\n",
        "The following uses [LangChain Expression Language, or LCEL](https://python.langchain.com/docs/expression_language/), a declarative way to compose chains.\n",
        "\n",
        "We have several steps composed into a chain:\n",
        "\n",
        "- The prompt template and LLM model, in this case using OpenAI\n",
        "- The retriever that provides the retrieved context for the question\n",
        "- The formatting of the retrieved context\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "id": "e3dcf5b4",
      "metadata": {
        "id": "e3dcf5b4"
      },
      "outputs": [],
      "source": [
        "vespa_hybrid_retriever = VespaStreamingColBERTRetriever(app=app, user=\"jo-bergum\", chunks_per_page=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "id": "d95473dc",
      "metadata": {
        "id": "d95473dc"
      },
      "outputs": [],
      "source": [
        "from langchain_openai import ChatOpenAI\n",
        "from langchain.prompts import ChatPromptTemplate\n",
        "from langchain.schema import StrOutputParser\n",
        "from langchain.schema.runnable import RunnablePassthrough\n",
        "\n",
        "prompt_template = \"\"\"\n",
        "Answer the question based only on the following context.\n",
        "Cite the page number and the url of the document you are citing.\n",
        "\n",
        "{context}\n",
        "Question: {question}\n",
        "\"\"\"\n",
        "prompt = ChatPromptTemplate.from_template(prompt_template)\n",
        "model = ChatOpenAI(model=\"gpt-4-0125-preview\")\n",
        "\n",
        "def format_prompt_context(docs) -> str:\n",
        "    context = []\n",
        "    for d in docs:\n",
        "        context.append(f\"{d.metadata['title']} by {d.metadata['authors']}\\n\")\n",
        "        context.append(f\"url: {d.metadata['url']}\\n\")\n",
        "        context.append(f\"page: {d.metadata['page']}\\n\")\n",
        "        context.append(f\"{d.page_content}\\n\\n\")\n",
        "    return \"\".join(context)\n",
        "\n",
        "\n",
        "chain = (\n",
        "    {\"context\": vespa_hybrid_retriever | format_prompt_context, \"question\": RunnablePassthrough()}\n",
        "    | prompt\n",
        "    | model\n",
        "    | StrOutputParser()\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "562d2c7d",
      "metadata": {
        "id": "562d2c7d"
      },
      "source": [
        "### Interact with the chain\n",
        "\n",
        "Now, we can start asking questions using the `chain` define above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "id": "36f7f092",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "id": "36f7f092",
        "outputId": "d509cc65-1a08-4c39-b987-14c007d0b9bf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'ColBERT, introduced by Omar Khattab and Matei Zaharia, is a novel ranking model that employs contextualized late interaction over deep language models (LMs), specifically focusing on BERT (Bidirectional Encoder Representations from Transformers) for efficient and effective passage search. It achieves this by independently encoding queries and documents into fine-grained representations that interact via cheap and pruning-friendly computations. This approach allows ColBERT to leverage the expressiveness of deep LMs while significantly speeding up query processing compared to existing BERT-based models. ColBERT also enables end-to-end neural retrieval directly from a large document collection, offering more than 170 times faster performance and requiring 14,000 times fewer FLOPs (floating-point operations) per query than previous BERT-based models, with minimal impact on quality. It outperforms every non-BERT baseline in effectiveness (https://arxiv.org/pdf/2004.12832.pdf, page 18).\\n\\nColBERT differentiates itself with a mechanism that delays the query-document interaction, which allows for pre-computation of document representations for cheap neural re-ranking and supports practical end-to-end neural retrieval through pruning via vector-similarity search. This method preserves the effectiveness of state-of-the-art models that condition most of their computations on the joint query-document pair, making ColBERT a scalable solution for passage search challenges (https://arxiv.org/pdf/2004.12832.pdf, page 6).'"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chain.invoke(\"what is colbert?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "id": "569929de",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        },
        "id": "569929de",
        "outputId": "035e08d4-81e5-4421-a1d1-e1039fa6bd26"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'The ColBERT MaxSim operator is a mechanism for computing the maximum similarity between query embeddings and document embeddings. It operates by calculating the maximum similarity (e.g., cosine similarity) for each query embedding with all document embeddings, and then summing the scalar outputs of these operations across query terms. This paradigm enables the efficient and effective retrieval of documents by allowing for the interaction between deep language model-based representations of queries and documents to occur in a late stage of the processing pipeline, thereby shifting the cost of encoding documents offline and amortizing the cost of encoding the query across all ranked documents. Additionally, the MaxSim operator facilitates the use of vector-similarity search indexes to directly retrieve the top-k results from a large document collection, substantially improving recall over models that only re-rank the output of term-based retrieval. This operator is a key component of ColBERT\\'s approach to efficient and effective passage search.\\n\\nSource: \"ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT\" by Omar Khattab and Matei Zaharia, page 4, https://arxiv.org/pdf/2004.12832.pdf'"
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chain.invoke(\"what is the colbert maxsim operator\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "id": "fde46620",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        },
        "id": "fde46620",
        "outputId": "26ae73f5-931c-4dcb-cf0f-1d3bcacb9cd0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'The main difference between ColBERT and single-vector representational models lies in their approach to handling document and query representations for information retrieval tasks. ColBERT utilizes a multi-vector representation for both queries and documents, whereas single-vector models encode each query and each document into a single, dense vector.\\n\\n1. **Multi-Vector vs. Single-Vector Representations**: ColBERT leverages a late interaction mechanism that allows for fine-grained matching between the multiple embeddings of query terms and document tokens. This approach enables capturing the nuanced semantics of the text by considering the contextualized representation of each term separately. On the other hand, single-vector models compress the entire content of a document or a query into a single dense vector, which might lead to a loss of detail and context specificity.\\n\\n2. **Efficiency and Effectiveness**: While single-vector models might be simpler and potentially faster in some scenarios due to their straightforward matching mechanism (e.g., cosine similarity between query and document vectors), this simplicity could come at the cost of effectiveness. ColBERT, with its detailed interaction between term-level vectors, can offer more accurate retrieval results because it preserves and utilizes the rich semantic relationships within and across the text of queries and documents. However, ColBERT\\'s detailed approach initially required more storage and computational resources compared to single-vector models. Nonetheless, advancements like ColBERTv2 have significantly improved the efficiency, achieving competitive storage requirements and reducing the computational cost while maintaining or even enhancing retrieval effectiveness.\\n\\n3. **Compression and Storage**: Initial versions of multi-vector models like ColBERT required significantly more storage space compared to single-vector models due to storing multiple vectors per document. However, with the introduction of techniques like residual compression in ColBERTv2, the storage requirements have been drastically reduced to levels competitive with single-vector models. Single-vector models, while naturally more storage-efficient, can also be compressed, but aggressive compression might exacerbate the loss in quality.\\n\\n4. **Search Quality and Compression**: Despite the potential for aggressive compression in single-vector models, such approaches often lead to a more pronounced loss in quality compared to late interaction methods like ColBERTv2. ColBERTv2, even when employing compression techniques to reduce its storage footprint, can achieve higher quality across systems, showcasing the robustness of its retrieval capabilities even when optimizing for space efficiency.\\n\\nIn summary, the difference between ColBERT and single-vector representational models is primarily in their approach to encoding and matching queries and documents, with ColBERT focusing on detailed, term-level interactions for improved accuracy, and single-vector models emphasizing simplicity and compactness, which might come at the cost of retrieval effectiveness.\\n\\nCitations:\\n- Santhanam et al., \"ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction,\" p. 14, 15, 17, https://arxiv.org/pdf/2112.01488.pdf'"
            ]
          },
          "execution_count": 33,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chain.invoke(\"What is the difference between colbert and single vector representational models?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "id": "8852bca0",
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\"ColBERT is designed to efficiently handle the interaction between query and document representations through a mechanism called late interaction, which is particularly beneficial when dealing with longer documents. This is because ColBERT independently encodes queries and documents into fine-grained representations using BERT, and then employs a cheap yet powerful interaction step that models their fine-grained similarity. This approach allows for the pre-computation of document representations offline, significantly speeding up query processing by avoiding the need to feed each query-document pair through a massive neural network at query time.\\n\\nFor longer documents, the benefits of this approach are twofold:\\n\\n1. **Efficiency in Handling Long Documents**: Since ColBERT encodes document representations offline, it can efficiently manage longer documents without a proportional increase in computational cost at query time. This is unlike traditional BERT-based models that might require more computational resources to process longer documents due to their size and complexity.\\n\\n2. **Effectiveness in Capturing Fine-Grained Semantics**: The fine-grained representations and the late interaction mechanism enable ColBERT to effectively capture the nuances and detailed semantics of longer documents. This is crucial for maintaining high retrieval quality, as longer documents often contain more information and require a more nuanced understanding to match relevant queries accurately.\\n\\nThus, ColBERT's architecture, which leverages the strengths of BERT for deep language understanding while introducing efficiencies through late interaction, makes it particularly adept at handling longer documents. It achieves this by pre-computing and efficiently utilizing detailed semantic representations of documents, enabling both high-quality retrieval and significant speed-ups in query processing times compared to traditional BERT-based models.\\n\\nReference: ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT by ['Omar Khattab', 'Matei Zaharia'] (https://arxiv.org/pdf/2004.12832.pdf), page 4.\""
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chain.invoke(\"Why does ColBERT work better for longer documents?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7c8b8223",
      "metadata": {
        "id": "7c8b8223"
      },
      "source": [
        "\n",
        "## Summary\n",
        "\n",
        "Vespa’s streaming mode is a game-changer, enabling the creation of highly cost-effective RAG applications for naturally partitioned data. Now it is also possible to use ColBERT for re-ranking, without having to integrate any custom embedder or re-ranking code.\n",
        "\n",
        "\n",
        "In this notebook, we delved into the hands-on application of [LangChain](https://python.langchain.com/docs/get_started/introduction),\n",
        "leveraging document loaders and transformers. Finally, we showcased a custom LangChain retriever that connected\n",
        "all the functionality of LangChain with Vespa.\n",
        "\n",
        "For those interested in learning more about Vespa, join the [Vespa community on Slack](https://vespatalk.slack.com/) to exchange ideas,\n",
        "seek assistance, or stay in the loop on the latest Vespa developments.\n",
        "\n",
        "\n",
        "We can now delete the cloud instance:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "71e310e3",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "71e310e3",
        "outputId": "991b1965-6c33-4985-e873-a92c43695528"
      },
      "outputs": [],
      "source": [
        "vespa_cloud.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.11.4 64-bit",
      "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.13"
    },
    "vscode": {
      "interpreter": {
        "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}