load_sentence_transformer_from_hf_to_elastic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": " Load Sentence Tranfromer from Hugging Face into Elasticsearch",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "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/load_sentence_transformer_from_hf_to_elastic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Loading a Sentense Tranformer from Hugging Face into Elasticsearch\n",
        "\n",
        "This code will show you how to load a supported embedding model from Hugging Face into an elasticsearch cluster in [Elastic Cloud](https://cloud.elastic.co/)\n",
        "\n",
        "First we will configure our python environment\n",
        "\n",
        "1. Set up our python environment\n",
        "2. Load the chosed model from HF into Elasticsearch\n",
        "3. Deploy and start the model on ML nodes\n",
        "4. Generate a test vector\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)."
      ],
      "metadata": {
        "id": "6xoLDtS_6Df1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n",
        "This section will set up the python environment with the required libraries"
      ],
      "metadata": {
        "id": "DgxCKQS7mCZw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install and import required python libraries"
      ],
      "metadata": {
        "id": "Ly1f1P-l9ri8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Elastic uses the [eland python library](https://github.com/elastic/eland) to download modesl from Hugging Face hub and load them into elasticsearch"
      ],
      "metadata": {
        "id": "MJAb_8zlPFhQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rUedSzQW9FIF"
      },
      "outputs": [],
      "source": [
        "pip install eland"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install elasticsearch"
      ],
      "metadata": {
        "id": "NK3Wx1I199yB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install transformers"
      ],
      "metadata": {
        "id": "cEfiiFXakzdP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install sentence_transformers"
      ],
      "metadata": {
        "id": "I20mDmJboKZw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install torch==1.11"
      ],
      "metadata": {
        "id": "uqcpWrbkBEB9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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\n",
        "from elasticsearch.client import MlClient"
      ],
      "metadata": {
        "id": "wyUZXUi4RWWL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "r7nMIbHke37Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import getpass"
      ],
      "metadata": {
        "id": "Xsd2m7HoTCLm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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:  ')"
      ],
      "metadata": {
        "id": "SSGgYHome69o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Connect to Elastic Cloud"
      ],
      "metadata": {
        "id": "jL4VDnVp96lf"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "I8mVJkKmetXo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Loading a model from Hugging Face model hub"
      ],
      "metadata": {
        "id": "pOuF_1VnmVU-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The Elastic Stack machine learning features support transformer models that conform to the standard BERT model interface and use the WordPiece tokenization algorithm.\n",
        "\n",
        "A current list of supported architectures can be viewed in the\n",
        "[Elastic NLP Model Support Docs](https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-model-ref.html)\n"
      ],
      "metadata": {
        "id": "NwIItJyhoWeT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download an embedding model from Hugging Face using the HF copy link\n",
        "\n",
        "For this example we will be using the [sentence-transformers/msmarco-MiniLM-L-12-v3](https://huggingface.co/sentence-transformers/msmarco-MiniLM-L-12-v3) model\n"
      ],
      "metadata": {
        "id": "10VvWJ87alld"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download the Model\n",
        "Here we specify the model id from Hugging Face. The easiest way to get this id is clicking the copy the model name icon next to the name on the model page. \n",
        "\n",
        "When calling `TransformerModel` you specify the HF model id and the task type. You can try specifying `auto` and eland will attempt to determine the correct type from info in the model config. This is not always possible so a list of specific `task_type` values can be viewed in the following code: \n",
        "[Supported values](https://github.com/elastic/eland/blob/15a300728876022b206161d71055c67b500a0192/eland/ml/pytorch/transformers.py#*L41*)"
      ],
      "metadata": {
        "id": "uBMWHj-ZmtvE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "hf_model_id='sentence-transformers/msmarco-MiniLM-L-12-v3'\n",
        "tm = TransformerModel(hf_model_id, \"text_embedding\")"
      ],
      "metadata": {
        "id": "zPV3oFsKiYFL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set and confirm the model ID\n",
        "To make the name compatible with elasticsearch, the '/' is replaced with '__'\n",
        "\n"
      ],
      "metadata": {
        "id": "sX-9dHuDmwgX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "es_model_id = tm.elasticsearch_model_id()\n",
        "es_model_id"
      ],
      "metadata": {
        "id": "XkIQBBCbdqvQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Export the model in a TorchScrpt representation which Elasticsearch uses"
      ],
      "metadata": {
        "id": "p0L2cfYwbIld"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tmp_path = \"models\"\n",
        "Path(tmp_path).mkdir(parents=True, exist_ok=True)\n",
        "model_path, config, vocab_path = tm.save(tmp_path)"
      ],
      "metadata": {
        "id": "GsSpvvP-nbCK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load the model into Elasticsearch\n",
        "Model should not already exist in elasticsearch"
      ],
      "metadata": {
        "id": "k1a_yNo6ba2E"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ptm = PyTorchModel(es, es_model_id)\n",
        "ptm.import_model(model_path=model_path, config_path=None, vocab_path=vocab_path, config=config) "
      ],
      "metadata": {
        "id": "Z4QD71Apnj4j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Starting the Model"
      ],
      "metadata": {
        "id": "4UYSzFp3vHdB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## View information about the model\n",
        "This is not required but can be handy to get a model overivew"
      ],
      "metadata": {
        "id": "wQwfozwznK4Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# List the in elasticsearch\n",
        "m = MlClient.get_trained_models(es, model_id=es_model_id)\n",
        "m.body"
      ],
      "metadata": {
        "id": "b4Wv8EJvpfZI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Deploy the model\n",
        "This will load the model on the ML nodes and start the process(es) making it available for the NLP task"
      ],
      "metadata": {
        "id": "oMGw3sk-pbaN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "s = MlClient.start_trained_model_deployment(es, model_id=es_model_id)\n",
        "s.body"
      ],
      "metadata": {
        "id": "w5muJ1rLqvUW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Verify the model started without issue"
      ],
      "metadata": {
        "id": "ZytlELrsnn_O"
      }
    },
    {
      "cell_type": "code",
      "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']"
      ],
      "metadata": {
        "id": "ZaQUUWe0Hxwz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "---\n",
        "# Working with Vectors\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "8M-4Fc-zn7D_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setting up the index in elasticsearch to store vectors\n",
        "\n"
      ],
      "metadata": {
        "id": "qpm3jgJOzlxi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generate Vector for Query\n",
        "\n",
        "Before we can run a kNN query, we need to convert our query string to a vector."
      ],
      "metadata": {
        "id": "6Hu2n4bmGYkG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Creating a sample query sentence"
      ],
      "metadata": {
        "id": "XXAQXIS0pvfR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "docs =  [\n",
        "    {\n",
        "      \"text_field\": \"Last week I upgraded my iOS version and ever since then my phone has been overheating whenever I use your app.\"\n",
        "    }\n",
        "  ]"
      ],
      "metadata": {
        "id": "wBNV7q5Dwlz6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We call the `_infer` endpoint supplying the model_id and the doc[s] we want to vectorize. "
      ],
      "metadata": {
        "id": "T5_BYcpvrLsb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "z = MlClient.infer_trained_model(es, model_id=es_model_id, docs=docs, )"
      ],
      "metadata": {
        "id": "ZsWg7XPSGbiu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The vector for the first doc can be accessed in the response dict as shown below"
      ],
      "metadata": {
        "id": "CdC3PkTyrZEV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "doc_0_vector = z['inference_results'][0]['predicted_value']\n",
        "doc_0_vector"
      ],
      "metadata": {
        "id": "Rle3J5mJXbdf"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}