examples/Search_reranking_using_embeddings.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJt4OFKPgO8C"
      },
      "source": [
        "# Search re-ranking using Gemini embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5n0sCc_fKcB"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/Search_reranking_using_embeddings.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0UO0Cgq0mRlG"
      },
      "source": [
        "This notebook demonstrates the use of embeddings to re-rank search results. This walkthrough will focus on the following objectives:\n",
        "\n",
        "\n",
        "\n",
        "1.   Setting up your development environment and API access to use Gemini.\n",
        "2.   Using Gemini's function calling support to access the Wikipedia API.\n",
        "3.   Embedding content via Gemini API.\n",
        "4.   Re-ranking the search results.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-e5Aje-DTQOq"
      },
      "source": [
        "This is how you will implement search re-ranking:\n",
        "\n",
        "\n",
        "1.   The user will make a search query.\n",
        "2.   You will use Wikipedia API to return the relevant search results.\n",
        "3.   The search results will be embedded and their relevance will be evaluated by calculating distance metrics like cosine similarity.\n",
        "4.   The most relevant search result will be returned as the final answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpVL_QrF5-oF"
      },
      "source": [
        "> The non-source code materials in this notebook are licensed under Creative Commons - Attribution-ShareAlike CC-BY-SA 4.0, https://creativecommons.org/licenses/by-sa/4.0/legalcode."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pwImlRg9DUp7"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nv_6f_qNHe_2"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5999faf240a4"
      },
      "outputs": [],
      "source": [
        "!pip install -q wikipedia"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "739f0bb73f05"
      },
      "source": [
        "Note: The [`wikipedia` package](https://pypi.org/project/wikipedia/) notes that it was \"designed for ease of use and simplicity, not for advanced use\", and that production or heavy use should instead \"use [Pywikipediabot](http://www.mediawiki.org/wiki/Manual:Pywikipediabot) or one of the other more advanced [Python MediaWiki API wrappers](http://en.wikipedia.org/wiki/Wikipedia:Creating_a_bot#Python)\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D0qOtLI3FVid"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import textwrap\n",
        "\n",
        "import google.generativeai as genai\n",
        "\n",
        "import wikipedia\n",
        "from wikipedia.exceptions import DisambiguationError, PageError\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "from IPython.display import Markdown\n",
        "\n",
        "def to_markdown(text):\n",
        "  text = text.replace('•', '  *')\n",
        "  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aeo_9J9iDtCx"
      },
      "source": [
        "To run the following cell, your API key must be stored it in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) quickstart for an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yxzEH0GqLshI"
      },
      "outputs": [],
      "source": [
        "from google.colab import userdata\n",
        "GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
        "genai.configure(api_key=GOOGLE_API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0hkJSFIWI7G1"
      },
      "source": [
        "## Define tools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HjN6hsAJkOq"
      },
      "source": [
        "As stated earlier, this tutorial uses Gemini's function calling support to access the Wikipedia API. Please refer to the [docs](https://ai.google.dev/docs/function_calling) to learn more about function calling."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fpFkazXUenkK"
      },
      "source": [
        "### Define the search function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2fp-m-5Ke0XY"
      },
      "source": [
        "To cater to the search engine needs, you will design this function in the following way:\n",
        "\n",
        "\n",
        "*   For each search query, the search engine will use the `wikipedia.search` method to get relevant topics.\n",
        "*   From the relevant topics, the engine will choose `n_topics(int)` top candidates and will use `gemini-pro` to extract relevant information from the page.\n",
        "*   The engine will avoid duplicate entries by maintaining a search history.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2U27tdP8OBys"
      },
      "outputs": [],
      "source": [
        "def wikipedia_search(search_queries: list[str]) -> list[str]:\n",
        "  \"\"\"Search wikipedia for each query and summarize relevant docs.\"\"\"\n",
        "  n_topics=3\n",
        "  search_history = set() # tracking search history\n",
        "  search_urls = []\n",
        "  mining_model = genai.GenerativeModel('gemini-pro')\n",
        "  summary_results = []\n",
        "\n",
        "  for query in search_queries:\n",
        "    print(f'Searching for \"{query}\"')\n",
        "    search_terms = wikipedia.search(query)\n",
        "\n",
        "    print(f\"Related search terms: {search_terms[:n_topics]}\")\n",
        "    for search_term in search_terms[:n_topics]: # select first `n_topics` candidates\n",
        "      if search_term in search_history: # check if the topic is already covered\n",
        "        continue\n",
        "\n",
        "      print(f'Fetching page: \"{search_term}\"')\n",
        "      search_history.add(search_term) # add to search history\n",
        "\n",
        "      try:\n",
        "        # extract the relevant data by using `gemini-pro` model\n",
        "        page = wikipedia.page(search_term, auto_suggest=False)\n",
        "        url = page.url\n",
        "        print(f\"Information Source: {url}\")\n",
        "        search_urls.append(url)\n",
        "        page = page.content\n",
        "        response = mining_model.generate_content(textwrap.dedent(f\"\"\"\\\n",
        "            Extract relevant information\n",
        "            about user's query: {query}\n",
        "            From this source:\n",
        "\n",
        "            {page}\n",
        "\n",
        "            Note: Do not summarize. Only Extract and return the relevant information\n",
        "        \"\"\"))\n",
        "\n",
        "        urls = [url]\n",
        "        if response.candidates[0].citation_metadata:\n",
        "          extra_citations = response.candidates[0].citation_metadata.citation_sources\n",
        "          extra_urls = [source.url for source in extra_citations]\n",
        "          urls.extend(extra_urls)\n",
        "          search_urls.extend(extra_urls)\n",
        "          print(\"Additional citations:\", response.candidates[0].citation_metadata.citation_sources)\n",
        "        try:\n",
        "          text = response.text\n",
        "        except ValueError:\n",
        "          pass\n",
        "        else:\n",
        "          summary_results.append(text + \"\\n\\nBased on:\\n  \" + ',\\n  '.join(urls))\n",
        "\n",
        "      except DisambiguationError:\n",
        "        print(f\"\"\"Results when searching for \"{search_term}\" (originally for \"{query}\")\n",
        "        were ambiguous, hence skipping\"\"\")\n",
        "\n",
        "      except PageError:\n",
        "        print(f'{search_term} did not match with any page id, hence skipping.')\n",
        "\n",
        "  print(f\"Information Sources:\")\n",
        "  for url in search_urls:\n",
        "    print('    ', url)\n",
        "\n",
        "  return summary_results\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EWKkkKDXmzOX"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"What are LLMs?\"\n",
            "Related search terms: ['Large language model', 'Prompt engineering', 'Gemini (language model)']\n",
            "Fetching page: \"Large language model\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Large_language_model\n",
            "Fetching page: \"Prompt engineering\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Prompt_engineering\n",
            "Fetching page: \"Gemini (language model)\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Gemini_(language_model)\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Large_language_model\n",
            "     https://en.wikipedia.org/wiki/Prompt_engineering\n",
            "     https://en.wikipedia.org/wiki/Gemini_(language_model)\n"
          ]
        }
      ],
      "source": [
        "example = wikipedia_search([\"What are LLMs?\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8bPT114QEPgW"
      },
      "source": [
        "Here is what the search results look like:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8wy7YUFqEYTv"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> **Definition:**\n",
              "> \n",
              "> * A large language model (LLM) is an artificial neural network that can perform natural language processing tasks such as text generation, classification, and question answering.\n",
              "> \n",
              "> **Training and Architecture:**\n",
              "> \n",
              "> * LLMs are trained on massive text datasets using self-supervised and semi-supervised learning techniques.\n",
              "> * They are typically built with transformer-based architectures and can have billions or even trillions of parameters.\n",
              "> * Prompt engineering is used to optimize the performance of LLMs on specific tasks.\n",
              "> \n",
              "> **Capabilities:**\n",
              "> \n",
              "> * **Text Generation:** LLMs can generate human-like text, including stories, articles, and code.\n",
              "> * **Language Understanding:** LLMs can understand the meaning and structure of text, including context and sentiment.\n",
              "> * **Question Answering:** LLMs can answer questions based on their knowledge of the text they have been trained on.\n",
              "> * **Natural Language Processing:** LLMs can perform various NLP tasks, such as sentiment analysis, machine translation, and named entity recognition.\n",
              "> \n",
              "> **Multimodality:**\n",
              "> \n",
              "> * LLMs can be extended to handle multiple modalities, such as images, audio, and video, by using tokenization methods.\n",
              "> * This enables them to perform tasks such as image captioning, video summarization, and audio transcription.\n",
              "> \n",
              "> **Properties:**\n",
              "> \n",
              "> * **Scaling Laws:** The performance of LLMs scales with the size of their training data, the number of parameters, and the computational power used for training.\n",
              "> * **Emergent Abilities:** LLMs exhibit emergent abilities that are not explicitly programmed into them, such as in-context learning and solving complex reasoning problems.\n",
              "> * **Interpretation:** Researchers are exploring methods to interpret the inner workings of LLMs and understand how they process and generate language.\n",
              "> \n",
              "> **Wider Impact:**\n",
              "> \n",
              "> * LLMs have potential applications in various industries, including healthcare, education, and customer service.\n",
              "> * They also raise concerns about copyright infringement, algorithmic bias, and the spread of misinformation.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Large_language_model"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/markdown": [
              "> **What are LLMs?**\n",
              "> \n",
              "> Large Language Models (LLMs) are a type of generative AI model that can be trained on large amounts of text data to understand and generate human-like language. They are used in a variety of applications, including:\n",
              "> \n",
              "> * **Text generation:** LLMs can generate text that is indistinguishable from human-written text. This can be used for a variety of purposes, such as writing articles, stories, and marketing copy.\n",
              "> * **Machine translation:** LLMs can be used to translate text from one language to another.\n",
              "> * **Question answering:** LLMs can be used to answer questions by searching through large amounts of text data.\n",
              "> * **Summarization:** LLMs can be used to summarize large amounts of text into shorter, more manageable summaries.\n",
              "> \n",
              "> LLMs are still under development, but they have the potential to revolutionize the way we interact with computers. They could make it possible for us to communicate with computers in a more natural way, and they could help us to access and understand information more easily.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Prompt_engineering"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/markdown": [
              "> - Large language models (LLMs) are advanced AI models that can process and generate human-like text.\n",
              "> - LLMs are trained on massive datasets of text and code, which allows them to understand and respond to a wide range of natural language queries.\n",
              "> - LLMs are used in a variety of applications, including search engines, chatbots, and machine translation.\n",
              "> - Gemini is a family of LLMs developed by Google DeepMind.\n",
              "> - Gemini comprises three models: Gemini Ultra, Gemini Pro, and Gemini Nano.\n",
              "> - Gemini Ultra is designed for \"highly complex tasks\", Gemini Pro is designed for \"a wide range of tasks\", and Gemini Nano is designed for \"on-device tasks\".\n",
              "> - Gemini was launched on December 6, 2023.\n",
              "> - Gemini is available through Google Cloud's Vertex AI service.\n",
              "> - Gemini is trained on and powered by Google's Tensor Processing Units (TPUs).\n",
              "> - Gemini has outperformed GPT-4, Anthropic's Claude 2, Inflection AI's Inflection-2, Meta's LLaMA 2, and xAI's Grok 1 on a variety of industry benchmarks.\n",
              "> - Gemini Pro is made available to Google Cloud customers on AI Studio and Vertex AI on December 13, while Gemini Nano will be made available to Android developers as well.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Gemini_(language_model)"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from IPython.display import display\n",
        "\n",
        "for e in example:\n",
        "  display(to_markdown(e))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SKcH1LICeg5Z"
      },
      "source": [
        "### Pass the tools to the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qQcxR7cjLjGM"
      },
      "source": [
        "If you pass a list of functions to the `GenerativeModel`'s `tools` argument,\n",
        "it will extract a schema from the function's signature and type hints, and then pass schema along to the API calls. In response the model may return a `FunctionCall` object asking to call the function.\n",
        "\n",
        "Note: This approach only handles annotations of `AllowedTypes = int | float | str | dict | list['AllowedTypes']`\n",
        "\n",
        "The `GenerativeModel` will keep a reference to the function inself, so that it _can_ execute the function locally later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39kqTnBRLDeQ"
      },
      "outputs": [],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    'gemini-pro',\n",
        "    tools=[wikipedia_search],\n",
        "    generation_config={'temperature': 0.6})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VedPbpzlh6jX"
      },
      "source": [
        "## Generate supporting search queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5sLKZik7isBW"
      },
      "source": [
        "In order to have multiple supporting search queries to the user's original query, you will ask the model to generate more such queries. This would help the engine to cover the asked question on comprehensive levels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z-Ym3H9KIosY"
      },
      "outputs": [],
      "source": [
        "instructions = \"\"\"You have access to the Wikipedia API which you will be using\n",
        "to answer a user's query. Your job is to generate a list of search queries which\n",
        "might answer a user's question. Be creative by using various key-phrases from\n",
        "the user's query. To generate variety of queries, ask questions which are\n",
        "related to  the user's query that might help to find the answer. The more\n",
        "queries you generate the better are the odds of you finding the correct answer.\n",
        "Here is an example:\n",
        "\n",
        "user: Tell me about Cricket World cup 2023 winners.\n",
        "\n",
        "function_call: wikipedia_search(['What is the name of the team that\n",
        "won the Cricket World Cup 2023?', 'Who was the captain of the Cricket World Cup\n",
        "2023 winning team?', 'Which country hosted the Cricket World Cup 2023?', 'What\n",
        "was the venue of the Cricket World Cup 2023 final match?', 'Cricket World cup 2023',\n",
        "'Who lifted the Cricket World Cup 2023 trophy?'])\n",
        "\n",
        "The search function will return a list of article summaries, use these to\n",
        "answer the  user's question.\n",
        "\n",
        "Here is the user's query: {query}\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Wyn3lV-d5S_"
      },
      "source": [
        "In order to yield creative and a more random variety of questions, you will set the model's temperature parameter to a value higher. Values can range from [0.0,1.0], inclusive. A value closer to 1.0 will produce responses that are more varied and creative, while a value closer to 0.0 will typically result in more straightforward responses from the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TD1vU6FXhYO5"
      },
      "source": [
        "## Enable automatic function calling and call the API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qyUSnXJ5hg6x"
      },
      "source": [
        "Now start a new chat with `enable_automatic_function_calling=True`. With it enabled, the `genai.ChatSession` will handle the back and forth required to call the function, and return the final response:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZNJqA60yKNpT"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"How do deep-sea creatures survive in the deep sea?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep-sea community', 'Deep sea']\n",
            "Fetching page: \"Deep-sea fish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "Fetching page: \"Deep-sea community\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "Fetching page: \"Deep sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep_sea\n",
            "Searching for \"What adaptations have deep-sea creatures developed to survive in the deep sea?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep-sea community', 'Deep sea']\n",
            "Searching for \"What is the pressure at the bottom of the ocean?\"\n",
            "Related search terms: ['DSV Limiting Factor', 'Titan submersible implosion', 'Seabed']\n",
            "Fetching page: \"DSV Limiting Factor\"\n",
            "Information Source: https://en.wikipedia.org/wiki/DSV_Limiting_Factor\n",
            "Fetching page: \"Titan submersible implosion\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Titan_submersible_implosion\n",
            "Fetching page: \"Seabed\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Seabed\n",
            "Searching for \"How do deep-sea creatures cope with the lack of sunlight?\"\n",
            "Related search terms: ['Deep-sea community', 'Deep sea', 'Deep-sea fish']\n",
            "Searching for \"What is the food chain like in the deep sea?\"\n",
            "Related search terms: ['Deep-sea community', 'Marine food web', 'Deep sea mining']\n",
            "Fetching page: \"Marine food web\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_food_web\n",
            "Fetching page: \"Deep sea mining\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep_sea_mining\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "     https://en.wikipedia.org/wiki/Deep_sea\n",
            "     https://en.wikipedia.org/wiki/DSV_Limiting_Factor\n",
            "     https://en.wikipedia.org/wiki/Titan_submersible_implosion\n",
            "     https://en.wikipedia.org/wiki/Seabed\n",
            "     https://en.wikipedia.org/wiki/Marine_food_web\n",
            "     https://en.wikipedia.org/wiki/Deep_sea_mining\n"
          ]
        }
      ],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    'gemini-pro', tools=[wikipedia_search], generation_config={'temperature': 0.6})\n",
        "\n",
        "chat = model.start_chat(enable_automatic_function_calling=True)\n",
        "\n",
        "query = \"Explain how deep-sea life survives.\"\n",
        "\n",
        "res = chat.send_message(instructions.format(query=query))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1l8KWb13M_lJ"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> Deep-sea life survives in the extreme conditions of the deep sea by developing various adaptations, including:\n",
              "> \n",
              "> * **Pressure Resistance:** Deep-sea organisms have evolved to withstand the crushing pressure of the deep sea, which can reach up to 1,000 times the pressure at sea level. Their bodies have a high internal pressure to balance the external pressure and prevent them from being crushed.\n",
              "> \n",
              "> * **Low Light Adaptation:** Many deep-sea fish have evolved large and sensitive eyes to maximize light absorption in the dark depths. Some species have multiple Rh1 opsin genes, which enhance their ability to see in low light conditions.\n",
              "> \n",
              "> * **Bioluminescence:** Many deep-sea creatures produce their own light through bioluminescent organs to attract prey, communicate, or camouflage themselves.\n",
              "> \n",
              "> * **Neutral Buoyancy:** Deep-sea fish often have low tissue density and reduced skeletal weight to maintain neutral buoyancy in the water column. Some species have swim bladders to help them adjust their buoyancy.\n",
              "> \n",
              "> * **Sensing Adaptations:** To navigate and locate prey in the dark, deep-sea creatures rely on highly sensitive senses such as touch, smell, and electroreception.\n",
              "> \n",
              "> * **Slow Metabolism and Energy Conservation:** Deep-sea organisms have adapted to survive on limited food sources. They have slower metabolisms and specialized feeding behaviors, such as lying in wait for prey or using lures to attract it."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_markdown(res.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ua7DPlj7HZyO"
      },
      "source": [
        "Check for additional citations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ASotq7EcHAeo"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'No citations found'"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "res.candidates[0].citation_metadata or 'No citations found'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A0CfKMjYFSQm"
      },
      "source": [
        "That looks like it worked. You can go through the chat history to see the details of what was sent and received in the function calls:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_9h2b8saNnxh"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "user -> {\n",
            "  \"text\": \"You have access to the Wikipedia API which you will be using\\nto answer a user's query. Your job is to generate a list of search queries which\\nmight answer a user's question. Be creative by using various key-phrases from\\nthe user's query. To generate variety of queries, ask questions which are\\nrelated to  the user's query that might help to find the answer. The more\\nqueries you generate the better are the odds of you finding the correct answer.\\nHere is an example:\\n\\nuser: Tell me about Cricket World cup 2023 winners.\\n\\nfunction_call: wikipedia_search(['What is the name of the team that\\nwon the Cricket World Cup 2023?', 'Who was the captain of the Cricket World Cup\\n2023 winning team?', 'Which country hosted the Cricket World Cup 2023?', 'What\\nwas the venue of the Cricket World Cup 2023 final match?', 'Cricket World cup 2023',\\n'Who lifted the Cricket World Cup 2023 trophy?'])\\n\\nThe search function will return a list of article summaries, use these to\\nanswer the  user's question.\\n\\nHere is the user's query: Explain how deep-sea life survives.\\n\"\n",
            "}\n",
            "------------------------------------------------------------\n",
            "model -> {\n",
            "  \"function_call\": {\n",
            "    \"name\": \"wikipedia_search\",\n",
            "    \"args\": {\n",
            "      \"search_queries\": [\n",
            "        \"How do deep-sea creatures survive in the deep sea?\",\n",
            "        \"What adaptations have deep-sea creatures developed to survive in the deep sea?\",\n",
            "        \"What is the pressure at the bottom of the ocean?\",\n",
            "        \"How do deep-sea creatures cope with the lack of sunlight?\",\n",
            "        \"What is the food chain like in the deep sea?\"\n",
            "      ]\n",
            "    }\n",
            "  }\n",
            "}\n",
            "------------------------------------------------------------\n",
            "user -> {\n",
            "  \"function_response\": {\n",
            "    \"name\": \"wikipedia_search\",\n",
            "    \"response\": {\n",
            "      \"result\": [\n",
            "        \"Deep-sea creatures have evolved various adaptations to survive in the extreme conditions of the deep sea.\\n\\n* **Pressure Resistance:** Deep-sea organisms have adapted to withstand the crushing pressure of the deep sea, which can reach up to 1,000 times the pressure at sea level. Their bodies have a high internal pressure to balance the external pressure and prevent them from being crushed.\\n\\n* **Low Light Adaptation:** Many deep-sea fish have evolved large and sensitive eyes to maximize light absorption in the dark depths. Some species have multiple Rh1 opsin genes, which enhance their ability to see in low light conditions.\\n\\n* **Bioluminescence:** Many deep-sea creatures produce their own light through bioluminescent organs to attract prey, communicate, or camouflage themselves.\\n\\n* **Neutral Buoyancy:** Deep-sea fish often have low tissue density and reduced skeletal weight to maintain neutral buoyancy in the water column. Some species have swim bladders to help them adjust their buoyancy.\\n\\n* **Sensing Adaptations:** To navigate and locate prey in the dark, deep-sea creatures rely on highly sensitive senses such as touch, smell, and electroreception.\\n\\n* **Slow Metabolism and Energy Conservation:** Deep-sea organisms have adapted to survive on limited food sources. They have slower metabolisms and specialized feeding behaviors, such as lying in wait for prey or using lures to attract it.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep-sea_fish\",\n",
            "        \"Deep-sea creatures have evolved adaptations to survive the extreme conditions of the deep sea, such as:\\n\\n* **Hyperbaricity:** Pressure increases by about one atmosphere every ten meters in the deep sea. Fish in these regions are usually small, with minimal skeletal structure and no swim bladders to cope with the pressure.\\n\\n* **Temperature:** Below the thermocline, the water mass of the deep ocean is cold and homogeneous. It decreases gradually from 5 or 6 \\u00b0C at 1,000 meters to a constant temperature below 3,000 to 4,000 meters.\\n\\n* **Salinity:** Salinity is relatively constant throughout the depths of the deep sea, except in regions such as the Mediterranean Sea, where evaporation causes high salinity, and in brine pools, which have higher salinity due to salt deposits.\\n\\n* **Energy sources:** Deep-sea creatures rely on various energy sources, including:\\n\\n    * **Marine snow:** Repackaged particles of organic matter that sink at a faster rate, providing food for organisms on the bottom.\\n    * **Whale falls:** Dead whales provide a significant amount of organic matter to the deep sea, supporting a variety of organisms in different stages of decomposition.\\n    * **Chemosynthesis:** Bacteria at hydrothermal vents and cold seeps use chemical energy to produce organic matter, providing a food source for vent-associated organisms.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep-sea_community\",\n",
            "        \"**How deep-sea creatures survive in the deep sea:**\\n\\n1. **Low temperatures**: Deep-sea organisms have adapted to survive in cold temperatures by developing unique physiological and metabolic mechanisms.\\n\\n2. **Darkness**: Many deep-sea creatures have developed large eyes and sensitive photoreceptors to make the most of the limited light available. Some species also use bioluminescence to attract prey or communicate with each other.\\n\\n3. **High pressure**: Deep-sea organisms have evolved adaptations to withstand the extreme pressure of the deep sea environment. Their bodies and cells have structural and biochemical adaptations that allow them to maintain their integrity under high pressure.\\n\\n4. **Limited food availability**: Deep-sea creatures have adapted to survive with limited food sources. They may rely on scavenging, predation, or filter feeding to obtain nutrients.\\n\\n5. **Unique adaptations**: Deep-sea creatures have evolved a variety of adaptations, such as gelatinous flesh, bioluminescence, and enlarged sensory organs, to survive in the harsh conditions of the deep sea.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep_sea\",\n",
            "        \"I'm sorry, but the provided text discusses deep-sea exploration and diving expeditions and does not contain any information about the pressure at the bottom of the ocean. Therefore, I cannot extract the requested data from the provided context.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/DSV_Limiting_Factor\",\n",
            "        \"Unfortunately, I'm unable to extract information about the pressure at the bottom of the ocean from the text you've provided.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Titan_submersible_implosion\",\n",
            "        \"The pressure at the bottom of the ocean varies depending on the depth. The deeper the water, the greater the pressure. At the bottom of the Mariana Trench, the deepest point in the ocean, the pressure is over 1,000 times greater than the pressure at sea level. This extreme pressure makes it difficult for most organisms to survive in the deepest parts of the ocean.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Seabed\",\n",
            "        \"**Food chain in the deep sea:**\\n\\nIn the deep sea, phytoplankton are the primary producers, followed by zooplankton, filter feeders, predatory vertebrates, and apex predators. This food chain is inverted compared to terrestrial environments, where primary producers (plants) have a larger biomass than consumers. \\n\\n**Importance of zooplankton:**\\nZooplankton form the second trophic level in the food chain and play a crucial role in marine ecosystems. They consume phytoplankton and are the primary food source for larger animals, such as fish, marine mammals, and seabirds.\\n\\n**Microbial loop:**\\nThe microbial loop is a process in which marine bacteria and archaea decompose organic matter and recycle nutrients back into the marine food web. This process helps maintain ocean productivity and supports higher trophic levels.\\n\\n**Cryptic interactions:**\\nCryptic interactions are hidden interactions within the marine food web that have been overlooked by traditional methods. These interactions, such as mixotrophy and auxotrophy, can have significant impacts on food web dynamics and model results.\\n\\n**Anthropogenic effects:**\\nHuman activities, such as overfishing and climate change, can disrupt marine food webs. Overfishing removes top predators, which can lead to trophic cascades and ecosystem imbalances. Climate change can alter ocean temperatures and acidification, which can negatively impact marine organisms and disrupt food web dynamics.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Marine_food_web\",\n",
            "        \"The food chain in the deep sea is not directly mentioned in the given context.\\n\\nBased on:\\n  https://en.wikipedia.org/wiki/Deep_sea_mining\"\n",
            "      ]\n",
            "    }\n",
            "  }\n",
            "}\n",
            "------------------------------------------------------------\n",
            "model -> {\n",
            "  \"text\": \"Deep-sea life survives in the extreme conditions of the deep sea by developing various adaptations, including:\\n\\n* **Pressure Resistance:** Deep-sea organisms have evolved to withstand the crushing pressure of the deep sea, which can reach up to 1,000 times the pressure at sea level. Their bodies have a high internal pressure to balance the external pressure and prevent them from being crushed.\\n\\n* **Low Light Adaptation:** Many deep-sea fish have evolved large and sensitive eyes to maximize light absorption in the dark depths. Some species have multiple Rh1 opsin genes, which enhance their ability to see in low light conditions.\\n\\n* **Bioluminescence:** Many deep-sea creatures produce their own light through bioluminescent organs to attract prey, communicate, or camouflage themselves.\\n\\n* **Neutral Buoyancy:** Deep-sea fish often have low tissue density and reduced skeletal weight to maintain neutral buoyancy in the water column. Some species have swim bladders to help them adjust their buoyancy.\\n\\n* **Sensing Adaptations:** To navigate and locate prey in the dark, deep-sea creatures rely on highly sensitive senses such as touch, smell, and electroreception.\\n\\n* **Slow Metabolism and Energy Conservation:** Deep-sea organisms have adapted to survive on limited food sources. They have slower metabolisms and specialized feeding behaviors, such as lying in wait for prey or using lures to attract it.\"\n",
            "}\n",
            "------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "for content in chat.history:\n",
        "  part = content.parts[0]\n",
        "\n",
        "  print(f'{content.role} -> ', end='')\n",
        "  print(json.dumps(type(part).to_dict(part), indent=2))\n",
        "  print('---' * 20)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WDWc9Fj9Ig6A"
      },
      "source": [
        "In the chat history you can see all 4 steps:\n",
        "\n",
        "1. The user sent the query.\n",
        "2. The model replied with a `glm.FunctionCall` calling the `wikipedia_search` with a number of relevant searches.\n",
        "3. Because you set `enable_automatic_function_calling=True` when creating the `genai.ChatSession`, it  executed the search function and returned the list of article summaries to the model.\n",
        "4. Folliwing the instructions in the prompt, the model generated a final answer based on those summaries.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJP1EAgfnPUA"
      },
      "source": [
        "## [Optional] Manually execute the function call"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sa5ke2ssKl7M"
      },
      "source": [
        "If you want to understand what happened behind the scenes, this section executes the `FunctionCall` manually to demonstrate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Q4p-zSP7ZBs"
      },
      "outputs": [],
      "source": [
        "import google.ai.generativelanguage as glm # for lower level code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wavbHrL3K5vo"
      },
      "outputs": [],
      "source": [
        "chat = model.start_chat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ON4LTcmiLs2E"
      },
      "outputs": [],
      "source": [
        "result = chat.send_message(instructions.format(query=query))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__JN3YuHe7fR"
      },
      "source": [
        "Initially the model returns a FunctionCall:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lgngdvcdi06F"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"name\": \"wikipedia_search\",\n",
            "  \"args\": {\n",
            "    \"search_queries\": [\n",
            "      \"How do deep-sea creatures survive the immense pressure of the deep sea?\",\n",
            "      \"What adaptations have deep-sea creatures evolved to withstand the extreme conditions of their environment?\",\n",
            "      \"How do deep-sea creatures generate their own light?\",\n",
            "      \"What is the role of bioluminescence in the deep sea?\",\n",
            "      \"How do deep-sea creatures find food in the darkness of the deep sea?\",\n",
            "      \"What are the unique challenges faced by deep-sea creatures?\",\n",
            "      \"How do deep-sea creatures reproduce in the extreme conditions of their environment?\"\n",
            "    ]\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "fc = result.candidates[0].content.parts[0].function_call\n",
        "fc = type(fc).to_dict(fc)\n",
        "print(json.dumps(fc, indent=2))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bfp1Kqv7PE8N"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'wikipedia_search'"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fc['name']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpD3axtDmfZb"
      },
      "source": [
        "Call the function with generated arguments to get the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ek4g-CTAPSou"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Searching for \"How do deep-sea creatures survive the immense pressure of the deep sea?\"\n",
            "Related search terms: ['Deep-sea community', 'Sea', 'Ocean']\n",
            "Fetching page: \"Deep-sea community\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "Fetching page: \"Sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Sea\n",
            "Fetching page: \"Ocean\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Ocean\n",
            "Searching for \"What adaptations have deep-sea creatures evolved to withstand the extreme conditions of their environment?\"\n",
            "Related search terms: ['Deep-sea fish', 'Astrobiology', 'Marine life']\n",
            "Fetching page: \"Deep-sea fish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "Fetching page: \"Astrobiology\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Astrobiology\n",
            "Fetching page: \"Marine life\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_life\n",
            "Searching for \"How do deep-sea creatures generate their own light?\"\n",
            "Related search terms: ['Bioluminescence', 'The Deep (TV series)', 'Leviathan']\n",
            "Fetching page: \"Bioluminescence\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Bioluminescence\n",
            "Fetching page: \"The Deep (TV series)\"\n",
            "Information Source: https://en.wikipedia.org/wiki/The_Deep_(TV_series)\n",
            "Fetching page: \"Leviathan\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Leviathan\n",
            "Searching for \"What is the role of bioluminescence in the deep sea?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep sea', 'Bioluminescence']\n",
            "Fetching page: \"Deep sea\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Deep_sea\n",
            "Searching for \"How do deep-sea creatures find food in the darkness of the deep sea?\"\n",
            "Related search terms: ['Deep-sea fish', 'Deep-sea community', 'Anglerfish']\n",
            "Fetching page: \"Anglerfish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Anglerfish\n",
            "Searching for \"What are the unique challenges faced by deep-sea creatures?\"\n",
            "Related search terms: ['Deep sea', 'List of legendary creatures by type', 'Marine biology']\n",
            "Fetching page: \"List of legendary creatures by type\"\n",
            "Information Source: https://en.wikipedia.org/wiki/List_of_legendary_creatures_by_type\n",
            "Fetching page: \"Marine biology\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Marine_biology\n",
            "Searching for \"How do deep-sea creatures reproduce in the extreme conditions of their environment?\"\n",
            "Related search terms: ['Jellyfish', 'Marine life', 'Microorganism']\n",
            "Fetching page: \"Jellyfish\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Jellyfish\n",
            "Fetching page: \"Microorganism\"\n",
            "Information Source: https://en.wikipedia.org/wiki/Microorganism\n",
            "Information Sources:\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_community\n",
            "     https://en.wikipedia.org/wiki/Sea\n",
            "     https://en.wikipedia.org/wiki/Ocean\n",
            "     https://en.wikipedia.org/wiki/Deep-sea_fish\n",
            "     https://en.wikipedia.org/wiki/Astrobiology\n",
            "     https://en.wikipedia.org/wiki/Marine_life\n",
            "     https://en.wikipedia.org/wiki/Bioluminescence\n",
            "     https://en.wikipedia.org/wiki/The_Deep_(TV_series)\n",
            "     https://en.wikipedia.org/wiki/Leviathan\n",
            "     https://en.wikipedia.org/wiki/Deep_sea\n",
            "     https://en.wikipedia.org/wiki/Anglerfish\n",
            "     https://en.wikipedia.org/wiki/List_of_legendary_creatures_by_type\n",
            "     https://en.wikipedia.org/wiki/Marine_biology\n",
            "     https://en.wikipedia.org/wiki/Jellyfish\n",
            "     https://en.wikipedia.org/wiki/Microorganism\n"
          ]
        }
      ],
      "source": [
        "summaries = wikipedia_search(**fc['args'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kv4WGnG_gT3F"
      },
      "source": [
        "Now send the `FunctionResult` to the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bcLuieHqj9PW"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> Deep-sea creatures have evolved unique adaptations to survive the immense pressure of the deep sea. These adaptations include having flexible bodies, using hydrostatic pressure to their advantage, and having specialized proteins that protect their cells from damage.\n",
              "> \n",
              "> Deep-sea creatures also have a number of adaptations that help them find food in the darkness of the deep sea. These adaptations include having large eyes that are sensitive to low levels of light, and having long, sensitive tentacles that can detect prey.\n",
              "> \n",
              "> Some deep-sea creatures also have bioluminescent organs that they use to attract prey or communicate with each other."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response = chat.send_message(\n",
        "    glm.Content(\n",
        "      parts=[glm.Part(\n",
        "          function_response = glm.FunctionResponse(\n",
        "            name='wikipedia_search',\n",
        "            response={'result': summaries}\n",
        "          )\n",
        "      )]\n",
        "    )\n",
        ")\n",
        "\n",
        "to_markdown(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2Vfv8xpmuV1"
      },
      "source": [
        "## Re-ranking the search results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cQ3JUJLeGGzA"
      },
      "source": [
        "Helper function to embed the content:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mSkE7EynJBwF"
      },
      "outputs": [],
      "source": [
        "def get_embeddings(content: list[str]) -> np.ndarray:\n",
        "  embeddings = genai.embed_content('models/embedding-001', content, 'SEMANTIC_SIMILARITY')\n",
        "  embds = embeddings.get('embedding', None)\n",
        "  embds = np.array(embds).reshape(len(embds), -1)\n",
        "  return embds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tip8ArqJf_ep"
      },
      "source": [
        "Please refer to the [embeddings guide](https://ai.google.dev/docs/embeddings_guide) for more information on embeddings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nSPyycFuFj-_"
      },
      "source": [
        "Your next step is to define functions that you can use to calculate similarity scores between two embedding vectors. These scores will help you decide which embedding vector is the most relevant vector to the user's query.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ltbB0vDsKQtI"
      },
      "source": [
        "You will now implement cosine similarity as your metric. Here returned embedding vectors will be of unit length and hence their L1 norm (`np.linalg.norm()`) will be ~1. Hence, calculating cosine similarity is esentially same as calculating their dot product score."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9iDFdzq_JWJW"
      },
      "outputs": [],
      "source": [
        "def dot_product(a: np.ndarray, b: np.ndarray):\n",
        "  return (a @ b.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MrF_1c_M_Hw3"
      },
      "source": [
        "### Similarity with user's query\n",
        "\n",
        "Now it's time to find the most relevant search result returned by the Wikipedia API.\n",
        "\n",
        "Use Gemini API to get embeddings for user's query and search results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gK9ryjftGDNe"
      },
      "outputs": [],
      "source": [
        "search_res = get_embeddings(summaries)\n",
        "embedded_query = get_embeddings([query])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wwWq30uGRG3"
      },
      "source": [
        "Calculate similarity score:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jWlFNYIsGV0X"
      },
      "outputs": [],
      "source": [
        "sim_value = dot_product(search_res, embedded_query)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJW1pQQXG2w2"
      },
      "source": [
        "using `np.argmax` best candidate is selected.\n",
        "\n",
        "**Users's Input:** Explain how deep-sea life survives.\n",
        "\n",
        "**Answer:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8vDMDnCsG8Wn"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Deep-sea creatures survive the immense pressure of the deep sea by evolving adaptations that allow them to withstand the extreme conditions. These adaptations include having flexible bodies, using hydrostatic pressure to their advantage, and having specialized proteins that protect their cells from damage.\n",
            "\n",
            "Based on:\n",
            "  https://en.wikipedia.org/wiki/Sea\n"
          ]
        }
      ],
      "source": [
        "print(summaries[np.argmax(sim_value)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ozn6mIFvoyJU"
      },
      "source": [
        "### Similarity with Hypothetical Document Embeddings (HyDE)\n",
        "\n",
        "Drawing inspiration from [Gao et al](https://arxiv.org/abs/2212.10496) the objective here is to generate a template answer to the user's query using `gemini-pro`'s internal knowledge. This hypothetical answer will serve as a baseline to calculate relevance of all the search results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J7m5KAkMREBH"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> In the abyssal depths of the ocean, where sunlight fades into eternal darkness and pressure increases with each descending fathom, life adapts to the extreme with remarkable resilience.\n",
              "> \n",
              "> At the core of this survival lies a unique set of adaptations that optimize energy conservation and maximize resource utilization. Organisms here harness the power of [unique chemical reactions] and [specialized organs] to extract sustenance from the meager organic matter that drifts down from the sunlit surface.\n",
              "> \n",
              "> The frigid temperatures and high-pressure environment necessitate a [slowed metabolism] that dramatically reduces energy expenditure. Organisms have evolved [flexible structures] that withstand the crushing forces, and [specialized enzymes] that function optimally in the cold and darkness.\n",
              "> \n",
              "> Many deep-sea creatures have developed [bioluminescent capabilities], creating their own light to attract mates or lure prey. These glowing displays illuminate the otherwise pitch-black surroundings, facilitating communication and survival.\n",
              "> \n",
              "> Moreover, deep-sea life exhibits [symbiotic relationships] with other organisms. Some species form partnerships with [bacteria that provide nourishment], while others establish alliances with [larger predators that offer protection]. These codependencies foster survival in an environment where resources are scarce.\n",
              "> \n",
              "> With their exceptional adaptations, deep-sea creatures defy the odds, thriving in an environment that would crush most other life forms. Their resilience and ingenuity serve as a testament to the boundless wonders that lie hidden in the unfathomable depths of our oceans."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "hypothetical_ans_model = genai.GenerativeModel('gemini-pro')\n",
        "res = hypothetical_ans_model.generate_content(f\"\"\"Generate a hypothetical answer\n",
        "to the user's query by using your own knowledge. Assume that you know everything\n",
        "about the said topic. Do not use factual information, instead use placeholders\n",
        "to complete your answer. Your answer should feel like it has been written by a human.\n",
        "\n",
        "query: {query}\"\"\")\n",
        "\n",
        "to_markdown(res.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85f4MpLgrYbW"
      },
      "source": [
        "Use Gemini API to get embeddings for the baseline answer and compare them with search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fSI6-5eYI3me"
      },
      "outputs": [],
      "source": [
        "hypothetical_ans = get_embeddings([res.text])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WxYGEWbqrh44"
      },
      "source": [
        "Calculate similarity scores to rank the search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99sjvMtlJfL_"
      },
      "outputs": [],
      "source": [
        "sim_value = dot_product(search_res, hypothetical_ans)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C32IhIapQFNa"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([[0.74144238],\n",
              "       [0.81955475],\n",
              "       [0.83515347],\n",
              "       [0.79930045],\n",
              "       [0.78609111],\n",
              "       [0.83174707],\n",
              "       [0.77690633],\n",
              "       [0.69279512],\n",
              "       [0.69031457],\n",
              "       [0.73286281],\n",
              "       [0.73241528],\n",
              "       [0.74342376],\n",
              "       [0.79973473],\n",
              "       [0.67441962],\n",
              "       [0.70701801]])"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sim_value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJ7pklxursSk"
      },
      "source": [
        "using `np.argmax` best candidate is selected.\n",
        "\n",
        "**Users's Input:** Explain how deep-sea life survives.\n",
        "\n",
        "**Answer:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzTfyU_mJ8-M"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> Deep-sea creatures can endure the immense pressure in the deep sea due to a variety of adaptations. These adaptations include:\n",
              "> \n",
              "> * **Having a body that is made up mostly of water.** This reduces the amount of pressure that the creature's body experiences.\n",
              "> * **Having a thick, tough skin that is resistant to pressure.** This helps to protect the creature's body from the crushing pressure of the water.\n",
              "> * **Having a skeleton that is made of a material that is stronger than bone.** This helps to support the creature's body and prevent it from collapsing under pressure.\n",
              "> * **Having a circulatory system that is adapted to high pressure.** This helps to ensure that the creature can get the oxygen and nutrients it needs to survive.\n",
              "> * **Having a metabolism that is slowed down.** This helps to conserve energy and prevent the creature from overheating.\n",
              "> \n",
              "> These adaptations allow deep-sea creatures to survive in the extreme conditions of the deep sea.\n",
              "> \n",
              "> Based on:\n",
              ">   https://en.wikipedia.org/wiki/Ocean"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_markdown(summaries[np.argmax(sim_value)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UjDN2OUpsL6M"
      },
      "source": [
        "You have now created a search re-ranking engine using embeddings!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AKH2vs2Lc1R1"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "I hope you found this example helpful! Check out more examples in the [Gemini Guide](https://github.com/google-gemini/gemini-guide/) to learn more."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Search_reranking_using_embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}