notebooks/search/12-semantic-reranking-elastic-rerank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tn-NTlKI2EYR"
   },
   "source": [
    "# Semantic reranking with Elastic Rerank\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/notebooks/search/12-semantic-reranking-elastic-rerank.ipynb)\n",
    "\n",
    "In this notebook you'll learn how to implement semantic reranking in Elasticsearch using the built-in [Elastic Rerank model](https://www.elastic.co/guide/en/machine-learning/master/ml-nlp-rerank.html). You'll also learn about the `retriever` abstraction, a simpler syntax for crafting queries and combining different search operations.\n",
    "\n",
    "You will:\n",
    "\n",
    "- Create an inference endpoint to manage your `rerank` task. This will download and deploy the Elastic Rerank model.\n",
    "- Query your data using the `text_similarity_rerank` retriever, leveraging the Elastic Rerank model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-tir9w4Sz80v"
   },
   "source": [
    "## 🧰 Requirements\n",
    "\n",
    "For this example, you will need:\n",
    "\n",
    "- An Elastic deployment:\n",
    "\n",
    "  - We'll be using [Elastic Cloud](https://www.elastic.co/guide/en/cloud/current/ec-getting-started.html) for this example (available with a [free trial](https://cloud.elastic.co/registration?onboarding_token=vectorsearch&utm_source=github&utm_content=elasticsearch-labs-notebook))\n",
    "\n",
    "- Elasticsearch 8.17.0 or above, or [Elasticsearch serverless](https://www.elastic.co/elasticsearch/serverless)\n",
    "\n",
    "- A 4GB ML node\n",
    "\n",
    "> ℹ️ Deploying the Elastic Rerank model in combination with ELSER (or other hosted models) requires at minimum 8GB ML node. The current max size for trial ML nodes is 4GB (defaults to 1GB)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ut-7TfSB2EYS"
   },
   "source": [
    "## Install packages\n",
    "\n",
    "This will take a couple of minutes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AzuZInsj2EYS"
   },
   "outputs": [],
   "source": [
    "!pip install -qU elasticsearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CnEQCDrd2EYT"
   },
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bj8HebgN2EYT"
   },
   "outputs": [],
   "source": [
    "from elasticsearch import Elasticsearch, helpers, exceptions\n",
    "from urllib.request import urlopen\n",
    "from getpass import getpass\n",
    "import json\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FlJXqZkJ2EYT"
   },
   "source": [
    "## Initialize Elasticsearch Python client\n",
    "\n",
    "You need to connect to a running Elasticsearch instance. In this example we're using an Elastic Cloud deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "frMDUtzt2EYT",
    "outputId": "6339d445-e5ef-41e4-83f2-100b58875c27"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elastic Cloud ID: ··········\n",
      "Elastic Api Key: ··········\n"
     ]
    }
   ],
   "source": [
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n",
    "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n",
    "\n",
    "# Create the client instance\n",
    "client = Elasticsearch(\n",
    "    # For local development\n",
    "    # hosts=[\"http://localhost:9200\"]\n",
    "    cloud_id=ELASTIC_CLOUD_ID,\n",
    "    api_key=ELASTIC_API_KEY,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rL5AEaDrzj9t"
   },
   "source": [
    "## Test connection\n",
    "\n",
    "Confirm that the Python client has connected to your Elasticsearch instance with this test.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t_lhez7Tznhp"
   },
   "outputs": [],
   "source": [
    "print(client.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0Lt2KpOizHLe"
   },
   "source": [
    "## Enable Telemetry\n",
    "\n",
    "Knowing that you are using this notebook helps us decide where to invest our efforts to improve our products. We would like to ask you that you run the following code to let us gather anonymous usage statistics. See telemetry.py for details. Thank you!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "08qJdCH3zNSC"
   },
   "outputs": [],
   "source": [
    "!curl -O -s https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/telemetry/telemetry.py\n",
    "from telemetry import enable_telemetry\n",
    "\n",
    "es_client = enable_telemetry(client, \"12-semantic-reranking-elastic-rerank\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yR-zE2ii2EYU"
   },
   "source": [
    "## Upload sample data\n",
    "\n",
    "This examples uses a small dataset of movies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bWEpw_X52EYU",
    "outputId": "e4faa762-cb9d-487f-9055-54b2e4476a53"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done indexing documents into `movies` index!\n"
     ]
    }
   ],
   "source": [
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/notebooks/search/movies.json\"\n",
    "response = urlopen(url)\n",
    "\n",
    "# Load the response data into a JSON object\n",
    "data_json = json.loads(response.read())\n",
    "\n",
    "# Prepare the documents to be indexed\n",
    "documents = []\n",
    "for doc in data_json:\n",
    "    documents.append(\n",
    "        {\n",
    "            \"_index\": \"movies\",\n",
    "            \"_source\": doc,\n",
    "        }\n",
    "    )\n",
    "\n",
    "# Use helpers.bulk to index\n",
    "helpers.bulk(client, documents)\n",
    "\n",
    "print(\"Done indexing documents into `movies` index!\")\n",
    "time.sleep(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DRIABkGAgV_Q"
   },
   "source": [
    "## Create inference endpoint\n",
    "\n",
    "Next we'll create an inference endpoint for the `rerank` task to deploy and manage our model and, if necessary, spin up the necessary ML resources behind the scenes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DiKsd3YygV_Q",
    "outputId": "c3c46c6b-b502-4167-c98c-d2e2e0a4613c"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    client.inference.delete(inference_id=\"my-elastic-reranker\")\n",
    "except exceptions.NotFoundError:\n",
    "    # Inference endpoint does not exist\n",
    "    pass\n",
    "\n",
    "try:\n",
    "    client.options(\n",
    "        request_timeout=60, max_retries=3, retry_on_timeout=True\n",
    "    ).inference.put(\n",
    "        task_type=\"rerank\",\n",
    "        inference_id=\"my-elastic-reranker\",\n",
    "        inference_config={\n",
    "            \"service\": \"elasticsearch\",\n",
    "            \"service_settings\": {\n",
    "                \"model_id\": \".rerank-v1\",\n",
    "                \"num_threads\": 1,\n",
    "                \"adaptive_allocations\": {\n",
    "                    \"enabled\": True,\n",
    "                    \"min_number_of_allocations\": 1,\n",
    "                    \"max_number_of_allocations\": 4,\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "    )\n",
    "    print(\"Inference endpoint created successfully\")\n",
    "except exceptions.BadRequestError as e:\n",
    "    if e.error == \"resource_already_exists_exception\":\n",
    "        print(\"Inference endpoint created successfully\")\n",
    "    else:\n",
    "        raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following command to confirm your inference endpoint is deployed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.inference.get().body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "⚠️ When you deploy your model, you might need to sync your ML saved objects in the Kibana (or Serverless) UI.\n",
    "Go to **Trained Models** and select **Synchronize saved objects**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pDAB9pX3VGKE"
   },
   "source": [
    "## Lexical queries\n",
    "\n",
    "First let's use a `standard` retriever to test out some lexical (or full-text) searchs and then we'll compare the improvements when we layer in semantic reranking.\n",
    "\n",
    "### Lexical match with `query_string` query\n",
    "\n",
    "Let's say we vaguely remember that there is a famous movie about a killer who eats his victims. For the sake of argument, pretend we've momentarily forgotten the word \"cannibal\".\n",
    "\n",
    "Let's perform a [`query_string` query](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html) to find the phrase \"flesh-eating bad guy\" in the `plot` fields of our Elasticsearch documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EAG0yZ_bVX-j",
    "outputId": "a8446bbf-ece5-47db-c7cc-ce9856ab9811"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No search results found\n"
     ]
    }
   ],
   "source": [
    "resp = client.search(\n",
    "    index=\"movies\",\n",
    "    retriever={\n",
    "        \"standard\": {\n",
    "            \"query\": {\n",
    "                \"query_string\": {\n",
    "                    \"query\": \"flesh-eating bad guy\",\n",
    "                    \"default_field\": \"plot\",\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    },\n",
    ")\n",
    "\n",
    "if resp[\"hits\"][\"hits\"]:\n",
    "    for hit in resp[\"hits\"][\"hits\"]:\n",
    "        title = hit[\"_source\"][\"title\"]\n",
    "        plot = hit[\"_source\"][\"plot\"]\n",
    "        print(f\"Title: {title}\\nPlot: {plot}\\n\")\n",
    "else:\n",
    "    print(\"No search results found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "48_nf56d2EYV"
   },
   "source": [
    "No results! Unfortunately we don't have any near exact matches for \"flesh-eating bad guy\". Because we don't have any more specific information about the exact phrasing in the Elasticsearch data, we'll need to cast our search net wider.\n",
    "\n",
    "### Simple `match` query\n",
    "\n",
    "This lexical query performs a standard keyword search for the term \"crime\" within the \"plot\" and \"genre\" fields of our Elasticsearch documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XkAlzHW82EYV",
    "outputId": "f01c6ca5-f2fd-4cba-f5bf-7b8f6914eef6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Title: The Godfather\n",
      "Plot: An organized crime dynasty's aging patriarch transfers control of his clandestine empire to his reluctant son.\n",
      "\n",
      "Title: Goodfellas\n",
      "Plot: The story of Henry Hill and his life in the mob, covering his relationship with his wife Karen Hill and his mob partners Jimmy Conway and Tommy DeVito in the Italian-American crime syndicate.\n",
      "\n",
      "Title: The Silence of the Lambs\n",
      "Plot: A young F.B.I. cadet must receive the help of an incarcerated and manipulative cannibal killer to help catch another serial killer, a madman who skins his victims.\n",
      "\n",
      "Title: Pulp Fiction\n",
      "Plot: The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.\n",
      "\n",
      "Title: Se7en\n",
      "Plot: Two detectives, a rookie and a veteran, hunt a serial killer who uses the seven deadly sins as his motives.\n",
      "\n",
      "Title: The Departed\n",
      "Plot: An undercover cop and a mole in the police attempt to identify each other while infiltrating an Irish gang in South Boston.\n",
      "\n",
      "Title: The Usual Suspects\n",
      "Plot: A sole survivor tells of the twisty events leading up to a horrific gun battle on a boat, which began when five criminals met at a seemingly random police lineup.\n",
      "\n",
      "Title: The Dark Knight\n",
      "Plot: When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "resp = client.search(\n",
    "    index=\"movies\",\n",
    "    retriever={\n",
    "        \"standard\": {\n",
    "            \"query\": {\"multi_match\": {\"query\": \"crime\", \"fields\": [\"plot\", \"genre\"]}}\n",
    "        }\n",
    "    },\n",
    ")\n",
    "\n",
    "for hit in resp[\"hits\"][\"hits\"]:\n",
    "    title = hit[\"_source\"][\"title\"]\n",
    "    plot = hit[\"_source\"][\"plot\"]\n",
    "    print(f\"Title: {title}\\nPlot: {plot}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2K1C_Pmzmrxw"
   },
   "source": [
    "That's better! At least we've got some results now. We broadened our search criteria to increase the chances of finding relevant results.\n",
    "\n",
    "But these results aren't very precise in the context of our original query \"flesh-eating bad guy\". We can see that \"The Silence of the Lambs\" is returned in the middle of the results set with this generic `match` query. Let's see if we can use our semantic reranking model to get closer to the searcher's original intent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nd8gAE6H2EYV"
   },
   "source": [
    "## Semantic reranker\n",
    "\n",
    "In the following `retriever` syntax, we wrap our standard `match` query retriever in a `text_similarity_reranker`. This allows us to leverage the NLP model we deployed to Elasticsearch to rerank the results based on the phrase \"flesh-eating bad guy\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z9DegKqb2EYV",
    "outputId": "fc3f8827-4ddd-46d0-cccd-5376e0a81bbf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Title: The Silence of the Lambs\n",
      "Plot: A young F.B.I. cadet must receive the help of an incarcerated and manipulative cannibal killer to help catch another serial killer, a madman who skins his victims.\n",
      "\n",
      "Title: Pulp Fiction\n",
      "Plot: The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.\n",
      "\n",
      "Title: Se7en\n",
      "Plot: Two detectives, a rookie and a veteran, hunt a serial killer who uses the seven deadly sins as his motives.\n",
      "\n",
      "Title: Goodfellas\n",
      "Plot: The story of Henry Hill and his life in the mob, covering his relationship with his wife Karen Hill and his mob partners Jimmy Conway and Tommy DeVito in the Italian-American crime syndicate.\n",
      "\n",
      "Title: The Dark Knight\n",
      "Plot: When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.\n",
      "\n",
      "Title: The Godfather\n",
      "Plot: An organized crime dynasty's aging patriarch transfers control of his clandestine empire to his reluctant son.\n",
      "\n",
      "Title: The Departed\n",
      "Plot: An undercover cop and a mole in the police attempt to identify each other while infiltrating an Irish gang in South Boston.\n",
      "\n",
      "Title: The Usual Suspects\n",
      "Plot: A sole survivor tells of the twisty events leading up to a horrific gun battle on a boat, which began when five criminals met at a seemingly random police lineup.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "resp = client.search(\n",
    "    index=\"movies\",\n",
    "    retriever={\n",
    "        \"text_similarity_reranker\": {\n",
    "            \"retriever\": {\n",
    "                \"standard\": {\n",
    "                    \"query\": {\n",
    "                        \"multi_match\": {\"query\": \"crime\", \"fields\": [\"plot\", \"genre\"]}\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"field\": \"plot\",\n",
    "            \"inference_id\": \"my-elastic-reranker\",\n",
    "            \"inference_text\": \"flesh-eating bad guy\",\n",
    "        }\n",
    "    },\n",
    ")\n",
    "\n",
    "for hit in resp[\"hits\"][\"hits\"]:\n",
    "    title = hit[\"_source\"][\"title\"]\n",
    "    plot = hit[\"_source\"][\"plot\"]\n",
    "    print(f\"Title: {title}\\nPlot: {plot}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LqVh3qt82EYW"
   },
   "source": [
    "Success! \"The Silence of the Lambs\" is our top result. Semantic reranking helped us find the most relevant result by parsing a natural language query, overcoming the limitations of lexical search that relies on keyword matching.\n",
    "\n",
    "Semantic reranking enables semantic search in a few steps, without the need for generating and storing embeddings. This a great tool for testing and building hybrid search systems in Elasticsearch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6yYfJjjstxwK"
   },
   "source": [
    "## Learn more\n",
    "\n",
    "- [Elastic Rerank](https://www.elastic.co/guide/en/machine-learning/master/ml-nlp-rerank.html). Learn more about the Elastic Rerank model, including information about how the model is trained and how to deploy it in different environments.\n",
    "- [Semantic Reranking overview](https://www.elastic.co/guide/en/elasticsearch/reference/current/semantic-reranking.html). A high level overview of semantic reranking in Elasticsearch.\n",
    "- [`text_similarity_reranker` retriever ](https://www.elastic.co/guide/en/elasticsearch/reference/current/retriever.html#text-similarity-reranker-retriever). Detailed API syntax reference with examples.\n",
    "- [`elasticsearch-labs` notebooks](https://github.com/elastic/elasticsearch-labs/tree/main/notebooks). Check out our full catalogue of Python notebooks."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}