{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TWbOzPWwDT_D"
   },
   "source": [
    "# Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9yWBrF5b3Nu9"
   },
   "outputs": [],
   "source": [
    "!pip install tdqm numpy faiss-cpu datasets requests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jD--Q-b4EGqw"
   },
   "source": [
    "# Load the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GtkNIvIY3NvB"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"sander-wood/wikimusictext\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DT26plJnETaj"
   },
   "source": [
    "# Start Jina Embeddings v2 Endpoint and Reranker Endpoint\n",
    "\n",
    "**IMPORTANT**: Please deploy the embedding and reranker endpoints in the [Azure portal](https://azure.microsoft.com/en-us/get-started/azure-portal). You will need to decide what **region** to use and assign a **DNS prefix** to the embedding service and another one to the reranker service. Then, add that information to the variables `embeddings_url` and `reranker_url` in the code below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OXTk7ASC3NvD"
   },
   "outputs": [],
   "source": [
    "\n",
    "import json\n",
    "\n",
    "import requests\n",
    "\n",
    "embeddings_url = \"http://<Insert here your DNS prefix>.<Insert here your region>.azurecontainer.io:8080/invocations\"\n",
    "reranker_url = \"http://<Insert here your DNS prefix>.<Insert here your region>.azurecontainer.io:8080/invocations\"\n",
    "\n",
    "def jina_embed(text):\n",
    "    headers = {\"Content-Type\": \"application/json\"}\n",
    "    json_data = {\"data\": [{\"text\": text}]}\n",
    "\n",
    "    response = requests.post(embeddings_url, headers=headers, data=json.dumps(json_data))\n",
    "    return response.json()[\"data\"][0][\"embedding\"]\n",
    "\n",
    "def jina_rerank(query, search_results):\n",
    "    headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "    json_data = {\n",
    "        \"data\": {\n",
    "            \"documents\": [\n",
    "                {\"text\": search_result[0]} for search_result in search_results\n",
    "            ],\n",
    "            \"query\": query,\n",
    "            \"top_n\": 3,\n",
    "        }\n",
    "    }\n",
    "\n",
    "    response = requests.post(reranker_url, headers=headers, data=json.dumps(json_data))\n",
    "    return response.json()[\"data\"][0][\"results\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A0rAFk4XEdfB"
   },
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zDU0Jkif3NvE"
   },
   "outputs": [],
   "source": [
    "ds = dataset[\"train\"]\n",
    "\n",
    "input_df = ds.to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uazlql8XEj3q"
   },
   "source": [
    "# Generate Embeddings and Index in FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mLUsPkk53NvE",
    "outputId": "313fd477-efaf-4e9b-9f7f-9a612eeea852"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding text chunks ...\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "tqdm.pandas()\n",
    "\n",
    "\n",
    "def generate_embeddings(input_df):\n",
    "    all_embeddings = []\n",
    "\n",
    "    for t in input_df.text:\n",
    "        review_embeddings = []\n",
    "        all_embeddings.append(np.array(jina_embed(t)))\n",
    "\n",
    "    input_df[\"embeddings\"] = all_embeddings\n",
    "\n",
    "    return input_df\n",
    "\n",
    "\n",
    "print(\"Embedding text chunks ...\")\n",
    "enhanced_dataframe = generate_embeddings(input_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 597
    },
    "id": "fm4nZGF-Zz9y",
    "outputId": "8e8a4219-697d-4662-a15b-d88dab281a56"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"input_df\",\n  \"rows\": 1010,\n  \"fields\": [\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1010,\n        \"samples\": [\n          \"Old Folks at Home\",\n          \"Suzy Snowflake\",\n          \"Que Sera, Sera.m\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"artist\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 828,\n        \"samples\": [\n          \"Dick Reynolds & Jack Rhodes\",\n          \"Kurt Weill, Marc Blitzstein\",\n          \"George Frideric Handel\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"genre\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 8,\n        \"samples\": [\n          \"Country\",\n          \"Rock\",\n          \"Jazz\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 993,\n        \"samples\": [\n          \"\\\"Three Kings of Orient\\\" was written by John Henry Hopkins Jr. in 1857. It was the first widely popular Christmas carol written in America. The carol centres around the Biblical Magi, who visited Jesus as a child in a house and gave him gifts.\",\n          \"\\\"Lovesick Blues\\\" is a Tin Pan Alley song, composed by Cliff Friend, with lyrics by Irving Mills. It first appeared in the 1922 musical \\\"Oh, Ernest\\\", and was recorded that year by Elsie Clark and Jack Shea. Emmett Miller recorded it in 1925 and 1928, followed by country music singer Rex Griffin in 1939. The recordings by Griffin and Miller inspired Hank Williams to perform the song during his first appearances on the Louisiana Hayride radio show.\",\n          \"\\\"Morning Has Broken\\\" is a Christian hymn first published in 1931. It has words by English author Eleanor Farjeon and was inspired by the village of Alfriston in East Sussex, then set to a traditional Scottish Gaelic tune, \\\"Bunessan\\\" It is often sung in children's services and in funeral services. Cat Stevens included a version on his album Teaser and the Firecat (1971) The song became identified with Stevens due to the popularity of this recording.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"embeddings\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe",
       "variable_name": "input_df"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-e851aed4-4d82-409a-a264-714ce248f7ea\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>artist</th>\n",
       "      <th>genre</th>\n",
       "      <th>text</th>\n",
       "      <th>embeddings</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>'Deed I Do</td>\n",
       "      <td>Fred Rose, Walter Hirsch</td>\n",
       "      <td>Jazz</td>\n",
       "      <td>\"Deed I Do\" is a 1926 jazz standard composed b...</td>\n",
       "      <td>[-0.21760905, -0.7139537, 0.0145562105, 0.6822...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>(Now And Then There's) A  fool such as I</td>\n",
       "      <td>Bill Trader, Bill Trader 1952</td>\n",
       "      <td>Country</td>\n",
       "      <td>\"Now and Then There's A Fool Such as I\" is a p...</td>\n",
       "      <td>[-0.41257474, -0.7050503, 0.28502792, 0.445489...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>(Remember Me) I'm The One Who Loves You</td>\n",
       "      <td>Stuart Hamblen</td>\n",
       "      <td>Folk</td>\n",
       "      <td>\"(Remember Me) I'm The One Who Loves You\" is ...</td>\n",
       "      <td>[-1.0425572, -0.28404456, 0.21446781, 0.321308...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>(Sittin' On) The Dock of the Bay</td>\n",
       "      <td>Otis Reading, Steve Cropper</td>\n",
       "      <td>R&amp;B</td>\n",
       "      <td>\"The Dock of the Bay\" is a song co-written by ...</td>\n",
       "      <td>[-0.6914139, -0.2128802, 0.91943, 0.41803297, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>(There'll be)  Peace in the Valley</td>\n",
       "      <td>Thomas A. Dorsey</td>\n",
       "      <td>R&amp;B</td>\n",
       "      <td>\"There'll Be Peace in the Valley for Me\" is a ...</td>\n",
       "      <td>[-0.6469957, -0.8650849, 0.52424073, 0.3500816...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1005</th>\n",
       "      <td>Younger Than Springtime</td>\n",
       "      <td>Richard Rodgers, Oscar Hammerstein II</td>\n",
       "      <td>Jazz</td>\n",
       "      <td>\"Younger Than Springtime\" is a show tune from ...</td>\n",
       "      <td>[-0.52188784, -0.30447546, -0.055954184, 0.304...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1006</th>\n",
       "      <td>Your Cheatin' Heart</td>\n",
       "      <td>Hank Williams</td>\n",
       "      <td>Country</td>\n",
       "      <td>\"Your Cheatin' Heart\" is a song written and re...</td>\n",
       "      <td>[-0.65648764, -0.6509074, 0.45983896, -0.16062...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1007</th>\n",
       "      <td>Your Song</td>\n",
       "      <td>Elton John</td>\n",
       "      <td>Rock</td>\n",
       "      <td>\"Your Song\" is a song written by Elton John an...</td>\n",
       "      <td>[-0.742592, -0.60632277, 0.3449046, 0.22714248...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1008</th>\n",
       "      <td>Yours Is My Heart Alone</td>\n",
       "      <td>Franz Lehar, Fritz Lohner, Harry B. Smith , Lu...</td>\n",
       "      <td>Pop</td>\n",
       "      <td>\"Yours Is My Heart Alone\" or \"You Are My Heart...</td>\n",
       "      <td>[-0.87446344, -0.52536905, 0.37290257, 0.18601...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1009</th>\n",
       "      <td>Zip-a-dee-doo-dah</td>\n",
       "      <td>Allie Wrubel, Ray Gilbert</td>\n",
       "      <td>R&amp;B</td>\n",
       "      <td>\"Zip-a-Dee-Doo-Dah\" is a song composed by Alli...</td>\n",
       "      <td>[-0.79368395, -0.50046057, 0.14534478, 0.19929...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1010 rows × 5 columns</p>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-e851aed4-4d82-409a-a264-714ce248f7ea')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-e851aed4-4d82-409a-a264-714ce248f7ea button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-e851aed4-4d82-409a-a264-714ce248f7ea');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-fedd7ea4-c2a1-479a-88c3-ec2faef4540f\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-fedd7ea4-c2a1-479a-88c3-ec2faef4540f')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-fedd7ea4-c2a1-479a-88c3-ec2faef4540f button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                          title  \\\n",
       "0                                    'Deed I Do   \n",
       "1     (Now And Then There's) A  fool such as I    \n",
       "2       (Remember Me) I'm The One Who Loves You   \n",
       "3              (Sittin' On) The Dock of the Bay   \n",
       "4            (There'll be)  Peace in the Valley   \n",
       "...                                         ...   \n",
       "1005                    Younger Than Springtime   \n",
       "1006                        Your Cheatin' Heart   \n",
       "1007                                  Your Song   \n",
       "1008                    Yours Is My Heart Alone   \n",
       "1009                          Zip-a-dee-doo-dah   \n",
       "\n",
       "                                                 artist    genre  \\\n",
       "0                              Fred Rose, Walter Hirsch     Jazz   \n",
       "1                         Bill Trader, Bill Trader 1952  Country   \n",
       "2                                        Stuart Hamblen     Folk   \n",
       "3                           Otis Reading, Steve Cropper      R&B   \n",
       "4                                      Thomas A. Dorsey      R&B   \n",
       "...                                                 ...      ...   \n",
       "1005              Richard Rodgers, Oscar Hammerstein II     Jazz   \n",
       "1006                                      Hank Williams  Country   \n",
       "1007                                         Elton John     Rock   \n",
       "1008  Franz Lehar, Fritz Lohner, Harry B. Smith , Lu...      Pop   \n",
       "1009                          Allie Wrubel, Ray Gilbert      R&B   \n",
       "\n",
       "                                                   text  \\\n",
       "0     \"Deed I Do\" is a 1926 jazz standard composed b...   \n",
       "1     \"Now and Then There's A Fool Such as I\" is a p...   \n",
       "2      \"(Remember Me) I'm The One Who Loves You\" is ...   \n",
       "3     \"The Dock of the Bay\" is a song co-written by ...   \n",
       "4     \"There'll Be Peace in the Valley for Me\" is a ...   \n",
       "...                                                 ...   \n",
       "1005  \"Younger Than Springtime\" is a show tune from ...   \n",
       "1006  \"Your Cheatin' Heart\" is a song written and re...   \n",
       "1007  \"Your Song\" is a song written by Elton John an...   \n",
       "1008  \"Yours Is My Heart Alone\" or \"You Are My Heart...   \n",
       "1009  \"Zip-a-Dee-Doo-Dah\" is a song composed by Alli...   \n",
       "\n",
       "                                             embeddings  \n",
       "0     [-0.21760905, -0.7139537, 0.0145562105, 0.6822...  \n",
       "1     [-0.41257474, -0.7050503, 0.28502792, 0.445489...  \n",
       "2     [-1.0425572, -0.28404456, 0.21446781, 0.321308...  \n",
       "3     [-0.6914139, -0.2128802, 0.91943, 0.41803297, ...  \n",
       "4     [-0.6469957, -0.8650849, 0.52424073, 0.3500816...  \n",
       "...                                                 ...  \n",
       "1005  [-0.52188784, -0.30447546, -0.055954184, 0.304...  \n",
       "1006  [-0.65648764, -0.6509074, 0.45983896, -0.16062...  \n",
       "1007  [-0.742592, -0.60632277, 0.3449046, 0.22714248...  \n",
       "1008  [-0.87446344, -0.52536905, 0.37290257, 0.18601...  \n",
       "1009  [-0.79368395, -0.50046057, 0.14534478, 0.19929...  \n",
       "\n",
       "[1010 rows x 5 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enhanced_dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3sq7llwO3NvF"
   },
   "outputs": [],
   "source": [
    "import faiss\n",
    "\n",
    "dim = 768  # dimension of Jina v2 embeddings\n",
    "index_with_ids = faiss.IndexIDMap(faiss.IndexFlatIP(dim))\n",
    "\n",
    "for idx, row in enhanced_dataframe.iterrows():\n",
    "    embeddings = row[\"embeddings\"]\n",
    "    normalized_embedding = np.ascontiguousarray(\n",
    "        np.array(embeddings, dtype=\"float32\").reshape(1, -1)\n",
    "    )\n",
    "    faiss.normalize_L2(normalized_embedding)\n",
    "    index_with_ids.add_with_ids(normalized_embedding, idx)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kBB3Pdf0Eunj"
   },
   "source": [
    "# Retrieve Matches to Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M6fwVBOC3NvG"
   },
   "outputs": [],
   "source": [
    "def find_similar_texts(query, n=20):\n",
    "    query_embedding = jina_embed(query)\n",
    "    query_embedding = np.ascontiguousarray(\n",
    "        np.array(query_embedding, dtype=\"float32\").reshape(1, -1)\n",
    "    )\n",
    "    faiss.normalize_L2(query_embedding)\n",
    "\n",
    "    similarities, indices = index_with_ids.search(query_embedding, n)\n",
    "\n",
    "    results = []\n",
    "    for i in range(n):\n",
    "        similarity = similarities[0][i]\n",
    "        index_id = indices[0][i]\n",
    "        results.append((enhanced_dataframe.loc[index_id, \"text\"], similarity))\n",
    "\n",
    "    return results\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m5kJ_74O3NvG"
   },
   "outputs": [],
   "source": [
    "query = \"What are some jazz songs that reached the top of the music charts in 1960s?\"\n",
    "search_results = find_similar_texts(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VaabIRUbbxEb",
    "outputId": "c212f250-be35-4218-e735-9cbbc680845b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('An instrumental version by Heywood and Hugo Winterhalter reached No. 2 on the Billboard Hot 100 chart and No. 7 on the R&B chart in 1956. A version sung by Andy Williams was also popular that year. The tune has been covered by a number of jazz performers beginning in the 1960s.',\n",
       "  0.84060496),\n",
       " ('\"The Sheik of Araby\" was written in 1921 by Harry B. Smith and Francis Wheeler, with music by Ted Snyder. The song was a Tin Pan Alley hit, and was also adopted by early jazz bands, especially in New Orleans, making it a jazz standard. The Beatles covered this song in 1962 at their unsuccessful Decca audition with George Harrison as the lead singer and Pete Best on the drums.',\n",
       "  0.80128336),\n",
       " ('\"C Jam Blues\" is a jazz standard composed in 1942 by Duke Ellington. The piece follows a twelve-bar blues form in the key of C major. It was also known as \"Duke\\'s Place\", with lyrics added by Bill Katts, Bob Thiele and Ruth Roberts.',\n",
       "  0.7998555),\n",
       " ('\"Baby Love\" is a song recorded by the American music group the Supremes in 1964 for their second studio album titled, Where Did Our Love Go. It was written and produced by Motown\\'s main production team Holland-Dozier-Holland. The song topped the Billboard pop singles chart in the United States from October 25, 1964 through November 21, 1964. It is considered one of the most popular songs of the late 20th century, and was ranked #324 on the Rolling Stone list of The 500 Greatest Songs of All Time.',\n",
       "  0.7993697),\n",
       " ('\"Take Five\" is a jazz standard composed by Paul Desmond and originally recorded by the Dave Brubeck Quartet for their album Time Out on July 1, 1959. Two years later it became a surprise hit and the biggest-selling jazz single ever. The single was inducted into the Grammy Hall of Fame in 1996. It became the first jazz single to surpass a million in sales.',\n",
       "  0.79927284),\n",
       " ('\"Blue Bayou\" is a song written by Roy Orbison and Joe Melson. It was originally sung and recorded by Orbison, who had an international hit with his version in 1963. It later became Linda Ronstadt\\'s signature song, with which she scored a Top 5 hit with her cover in 1977.',\n",
       "  0.7968584),\n",
       " ('\"Sun Arise\" is the fourth single released by Australian singer-songwriter Rolf Harris. Released in January 1961 in Australia and October 1962 in the UK, it was Harris\\' third charting hit in Australia (following \"The Big Black Hat\" in 1960) It was also his first hit in the United States, at no. 61 on the Billboard Hot 100. The Godfathers covered the song as one of their first singles, which was later compiled onto Hit by Hit.',\n",
       "  0.7938112),\n",
       " ('\"Sir Duke\" is a song composed and performed by Stevie Wonder from his 1976 album Songs in the Key of Life. Released as a single in 1977, the track topped the U.S. Billboard Hot 100 and Black Singles charts, and reached number two in the UK Singles Chart. The song was written in tribute to Duke Ellington, the influential jazz legend who had died in 1974.',\n",
       "  0.79362196),\n",
       " ('\"Sophisticated Lady\" is a jazz standard, composed as an instrumental in 1932 by Duke Ellington. The words were added to the song by Mitchell Parish. The recording entered the charts on 27 May 1933 and rose to number three.',\n",
       "  0.79240394),\n",
       " ('Oh, You Crazy Moon is a jazz standard by Jimmy Van Heusen, with lyrics by Johnny Burke. It was recorded by Mel Torme in 1960 and Frank Sinatra in 1965.',\n",
       "  0.79155463),\n",
       " ('Giant Steps is the fifth studio album by jazz musician John Coltrane as leader. It was released in February 1960 on Atlantic Records, catalogue SD 1311. It is considered one of the most influential jazz albums of all time. Many of its tracks have become practice templates for jazz saxophonists.',\n",
       "  0.7910377),\n",
       " ('\"Polka Dots and Moonbeams\" is a popular song with music by Jimmy Van Heusen and lyrics by Johnny Burke, published in 1940. It was Frank Sinatra\\'s first hit recorded with the Tommy Dorsey Orchestra. Bob Dylan covered this song in his 2016 album Fallen Angels.',\n",
       "  0.7900705),\n",
       " ('\"Jitterbug Waltz\" is a 1942 jazz composition by Fats Waller. It is one of the first jazz records recorded with a Hammond organ. The song is in the key of E major and in 3/4 time.',\n",
       "  0.7893182),\n",
       " ('\"I Didn\\'t Know About You\" is a song composed by Duke Ellington, with lyrics written by Bob Russell. The recording by Count Basie & His Orchestra (vocal by Thelma Carpenter) briefly reached the No. 21 position in the Billboard charts in 1945. It was based on an instrumental first recorded by Ellingon in 1942 under the title \"Sentimental Lady\"',\n",
       "  0.7890542),\n",
       " ('\"Just Squeeze Me (But Please Don\\'t Tease Me)\" is a 1941 popular song composed by Duke Ellington. The song has been recorded numerous times by a number of artists in the years since, having become a jazz standard.',\n",
       "  0.7887314),\n",
       " ('First recorded by Toots Thielemans in 1961, the song became an international hit. It has since been covered by over one hundred artists.',\n",
       "  0.78823656),\n",
       " ('\"Piano in the Dark\" is a song by American singer-songwriter Brenda Russell. It was the first single to be taken from Russell\\'s 1988 album, Get Here. The song became Russell\\'s biggest hit, peaking at number 6 on the Billboard Hot 100, number 8 on the R&B Chart and number 3 on the Adult Contemporary Chart.',\n",
       "  0.7882278),\n",
       " ('\"Bennie and the Jets\" is a song written by Elton John and songwriter Bernie Taupin. The song first appeared on the Goodbye Yellow Brick Road album in 1973. The track was a massive hit in the United States and Canada, released in 1974 as an A-side. In most territories the track was released as the B-side to \"Candle in the Wind\", but spelled \\'Benny\\' It is ranked number 371 on Rolling Stone\\'s list of The 500 Greatest Songs of All Time.',\n",
       "  0.78712404),\n",
       " ('\"Bernie\\'s Tune\" is a 1952 jazz standard. The music was written by Bernie Miller, with lyrics added later by Jerry Leiber and Mike Stoller. It was popularised with a recording by the quartet of the American saxophonist and composer Gerry Mulligan.',\n",
       "  0.78656244),\n",
       " ('\"I Got Rhythm\" is a piece composed by George Gershwin with lyrics by Ira GersHwin and published in 1930. Its chord progression is the foundation for many other popular jazz tunes such as Charlie Parker\\'s and Dizzy Gillespie\\'s bebop standard \"Anthropology (Thrivin\\' on a Riff)\". The song came from the musical Girl Crazy and has been sung by many jazz singers since.',\n",
       "  0.7855013)]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "search_results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yaz4YcNRb6wy"
   },
   "source": [
    "# Rerank to Get Most Relevant Matches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tArfE1PK1js7"
   },
   "outputs": [],
   "source": [
    "reranked_results = jina_rerank(query, search_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4y3N0g2qcCAY",
    "outputId": "98e74d98-c6f1-4006-ca1a-a973ef076016"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': 'c26a67d979cb73474e9f80221b14b5c9',\n",
       "  'index': 0,\n",
       "  'document': {'id': 'd2183fd857661fbf9ca60a22e91888a0',\n",
       "   'text': 'An instrumental version by Heywood and Hugo Winterhalter reached No. 2 on the Billboard Hot 100 chart and No. 7 on the R&B chart in 1956. A version sung by Andy Williams was also popular that year. The tune has been covered by a number of jazz performers beginning in the 1960s.'},\n",
       "  'relevance_score': 0.7132052183151245,\n",
       "  'usage': {'id': '037b9d22a5f13b68258ab51cbab1a7ad', 'total_tokens': 64}},\n",
       " {'id': 'a9205e69a4e76ca49717b8497a2798bf',\n",
       "  'index': 4,\n",
       "  'document': {'id': '25e78e92da17f01df111a7ed2716b057',\n",
       "   'text': '\"Take Five\" is a jazz standard composed by Paul Desmond and originally recorded by the Dave Brubeck Quartet for their album Time Out on July 1, 1959. Two years later it became a surprise hit and the biggest-selling jazz single ever. The single was inducted into the Grammy Hall of Fame in 1996. It became the first jazz single to surpass a million in sales.'},\n",
       "  'relevance_score': 0.204337015748024,\n",
       "  'usage': {'id': '6d55f32b339b83350ffb9489fbf31f5d', 'total_tokens': 80}},\n",
       " {'id': '50a610653b307f6f1ae6ec796b72ca83',\n",
       "  'index': 9,\n",
       "  'document': {'id': '70278633234c32775b1a28b364f6783a',\n",
       "   'text': 'Oh, You Crazy Moon is a jazz standard by Jimmy Van Heusen, with lyrics by Johnny Burke. It was recorded by Mel Torme in 1960 and Frank Sinatra in 1965.'},\n",
       "  'relevance_score': 0.16270869970321655,\n",
       "  'usage': {'id': '79eabc46bf3c659d3ad3e4d4d7e7a8f2', 'total_tokens': 40}}]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reranked_results"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}