Search_with_Reranking.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCbHhd3GbCwN"
      },
      "source": [
        "# Improving Semantic Search with Rerankers\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/togethercomputer/together-cookbook/blob/main/Search_with_Reranking.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uUQXiVX3bCwQ"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "In this notebook we will use a reranker model to improve the results produced from a simple semantic search workflow. To get a better understanding of how semantic search works please refer to the [Cookbook here](https://github.com/togethercomputer/together-cookbook/blob/main/Semantic_Search.ipynb).\n",
        "\n",
        "A reranker model operates by looking at the query and the retrieved results from the semantic search pipeline one by one and assesses how relevant the returned result is to the query. Because the reranker model can spend compute assessing the query with the returned result at the same time it can better judge how relevant the words and meanings in the query are to individual documents. This also means that rerankers are computationally expensive and slower - thus they cannot be used to rank every document in our database.\n",
        "\n",
        "We run a semantic search process to obtain a list of 15-25 candidate objects that are similar \"enough\" to the query and then use the reranker as a fine-toothed comb to pick the top 5-10 objects that are actually closest to our query.\n",
        "\n",
        "We will be using the [Salesforce Llama Rank reranker model](https://blog.salesforceairesearch.com/llamarank/).\n",
        "\n",
        "<img src=\"images/reranking.png\" width=\"1000\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2kGNT6S6bCwR"
      },
      "source": [
        "### Install relevant libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RCWT_Xu6Z8fg",
        "outputId": "b1825654-f15f-4d90-d4a1-e53abf1542b4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting together\n",
            "  Downloading together-1.3.3-py3-none-any.whl.metadata (11 kB)\n",
            "Requirement already satisfied: aiohttp<4.0.0,>=3.9.3 in /usr/local/lib/python3.10/dist-packages (from together) (3.10.10)\n",
            "Requirement already satisfied: click<9.0.0,>=8.1.7 in /usr/local/lib/python3.10/dist-packages (from together) (8.1.7)\n",
            "Requirement already satisfied: eval-type-backport<0.3.0,>=0.1.3 in /usr/local/lib/python3.10/dist-packages (from together) (0.2.0)\n",
            "Requirement already satisfied: filelock<4.0.0,>=3.13.1 in /usr/local/lib/python3.10/dist-packages (from together) (3.16.1)\n",
            "Requirement already satisfied: numpy>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from together) (1.26.4)\n",
            "Requirement already satisfied: pillow<11.0.0,>=10.3.0 in /usr/local/lib/python3.10/dist-packages (from together) (10.4.0)\n",
            "Requirement already satisfied: pyarrow>=10.0.1 in /usr/local/lib/python3.10/dist-packages (from together) (16.1.0)\n",
            "Requirement already satisfied: pydantic<3.0.0,>=2.6.3 in /usr/local/lib/python3.10/dist-packages (from together) (2.9.2)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.31.0 in /usr/local/lib/python3.10/dist-packages (from together) (2.32.3)\n",
            "Requirement already satisfied: rich<14.0.0,>=13.8.1 in /usr/local/lib/python3.10/dist-packages (from together) (13.9.2)\n",
            "Requirement already satisfied: tabulate<0.10.0,>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from together) (0.9.0)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.66.2 in /usr/local/lib/python3.10/dist-packages (from together) (4.66.5)\n",
            "Requirement already satisfied: typer<0.13,>=0.9 in /usr/local/lib/python3.10/dist-packages (from together) (0.12.5)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (2.4.3)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (6.1.0)\n",
            "Requirement already satisfied: yarl<2.0,>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.15.2)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (4.0.3)\n",
            "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.23.4 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (2.23.4)\n",
            "Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (2024.8.30)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich<14.0.0,>=13.8.1->together) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich<14.0.0,>=13.8.1->together) (2.18.0)\n",
            "Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer<0.13,>=0.9->together) (1.5.4)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich<14.0.0,>=13.8.1->together) (0.1.2)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from yarl<2.0,>=1.12.0->aiohttp<4.0.0,>=3.9.3->together) (0.2.0)\n",
            "Downloading together-1.3.3-py3-none-any.whl (68 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.1/68.1 kB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: together\n",
            "Successfully installed together-1.3.3\n"
          ]
        }
      ],
      "source": [
        "!pip install together"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "UdsbAO-rbCwT"
      },
      "outputs": [],
      "source": [
        "import together, os\n",
        "\n",
        "# Paste in your Together AI API Key or load it\n",
        "TOGETHER_API_KEY = os.environ.get(\"TOGETHER_API_KEY\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VJaG0uX_bCwT"
      },
      "source": [
        "## Download and view the Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FuVvUlNEbCwU"
      },
      "outputs": [],
      "source": [
        "# Let's get the movies dataset\n",
        "!wget https://raw.githubusercontent.com/togethercomputer/together-cookbook/refs/heads/main/datasets/movies.json\n",
        "!mkdir datasets\n",
        "!mv movies.json datasets/movies.json"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3wysxjIvLZ3w",
        "outputId": "8171c42d-2e4d-415b-b8ac-420d6306b832"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[{'title': 'Terminator Genisys',\n",
              "  'overview': \"The year is 2029. John Connor, leader of the resistance continues the war against the machines. At the Los Angeles offensive, John's fears of the unknown future begin to emerge when TECOM spies reveal a new plot by SkyNet that will attack him from both fronts; past and future, and will ultimately change warfare forever.\",\n",
              "  'director': 'Alan Taylor',\n",
              "  'genres': 'Science Fiction Action Thriller Adventure',\n",
              "  'tagline': 'Reset the future'},\n",
              " {'title': 'Captain America: Civil War',\n",
              "  'overview': 'Following the events of Age of Ultron, the collective governments of the world pass an act designed to regulate all superhuman activity. This polarizes opinion amongst the Avengers, causing two factions to side with Iron Man or Captain America, which causes an epic battle between former allies.',\n",
              "  'director': 'Anthony Russo',\n",
              "  'genres': 'Adventure Action Science Fiction',\n",
              "  'tagline': 'Divided We Fall'},\n",
              " {'title': 'Whiplash',\n",
              "  'overview': 'Under the direction of a ruthless instructor, a talented young drummer begins to pursue perfection at any cost, even his humanity.',\n",
              "  'director': 'Damien Chazelle',\n",
              "  'genres': 'Drama',\n",
              "  'tagline': 'The road to greatness can take you to the edge.'}]"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import json\n",
        "\n",
        "with open('./datasets/movies.json', 'r') as file:\n",
        "    movies_data = json.load(file)\n",
        "\n",
        "movies_data[10:13]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLWI2iIqbCwV"
      },
      "source": [
        "## Implement Semantic Search Pipeline\n",
        "\n",
        "Below we implement a simple semantic search pipeline:\n",
        "1. Embed movie documents + query\n",
        "2. Obtain a list of movies ranked based on cosine similarities between the query and movie vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "0CR4khxlMb2P"
      },
      "outputs": [],
      "source": [
        "# This function will be used to access the Together API to generate embeddings for the movie plots\n",
        "\n",
        "from typing import List\n",
        "\n",
        "def generate_embeddings(input_texts: List[str], model_api_string: str) -> List[List[float]]:\n",
        "    \"\"\"Generate embeddings from Together python library.\n",
        "\n",
        "    Args:\n",
        "        input_texts: a list of string input texts.\n",
        "        model_api_string: str. An API string for a specific embedding model of your choice.\n",
        "\n",
        "    Returns:\n",
        "        embeddings_list: a list of embeddings. Each element corresponds to the each input text.\n",
        "    \"\"\"\n",
        "    together_client = together.Together(api_key = TOGETHER_API_KEY)\n",
        "    outputs = together_client.embeddings.create(\n",
        "        input=input_texts,\n",
        "        model=model_api_string,\n",
        "    )\n",
        "    return [x.embedding for x in outputs.data]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C3YCz6oYMo7G",
        "outputId": "160fa2dc-0834-4048-ee3b-c6e3d1272925"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['Minions Minions Stuart, Kevin and Bob are recruited by Scarlet Overkill, a super-villain who, alongside her inventor husband Herb, hatches a plot to take over the world. Before Gru, they had a history of bad bosses',\n",
              " 'Interstellar Interstellar chronicles the adventures of a group of explorers who make use of a newly discovered wormhole to surpass the limitations on human space travel and conquer the vast distances involved in an interstellar voyage. Mankind was born on Earth. It was never meant to die here.',\n",
              " 'Deadpool Deadpool tells the origin story of former Special Forces operative turned mercenary Wade Wilson, who after being subjected to a rogue experiment that leaves him with accelerated healing powers, adopts the alter ego Deadpool. Armed with his new abilities and a dark, twisted sense of humor, Deadpool hunts down the man who nearly destroyed his life. Witness the beginning of a happy ending',\n",
              " 'Guardians of the Galaxy Light years from Earth, 26 years after being abducted, Peter Quill finds himself the prime target of a manhunt after discovering an orb wanted by Ronan the Accuser. All heroes start somewhere.',\n",
              " \"Mad Max: Fury Road An apocalyptic story set in the furthest reaches of our planet, in a stark desert landscape where humanity is broken, and most everyone is crazed fighting for the necessities of life. Within this world exist two rebels on the run who just might be able to restore order. There's Max, a man of action and a man of few words, who seeks peace of mind following the loss of his wife and child in the aftermath of the chaos. And Furiosa, a woman of action and a woman who believes her path to survival may be achieved if she can make it across the desert back to her childhood homeland. What a Lovely Day.\",\n",
              " 'Jurassic World Twenty-two years after the events of Jurassic Park, Isla Nublar now features a fully functioning dinosaur theme park, Jurassic World, as originally envisioned by John Hammond. The park is open.',\n",
              " \"Pirates of the Caribbean: The Curse of the Black Pearl Jack Sparrow, a freewheeling 17th-century pirate who roams the Caribbean Sea, butts heads with a rival pirate bent on pillaging the village of Port Royal. When the governor's daughter is kidnapped, Sparrow decides to help the girl's love save her. But their seafaring mission is hardly simple. Prepare to be blown out of the water.\",\n",
              " 'Dawn of the Planet of the Apes A group of scientists in San Francisco struggle to stay alive in the aftermath of a plague that is wiping out humanity, while Caesar tries to maintain dominance over his community of intelligent apes. One last chance for peace.',\n",
              " 'The Hunger Games: Mockingjay - Part 1 Katniss Everdeen reluctantly becomes the symbol of a mass rebellion against the autocratic Capitol. Fire burns brighter in the darkness',\n",
              " 'Big Hero 6 The special bond that develops between plus-sized inflatable robot Baymax, and prodigy Hiro Hamada, who team up with a group of friends to form a band of high-tech heroes. From the creators of Wreck-it Ralph and Frozen']"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Concatenate the title, overview, and tagline of each movie\n",
        "# this makes the text that will be embedded for each movie more informative\n",
        "# as a result the embeddings will be richer and capture this information.\n",
        "to_embed = []\n",
        "for movie in movies_data[:1000]:\n",
        "    text = ''\n",
        "    for field in ['title', 'overview', 'tagline']:\n",
        "        value = movie.get(field, '')\n",
        "        text += str(value) + ' '\n",
        "    to_embed.append(text.strip())\n",
        "\n",
        "to_embed[:10]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "PPkb_O6UNKEB"
      },
      "outputs": [],
      "source": [
        "# Use bge-base-en-v1.5 model to generate embeddings\n",
        "embeddings = generate_embeddings(to_embed, 'BAAI/bge-base-en-v1.5')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K59L7q_2Noer",
        "outputId": "b3118011-87b1-4eb9-ed01-6b3bb4e61658"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "768"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# bge-base-en-v1.5 model generates 768-dimensional embeddings\n",
        "len(embeddings[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "BdRaOLoBQF3o"
      },
      "outputs": [],
      "source": [
        "# Generate the vector embeddings for the query\n",
        "query = \"super hero mystery action movie about bats\"\n",
        "\n",
        "query_embedding = generate_embeddings([query], 'BAAI/bge-base-en-v1.5')[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "pCq1sU8kQ8WP"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "\n",
        "# Calculate cosine similarity between the query embedding and each movie embedding\n",
        "similarity_scores = cosine_similarity([query_embedding], embeddings)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z7Xp_afERxHJ",
        "outputId": "04de64fd-5150-48e0-d66a-7c57cb98b8da"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(1, 1000)"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# We get a similarity score for each of our 1000 movies - the higher the score, the more similar the movie is to the query\n",
        "similarity_scores.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZAYFBYa_SB2i",
        "outputId": "8a17f9b5-1431-4fe7-c659-9fa14d8d8ab6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([0.40771197, 0.38049766, 0.41696283, 0.45355013, 0.42999108,\n",
              "       0.38060093, 0.45121266, 0.45003822, 0.3967806 , 0.48520776,\n",
              "       0.40464459, 0.42882582, 0.38652118, 0.57960979, 0.34677179,\n",
              "       0.46018015, 0.35667508, 0.54620629, 0.46283555, 0.38756721,\n",
              "       0.45288009, 0.43207754, 0.44227613, 0.41876672, 0.49631512,\n",
              "       0.41780368, 0.45989605, 0.34773829, 0.38279619, 0.44958772,\n",
              "       0.47529959, 0.33805641, 0.35754316, 0.55132521, 0.3921352 ,\n",
              "       0.45416451, 0.43980237, 0.40724228, 0.38742214, 0.38542083,\n",
              "       0.34656122, 0.30797263, 0.3550458 , 0.34403634, 0.39187938,\n",
              "       0.34535796, 0.30361464, 0.43121332, 0.36820356, 0.45294667])"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "similarity_scores[0, :50]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B04i7oqLSEv_",
        "outputId": "ffad70a0-c11e-4dee-8556-5ff423ff6ea3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([ 13, 265, 451,  33,  56,  17, 140, 450,  58, 828, 227,  62, 337,\n",
              "       172, 724, 424, 585, 696, 933, 996, 932, 433, 883, 420, 744, 407,\n",
              "       633, 775, 746, 723, 312, 119, 325, 688, 606, 400, 653, 647, 175,\n",
              "       655, 187, 613, 948, 580, 772,  24, 751, 835, 476, 961])"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Sort the similarity scores in descending order, obtain the index of the movies\n",
        "indices = np.argsort(-similarity_scores)\n",
        "\n",
        "indices[0, :50]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jIPNb3ZXSZZD",
        "outputId": "d014c8ba-7f4d-46fc-ec58-224526e82ab2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['The Dark Knight',\n",
              " 'Watchmen',\n",
              " 'Predator',\n",
              " 'Despicable Me 2',\n",
              " 'Night at the Museum: Secret of the Tomb',\n",
              " 'Batman v Superman: Dawn of Justice',\n",
              " 'Penguins of Madagascar',\n",
              " 'Batman & Robin',\n",
              " 'Batman Begins',\n",
              " 'Super 8',\n",
              " 'Megamind',\n",
              " 'The Dark Knight Rises',\n",
              " 'Batman Returns',\n",
              " 'The Incredibles',\n",
              " 'The Raid',\n",
              " 'Die Hard: With a Vengeance',\n",
              " 'Kick-Ass',\n",
              " 'Fantastic Mr. Fox',\n",
              " 'Commando',\n",
              " 'Tremors',\n",
              " 'The Peanuts Movie',\n",
              " 'Kung Fu Panda 2',\n",
              " 'Crank: High Voltage',\n",
              " 'Men in Black 3',\n",
              " 'ParaNorman']"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Get the top 25 movie titles that are most similar to the query - these will be passed to the reranker\n",
        "top_25_sorted_titles = [movies_data[index]['title'] for index in indices[0]][:25]\n",
        "\n",
        "top_25_sorted_titles"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "misJuVc8cRfA"
      },
      "source": [
        "Notice here that not all movies in our top 25 have to do with our query - `super hero mystery action movie about bats`. This is because semantic search capture the \"approximate\" meaning of the query and movies.\n",
        "\n",
        "The reranker can more closely determine the similarity between these 25 candidates and rerank which ones deserve to be atop our list."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ThfbpUDMbCwX"
      },
      "source": [
        "## Use Llama Rank to Rerank Top 25 Movies\n",
        "\n",
        "Treating the top 25 matching movies as good candidate matches, potentially with irrelevant false positives, that might have snuck in we want to have the reranker model look and rerank each based on similarity to the query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8cgjxMCNToZK",
        "outputId": "24a19c19-b6f8-4c80-b86b-3473ae2585eb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Document Index: 12\n",
            "Document: Batman Returns\n",
            "Relevance Score: 0.35380946383813044\n",
            "Document Index: 8\n",
            "Document: Batman Begins\n",
            "Relevance Score: 0.339339115127178\n",
            "Document Index: 7\n",
            "Document: Batman & Robin\n",
            "Relevance Score: 0.33013392395016167\n",
            "Document Index: 5\n",
            "Document: Batman v Superman: Dawn of Justice\n",
            "Relevance Score: 0.3289763252445171\n",
            "Document Index: 9\n",
            "Document: Super 8\n",
            "Relevance Score: 0.258483721657576\n"
          ]
        }
      ],
      "source": [
        "from together import Together\n",
        "\n",
        "client = Together(api_key = TOGETHER_API_KEY)\n",
        "\n",
        "query = \"super hero mystery action movie about bats\" # we keep the same query - can change if we want\n",
        "\n",
        "response = client.rerank.create(\n",
        "  model=\"Salesforce/Llama-Rank-V1\",\n",
        "  query=query,\n",
        "  documents=top_25_sorted_titles,\n",
        "  top_n=5 # we only want the top 5 results\n",
        ")\n",
        "\n",
        "for result in response.results:\n",
        "    print(f\"Document Index: {result.index}\")\n",
        "    print(f\"Document: {top_25_sorted_titles[result.index]}\")\n",
        "    print(f\"Relevance Score: {result.relevance_score}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "39cczZqybCwY"
      },
      "source": [
        "Here we can see that that reranker was able to improve the list by demoting irrelevant movies like `'Watchmen'`, `'Predator'`, `'Despicable Me 2'`, `'Night at the Museum: Secret of the Tomb'`, `'Penguins of Madagascar'`, further down the list and promoting `'Batman Returns'`, `'Batman Begins'`, `'Batman & Robin'`, `'Batman v Superman: Dawn of Justice'` to the top of the list!\n",
        "\n",
        "The `bge-base-en-v1.5` embedding model gives us a fuzzy match to concepts mentioned in the query, the `Llama-Rank-V1` reranker then imrpoves the quality of our list further by spending more compute to resort the list of movies.\n",
        "\n",
        "Learn more about how to use reranker models in the [docs](https://docs.together.ai/docs/rerank-overview) here!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}