generate_vectors_for_knn_in_elastic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jeffvestal/elastic_jupyter_notebooks/blob/main/image_vector_search_load_and_search.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xoLDtS_6Df1"
      },
      "source": [
        "# Generating vectors on ingest and testing kNN search\n",
        "\n",
        "This code will show you how to set up an ingest pipeline to generate vectors for documents on ingest. We will then run a sample kNN search\n",
        "\n",
        "Overview of steps\n",
        "1. Set up our python environment\n",
        "2. Setup index mapping\n",
        "3. Configure ingest pipeline\n",
        "4. Index a couple test documents\n",
        "5. Run a sample kNN search\n",
        "\n",
        "### Requirements\n",
        "This notebook assumes you already have loaded an embedding model into elasticsearch. If you haven't, please start with [this notebook example](https://github.com/jeffvestal/elastic_jupyter_notebooks/blob/main/load_embedding_model_from_hf_to_elastic.ipynb)\n",
        "\n",
        "\n",
        "### Elastic version support\n",
        "Requires Elastic version 8.0+ with a platinum or enterprise license (or trial license)\n",
        "\n",
        "You can set up a [free trial elasticsearch Deployment in Elastic Cloud](https://cloud.elastic.co/registration)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DgxCKQS7mCZw"
      },
      "source": [
        "---\n",
        "---\n",
        "# Setup\n",
        "---\n",
        "---\n",
        "This section will set up the python environment with the required libraries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ly1f1P-l9ri8"
      },
      "source": [
        "## Install and import required python libraries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MJAb_8zlPFhQ"
      },
      "source": [
        "Elastic uses the [eland python library](https://github.com/elastic/eland) to download modesl from Hugging Face hub and load them into elasticsearch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rUedSzQW9FIF"
      },
      "outputs": [],
      "source": [
        "pip install eland"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NK3Wx1I199yB"
      },
      "outputs": [],
      "source": [
        "pip install elasticsearch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cEfiiFXakzdP"
      },
      "outputs": [],
      "source": [
        "pip install transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I20mDmJboKZw"
      },
      "outputs": [],
      "source": [
        "pip install sentence_transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uqcpWrbkBEB9"
      },
      "outputs": [],
      "source": [
        "pip install torch==1.11"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyUZXUi4RWWL"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "from eland.ml.pytorch import PyTorchModel\n",
        "from eland.ml.pytorch.transformers import TransformerModel\n",
        "from elasticsearch import Elasticsearch, helpers\n",
        "from elasticsearch.client import MlClient\n",
        "from pprint import pprint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7nMIbHke37Q"
      },
      "source": [
        "## Configure elasticsearch authentication. \n",
        "The recommended authentication approach is using the [Elastic Cloud ID](https://www.elastic.co/guide/en/cloud/current/ec-cloud-id.html) and a [cluster level API key](https://www.elastic.co/guide/en/kibana/current/api-keys.html)\n",
        "\n",
        "You can use any method you wish to set the required credentials. We are using getpass in this example to prompt for credentials to avoide storing them in github."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xsd2m7HoTCLm"
      },
      "outputs": [],
      "source": [
        "import getpass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SSGgYHome69o"
      },
      "outputs": [],
      "source": [
        "es_cloud_id = getpass.getpass('Enter Elastic Cloud ID:  ')\n",
        "es_api_id = getpass.getpass('Enter cluster API key ID:  ') \n",
        "es_api_key = getpass.getpass('Enter cluster API key:  ')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jL4VDnVp96lf"
      },
      "source": [
        "## Connect to Elastic Cloud"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I8mVJkKmetXo"
      },
      "outputs": [],
      "source": [
        "es = Elasticsearch(cloud_id=es_cloud_id, \n",
        "                   api_key=(es_api_id, es_api_key)\n",
        "                   )\n",
        "es.info() # should return cluster info"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4UYSzFp3vHdB"
      },
      "source": [
        "---\n",
        "---\n",
        "# Model Information and Status\n",
        "---\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wQwfozwznK4Y"
      },
      "source": [
        "## View information about the model\n",
        "This is not required but will allow us to get the model_id as it is stored in elasticsearch as well as verify the model is running / deployed and ready to use in our ingest pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4Wv8EJvpfZI"
      },
      "outputs": [],
      "source": [
        "m = MlClient.get_trained_models(es)\n",
        "m.body"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KbbdWiJetJV0"
      },
      "source": [
        "## Set the model_id for ease of reference later\n",
        "To make is easy for reference later, we will set  `es_model_id` to the `model_id` listed in the output above"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P8xW5_lCtUE2"
      },
      "outputs": [],
      "source": [
        "es_model_id = \"sentence-transformers__msmarco-minilm-l-12-v3\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oMGw3sk-pbaN"
      },
      "source": [
        "### *If* the model is not started we will need to deploy the model\n",
        "\n",
        "You will only need to run this if the model hasn't been deployed. \n",
        "\n",
        "This will load the model on the ML nodes and start the process(es) making it available for the NLP task\n",
        "\n",
        "uncomment the code below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w5muJ1rLqvUW"
      },
      "outputs": [],
      "source": [
        "#s = MlClient.start_trained_model_deployment(es, model_id=es_model_id)\n",
        "#s.body"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZytlELrsnn_O"
      },
      "source": [
        "#### Verify the model started without issue\n",
        "If you aren't sure if the model is started you can check here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZaQUUWe0Hxwz"
      },
      "outputs": [],
      "source": [
        "stats = MlClient.get_trained_models_stats(es, model_id=es_model_id)\n",
        "stats.body['trained_model_stats'][0]['deployment_stats']['nodes'][0]['routing_state']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEwsReS8zyOc"
      },
      "source": [
        "---\n",
        "---\n",
        "# Elasticsearch index setup\n",
        "---\n",
        "---\n",
        "Here we will configure an index template with settings and mappings to store our vectors and text data\n",
        "\n",
        "The **important** part here will be setting our vector field to be a `dense_vector` type. This will tell elasticsearch to build the HNSW graph for the vectors so we can then use kNN search later. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KQvNTOJQ2Jk1"
      },
      "source": [
        "## Define the index template\n",
        "We will have the following fields\n",
        "\n",
        "- `vectors` of type `dense_vector`\n",
        "-- it is important to set `dims` to the number of dimensions the model you will use outputs\n",
        "- `title` of type `text`\n",
        "- `summary` of type `text`\n",
        "\n",
        "We will have \n",
        "- 1 primary shard\n",
        "- 0 replica -> *note* in production you will want at least 1 replica\n",
        "\n",
        "This will match new indices with the name matching the pattern of `jupyter-vector-demo*`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e8079Ic44SEO"
      },
      "outputs": [],
      "source": [
        "index_patterns = \"jupyter-vector-demo*\"\n",
        "settings= {\n",
        "        \"number_of_shards\": 1,\n",
        "        \"number_of_replicas\": 0\n",
        "    }\n",
        "mappings= {\n",
        "        \"properties\": {\n",
        "            \"vectors\": {\n",
        "                \"type\": \"dense_vector\",\n",
        "                \"dims\": 384,\n",
        "                \"index\" : True,\n",
        "                \"similarity\" : \"cosine\"\n",
        "            },\n",
        "            \"title\": {\n",
        "                \"type\": \"text\"\n",
        "            },\n",
        "            \"summary\": {\n",
        "                \"type\": \"text\"\n",
        "            }\n",
        "        }\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vryddQGB3U6q"
      },
      "source": [
        "## Apply the template\n",
        "Here we apply the templat and give it a name of `jupyter-vector-demo`. This is just the name of the template if we need to modify it later on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LEmOQ4IT3XK8"
      },
      "outputs": [],
      "source": [
        "es.indices.put_template(name=\"jupyter-vector-demo-template\", \n",
        "                        index_patterns=index_patterns,\n",
        "                        settings=settings,\n",
        "                        mappings=mappings\n",
        "                        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3MZ6EBVUTjhN"
      },
      "source": [
        "---\n",
        "---\n",
        "# The Ingest Pipeline\n",
        "---\n",
        "---\n",
        "\n",
        "An ingest pipeline has one or more processors and processes documents before they are written into an elasticsearch index. \n",
        "\n",
        "Each processor is designed to perform a various task such as parsing fields or enriching data. \n",
        "\n",
        "The main processor for this pipeline is the `inference` processor. The inference processor sends a specified field to a supervised model and writes the output from the model to a new field along with the original fields in the document. \n",
        "\n",
        "To make it simpler to access the vector, we will copy the vectors to a field named `vectors` and them remove the `ml` field tree which is the default output."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3iOh80S0UhsU"
      },
      "source": [
        "## Configure the pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DwwyOBWEVd-P"
      },
      "outputs": [],
      "source": [
        "pipeline_definition = {\n",
        "    \"description\": \"A pipeline for generating and storing vectors on ingest\",\n",
        "    \"processors\": [\n",
        "      {\n",
        "       \"inference\": {\n",
        "          \"model_id\": \"sentence-transformers__msmarco-minilm-l-12-v3\",\n",
        "          \"field_map\": {\n",
        "           \"summary\": \"text_field\"\n",
        "          }\n",
        "       }\n",
        "     },\n",
        "     {\n",
        "      \"set\": {\n",
        "        \"field\": \"vectors\",\n",
        "        \"copy_from\": \"ml.inference.predicted_value\"\n",
        "        }\n",
        "     },\n",
        "    {\n",
        "      \"remove\": {\n",
        "        \"field\": \"ml\"\n",
        "      }\n",
        "    }\n",
        "  ]\n",
        "}\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dHzxqcvIVjyC"
      },
      "source": [
        "## Create the pipeline if it doesn't exist"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LXZUQI3IVp21"
      },
      "outputs": [],
      "source": [
        "if es.ingest.put_pipeline(id=\"jupyter-vector-demo-pipeline\", body=pipeline_definition):\n",
        "    print(\"Pipeline created successfully\")\n",
        "else:\n",
        "    print(\"Failed to create pipeline\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PRyS-1HjcqV7"
      },
      "source": [
        "## Verify the pipeline\n",
        "Not required but nice to verify everything looks correct"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pH59icc9czcD"
      },
      "outputs": [],
      "source": [
        "pipeline = es.ingest.get_pipeline(id=\"jupyter-vector-demo-pipeline\")\n",
        "pipeline.body"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ruM78vW_hTOy"
      },
      "source": [
        "---\n",
        "---\n",
        "# Ingest Docs and Generate Vectors\n",
        "---\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RY3kxbN_hYGC"
      },
      "source": [
        "## Create sample documents\n",
        "These aren't real blogs just sampls ChatGPT created for me :) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaWpOz6nhr-L"
      },
      "outputs": [],
      "source": [
        "samples = [[\"The Power of Word Embeddings in NLP\", \"Word embeddings have revolutionized the field of NLP.\"  ],  \n",
        "    [\"An Introduction to Transformer Models\", \"Transformer models have taken NLP by storm.\"  ],  \n",
        "    [\"Fine-Tuning BERT for Text Classification\", \"Fine-tuning BERT can lead to state-of-the-art results in text classification.\"  ],  \n",
        "    [\"Why GPT-3 is a Game Changer for NLP\", \"GPT-3 has set a new standard for language models in NLP.\"  ],  \n",
        "    [\"Using ELMO for Sentiment Analysis\", \"ELMO can effectively capture contextual information for sentiment analysis.\"  ],  \n",
        "    [\"The Rise of Pre-Trained Models in NLP\", \"Pre-trained models have become increasingly popular in NLP.\"  ]\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OsbHqadHmif5"
      },
      "source": [
        "## Create the list of docs to ingest"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-NT3J6FGmcmR"
      },
      "outputs": [],
      "source": [
        "docs = [\n",
        "    {   \"_index\": \"jupyter-vector-demo\",\n",
        "        \"_source\": {\n",
        "           \"title\": sample[0], \n",
        "           \"summary\": sample[1]\n",
        "        }\n",
        "    }\n",
        "    for sample in samples\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WmBfO3N6nf37"
      },
      "source": [
        "## Index the docs \n",
        "This will send a bulk index request to elastic, sending all the docs through the ingest pipeline, generating vectors, and storing them in elasticsearch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Jyd_xOyo4xK"
      },
      "outputs": [],
      "source": [
        "helpers.bulk(es, docs, pipeline=\"jupyter-vector-demo-pipeline\" )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-xZ8SyBTpOYh"
      },
      "source": [
        "## Verify one of the docs \n",
        "Let's take a look at one doc and see how it was indexed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EasfA-uyqnUu"
      },
      "outputs": [],
      "source": [
        "result = es.search(index='jupyter-vector-demo', body={}, size=1)\n",
        "result.body['hits']['hits'][0]['_source']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8M-4Fc-zn7D_"
      },
      "source": [
        "---\n",
        "---\n",
        "# Approximate k-Nearest Neighbor Search (knn)\n",
        "---\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Hu2n4bmGYkG"
      },
      "source": [
        "## Generate Vector for Query\n",
        "\n",
        "Before we can run an approximate k-nearest neighbor (kNN) query, we need to convert our query string to a vector."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XXAQXIS0pvfR"
      },
      "source": [
        "Set a sample query doc\n",
        "\n",
        "Depending on your specific model, you may need to change the field name from \"text_field\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wBNV7q5Dwlz6"
      },
      "outputs": [],
      "source": [
        "docs =  [\n",
        "    {\n",
        "      \"text_field\": \"State of the art nlp models\"\n",
        "    }\n",
        "  ]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T5_BYcpvrLsb"
      },
      "source": [
        "We call the `_infer` endpoint supplying the model_id and the doc[s] we want to vectorize. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZsWg7XPSGbiu"
      },
      "outputs": [],
      "source": [
        "vec = MlClient.infer_trained_model(es, model_id=es_model_id, docs=docs, )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CdC3PkTyrZEV"
      },
      "source": [
        "The vector for the first doc can be accessed in the response dict as shown below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rle3J5mJXbdf"
      },
      "outputs": [],
      "source": [
        "doc_0_vector = vec['inference_results'][0]['predicted_value']\n",
        "doc_0_vector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pyeoQ8TK8Ddr"
      },
      "source": [
        "## Run the Search\n",
        "\n",
        "We will call the `_search` api and specify the `knn` section. \n",
        "\n",
        "This is a simple example of a search query. Elastic supports combining kNN search with \"traditional\" BM25 search. You can also filter documents to reduce the number of docs that needs to be searched. See the [docs](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search-api.html) for more information.\n",
        "\n",
        "This will be a very simple example to get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31c_Z4eL8mou"
      },
      "source": [
        "### Create the search paramaters\n",
        "Here we are just specifying the `knn` section, but you can also set all the other search params to pass."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_xI-NOjv9bbP"
      },
      "outputs": [],
      "source": [
        "knn = {\n",
        "    \"field\": \"vectors\",\n",
        "    \"query_vector\": doc_0_vector,\n",
        "    \"k\": 2,\n",
        "    \"num_candidates\": 10\n",
        "  }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXFxIpJxQTol"
      },
      "source": [
        "Send in the search request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bwQoXata8pZY"
      },
      "outputs": [],
      "source": [
        "results = es.search(index='jupyter-vector-demo', knn=knn, size=1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wmST0UL1Eqv2"
      },
      "source": [
        "### View the results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BAoKozJWEvto"
      },
      "outputs": [],
      "source": [
        "results.body"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mewh4SoLFBlM"
      },
      "source": [
        "### Search Query without returning vectors\n",
        "\n",
        "Often when running kNN search, you don't actually need to return the vectors themselves, you just want to return the fields to display to the end user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qY_xGJ6wQed4"
      },
      "source": [
        "We are adding a new paramater `fields` which allows us to specify specific fields to return rather than all of them\n",
        "\n",
        "By setting `source` to False (_source:false) we save having to get the entire source payoad back in the response\n",
        "\n",
        "We are moving the `size` value here simply to gather the paramaters together"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fgV6DIoCFBlM"
      },
      "outputs": [],
      "source": [
        "knn = {\n",
        "    \"field\": \"vectors\",\n",
        "    \"query_vector\": doc_0_vector,\n",
        "    \"k\": 2,\n",
        "    \"num_candidates\": 10\n",
        "  }\n",
        "fields = [\"summary\", \"title\"]\n",
        "size = 1\n",
        "source = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zl9ys10oFBlM"
      },
      "outputs": [],
      "source": [
        "results = es.search(index='jupyter-vector-demo', \n",
        "                    knn=knn, \n",
        "                    source=source, \n",
        "                    fields=fields, \n",
        "                    size=size\n",
        "                  )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sqHp2PMuFBlN"
      },
      "source": [
        "### View the results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZqgB31ZdFBlN"
      },
      "outputs": [],
      "source": [
        "results.body"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": " Working with Vectors in Elasticsearch",
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyMtQ854rmcWpdrrwxJjYMxu",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.4"
    },
    "vscode": {
      "interpreter": {
        "hash": "3ad933181bd8a04b432d3370b9dc3b0662ad032c4dfaa4e4f1596c548f763858"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}