notebooks/llmu/RAG_with_Chat_Embed_and_Rerank.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "ctaLvRUsfpj8",
      "metadata": {
        "id": "ctaLvRUsfpj8"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/RAG_with_Chat_Embed_and_Rerank.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61bac5b5",
      "metadata": {
        "id": "61bac5b5"
      },
      "source": [
        "# RAG with Chat, Embed, and Rerank\n",
        "\n",
        "This notebook shows how to build a RAG-powered chatbot with Cohere's Chat endpoint.  The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "Read the accompanying [article here](https://txt.cohere.com/rag-chatbot/).\n",
        "\n",
        "This application will use several Cohere API endpoints:\n",
        "\n",
        "- Chat: For handling the main logic of the chatbot, including turning a user message into queries, generating responses, and producing citations\n",
        "- Embed: For turning textual documents into their embeddings representation, later to be used in retrieval (we’ll use the latest, state-of-the-art Embed v3 model)\n",
        "- Rerank: For reranking the retrieved documents according to their relevance to a query\n",
        "\n",
        "The diagram below provides an overview of what we’ll build."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "33327522",
      "metadata": {},
      "source": [
        "![Workflow](../images/llmu/rag/rag-workflow-2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6ab2d5d",
      "metadata": {
        "id": "f6ab2d5d"
      },
      "source": [
        "Here is a summary of the steps involved.\n",
        "\n",
        "Initial phase:\n",
        "- **Step 0**: Ingest the documents – get documents, chunk, embed, and index.\n",
        "\n",
        "For each user-chatbot interaction:\n",
        "- **Step 1**: Get the user message\n",
        "- **Step 2**: Call the Chat endpoint in query-generation mode\n",
        "- If at least one query is generated\n",
        "  - **Step 3**: Retrieve and rerank relevant documents\n",
        "  - **Step 4**: Call the Chat endpoint in document mode to generate a grounded response with citations\n",
        "- If no query is generated\n",
        "  - **Step 4**: Call the Chat endpoint in normal mode to generate a response"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "TWyo_5WoNUM-",
      "metadata": {
        "id": "TWyo_5WoNUM-"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5pLAhQmTOKiV",
      "metadata": {
        "id": "5pLAhQmTOKiV"
      },
      "outputs": [],
      "source": [
        "! pip install cohere hnswlib unstructured -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "f3a03a57",
      "metadata": {
        "id": "f3a03a57"
      },
      "outputs": [],
      "source": [
        "import cohere\n",
        "import uuid\n",
        "import hnswlib\n",
        "from typing import List, Dict\n",
        "from unstructured.partition.html import partition_html\n",
        "from unstructured.chunking.title import chunk_by_title\n",
        "\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "Dx1cncziCWBB",
      "metadata": {
        "cellView": "form",
        "id": "Dx1cncziCWBB"
      },
      "outputs": [],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9d34e4b7",
      "metadata": {},
      "source": [
        "# Create a vector store for ingestion and retrieval\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "588ed6d0",
      "metadata": {},
      "source": [
        "![RAG components - Vectorstore](../images/llmu/rag/rag-components-vectorstore.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2f7e7d1c",
      "metadata": {
        "id": "2f7e7d1c"
      },
      "source": [
        "\n",
        "First, we define the list of documents we want to ingest and make available for retrieval. As an example, we'll use the contents from the first module of Cohere's *LLM University: What are Large Language Models?*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "3dca4a88",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "3dca4a88",
        "outputId": "b05da1ee-0456-4387-c232-a43e0ffed54c"
      },
      "outputs": [],
      "source": [
        "raw_documents = [\n",
        "    {\n",
        "        \"title\": \"Text Embeddings\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/text-embeddings\"},\n",
        "    {\n",
        "        \"title\": \"Similarity Between Words and Sentences\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/similarity-between-words-and-sentences\"},\n",
        "    {\n",
        "        \"title\": \"The Attention Mechanism\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/the-attention-mechanism\"},\n",
        "    {\n",
        "        \"title\": \"Transformer Models\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/transformer-models\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e2a8968",
      "metadata": {
        "id": "5e2a8968"
      },
      "source": [
        "Usually the number of documents for practical applications is vast, and so we'll need to be able to search documents efficiently.  This involves breaking the documents into chunks, generating embeddings, and indexing the embeddings, as shown in the image below.  \n",
        "\n",
        "We implement this in the `Vectorstore` class below, which takes the `raw_documents` list as input.  Three methods are immediately called when creating an object of the `Vectorstore` class:\n",
        "\n",
        "\n",
        "`load_and_chunk()`  \n",
        "This method uses the `partition_html()` method from the `unstructured` library to load the documents from URL and break them into smaller chunks.  Each chunk is turned into a dictionary object with three fields:\n",
        "- `title` - the web page’s title,\n",
        "- `text` - the textual content of the chunk, and\n",
        "- `url` - the web page’s URL.  \n",
        "  \n",
        "  \n",
        "`embed()`  \n",
        "This method uses Cohere's `embed-english-v3.0` model to generate embeddings of the chunked documents.  Since our documents will be used for retrieval, we set `input_type=\"search_document\"`.  We send the documents to the Embed endpoint in batches, because the endpoint has a limit of 96 documents per call.\n",
        "\n",
        "`index()`  \n",
        "This method uses the `hsnwlib` package to index the document chunk embeddings.  This will ensure efficient similarity search during retrieval.  Note that `hnswlib` uses a vector library, and we have chosen it for its simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "7c33412c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "7c33412c",
        "outputId": "cf04f8ed-8000-4433-f976-2d37747f21e7"
      },
      "outputs": [],
      "source": [
        "class Vectorstore:\n",
        "    \"\"\"\n",
        "    A class representing a collection of documents indexed into a vectorstore.\n",
        "\n",
        "    Parameters:\n",
        "    raw_documents (list): A list of dictionaries representing the sources of the raw documents. Each dictionary should have 'title' and 'url' keys.\n",
        "\n",
        "    Attributes:\n",
        "    raw_documents (list): A list of dictionaries representing the raw documents.\n",
        "    docs (list): A list of dictionaries representing the chunked documents, with 'title', 'text', and 'url' keys.\n",
        "    docs_embs (list): A list of the associated embeddings for the document chunks.\n",
        "    docs_len (int): The number of document chunks in the collection.\n",
        "    idx (hnswlib.Index): The index used for document retrieval.\n",
        "\n",
        "    Methods:\n",
        "    load_and_chunk(): Loads the data from the sources and partitions the HTML content into chunks.\n",
        "    embed(): Embeds the document chunks using the Cohere API.\n",
        "    index(): Indexes the document chunks for efficient retrieval.\n",
        "    retrieve(): Retrieves document chunks based on the given query.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, raw_documents: List[Dict[str, str]]):\n",
        "        self.raw_documents = raw_documents\n",
        "        self.docs = []\n",
        "        self.docs_embs = []\n",
        "        self.retrieve_top_k = 10\n",
        "        self.rerank_top_k = 3\n",
        "        self.load_and_chunk()\n",
        "        self.embed()\n",
        "        self.index()\n",
        "\n",
        "\n",
        "    def load_and_chunk(self) -> None:\n",
        "        \"\"\"\n",
        "        Loads the text from the sources and chunks the HTML content.\n",
        "        \"\"\"\n",
        "        print(\"Loading documents...\")\n",
        "\n",
        "        for raw_document in self.raw_documents:\n",
        "            elements = partition_html(url=raw_document[\"url\"])\n",
        "            chunks = chunk_by_title(elements)\n",
        "            for chunk in chunks:\n",
        "                self.docs.append(\n",
        "                    {\n",
        "                        \"title\": raw_document[\"title\"],\n",
        "                        \"text\": str(chunk),\n",
        "                        \"url\": raw_document[\"url\"],\n",
        "                    }\n",
        "                )\n",
        "\n",
        "    def embed(self) -> None:\n",
        "        \"\"\"\n",
        "        Embeds the document chunks using the Cohere API.\n",
        "        \"\"\"\n",
        "        print(\"Embedding document chunks...\")\n",
        "\n",
        "        batch_size = 90\n",
        "        self.docs_len = len(self.docs)\n",
        "        for i in range(0, self.docs_len, batch_size):\n",
        "            batch = self.docs[i : min(i + batch_size, self.docs_len)]\n",
        "            texts = [item[\"text\"] for item in batch]\n",
        "            docs_embs_batch = co.embed(\n",
        "                texts=texts, model=\"embed-english-v3.0\", input_type=\"search_document\"\n",
        "            ).embeddings\n",
        "            self.docs_embs.extend(docs_embs_batch)\n",
        "\n",
        "    def index(self) -> None:\n",
        "        \"\"\"\n",
        "        Indexes the document chunks for efficient retrieval.\n",
        "        \"\"\"\n",
        "        print(\"Indexing document chunks...\")\n",
        "\n",
        "        self.idx = hnswlib.Index(space=\"ip\", dim=1024)\n",
        "        self.idx.init_index(max_elements=self.docs_len, ef_construction=512, M=64)\n",
        "        self.idx.add_items(self.docs_embs, list(range(len(self.docs_embs))))\n",
        "\n",
        "        print(f\"Indexing complete with {self.idx.get_current_count()} document chunks.\")\n",
        "\n",
        "    def retrieve(self, query: str) -> List[Dict[str, str]]:\n",
        "        \"\"\"\n",
        "        Retrieves document chunks based on the given query.\n",
        "\n",
        "        Parameters:\n",
        "        query (str): The query to retrieve document chunks for.\n",
        "\n",
        "        Returns:\n",
        "        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.\n",
        "        \"\"\"\n",
        "\n",
        "        # Dense retrieval\n",
        "        query_emb = co.embed(\n",
        "            texts=[query], model=\"embed-english-v3.0\", input_type=\"search_query\"\n",
        "        ).embeddings\n",
        "\n",
        "        doc_ids = self.idx.knn_query(query_emb, k=self.retrieve_top_k)[0][0]\n",
        "\n",
        "        # Reranking\n",
        "        rank_fields = [\"title\", \"text\"] # We'll use the title and text fields for reranking\n",
        "\n",
        "        docs_to_rerank = [self.docs[doc_id] for doc_id in doc_ids]\n",
        "\n",
        "        rerank_results = co.rerank(\n",
        "            query=query,\n",
        "            documents=docs_to_rerank,\n",
        "            top_n=self.rerank_top_k,\n",
        "            model=\"rerank-english-v3.0\",\n",
        "            rank_fields=rank_fields\n",
        "        )\n",
        "\n",
        "        doc_ids_reranked = [doc_ids[result.index] for result in rerank_results.results]\n",
        "\n",
        "        docs_retrieved = []\n",
        "        for doc_id in doc_ids_reranked:\n",
        "            docs_retrieved.append(\n",
        "                {\n",
        "                    \"title\": self.docs[doc_id][\"title\"],\n",
        "                    \"text\": self.docs[doc_id][\"text\"],\n",
        "                    \"url\": self.docs[doc_id][\"url\"],\n",
        "                }\n",
        "            )\n",
        "\n",
        "        return docs_retrieved"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e1bf5d85",
      "metadata": {
        "id": "e1bf5d85"
      },
      "source": [
        "In the code cell below, we initialize an instance of the `Vectorstore` class and pass in the `raw_documents` list as input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "4643e630",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 178
        },
        "id": "4643e630",
        "outputId": "fe01fcb6-3574-4322-d8d0-57d37aad397d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading documents...\n",
            "Embedding document chunks...\n",
            "Indexing document chunks...\n",
            "Indexing complete with 134 document chunks.\n"
          ]
        }
      ],
      "source": [
        "# Create an instance of the Vectorstore class with the given sources\n",
        "vectorstore = Vectorstore(raw_documents)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61928287",
      "metadata": {
        "id": "61928287"
      },
      "source": [
        "The `Vectorstore` class also has a `retrieve()` method, which we'll use to retrieve relevant document chunks given a query (as in Step 3 in the diagram shared at the beginning of this notebook).  This method has two components: (1) dense retrieval, and (2) reranking.\n",
        "\n",
        "### Dense retrieval\n",
        "\n",
        "First, we embed the query using the same `embed-english-v3.0` model we used to embed the document chunks, but this time we set `input_type=\"search_query\"`.\n",
        "\n",
        "Search is performed by the `knn_query()` method from the `hnswlib` library. Given a query, it returns the document chunks most similar to the query. We can define the number of document chunks to return using the attribute `self.retrieve_top_k=10`.\n",
        "\n",
        "### Reranking\n",
        "\n",
        "After semantic search, we implement a reranking step.  While our semantic search component is already highly capable of retrieving relevant sources, the [Rerank endpoint](https://cohere.com/rerank) provides an additional boost to the quality of the search results, especially for complex and domain-specific queries. It takes the search results and sorts them according to their relevance to the query.\n",
        "\n",
        "We call the Rerank endpoint with the `co.rerank()` method and define the number of top reranked document chunks to retrieve using the attribute `self.rerank_top_k=3`.  The model we use is `rerank-english-v2.0`.  \n",
        "\n",
        "This method returns the top retrieved document chunks `chunks_retrieved` so that they can be passed to the chatbot.\n",
        "\n",
        "In the code cell below, we check the document chunks that are retrieved for the query `\"multi-head attention definition\"`."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "OwozNf_uPEyX",
      "metadata": {
        "id": "OwozNf_uPEyX"
      },
      "source": [
        "## Test Retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "82617b91",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        },
        "id": "82617b91",
        "outputId": "7f1f2bc8-8ed9-4190-bd6b-7af2d9dc1980"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[{'title': 'Transformer Models',\n",
              "  'text': 'The attention step used in transformer models is actually much more powerful, and it’s called multi-head attention. In multi-head attention, several different embeddings are used to modify the vectors and add context to them. Multi-head attention has helped language models reach much higher levels of efficacy when processing and generating text.',\n",
              "  'url': 'https://docs.cohere.com/docs/transformer-models'},\n",
              " {'title': 'The Attention Mechanism',\n",
              "  'text': \"What you learned in this chapter is simple self-attention. However, we can do much better than that. There is a method called multi-head attention, in which one doesn't only consider one embedding, but several different ones. These are all obtained from the original by transforming it in different ways. Multi-head attention has been very successful at the task of adding context to text. If you'd like to learn more about the self and multi-head attention, you can check out the following two\",\n",
              "  'url': 'https://docs.cohere.com/docs/the-attention-mechanism'},\n",
              " {'title': 'Transformer Models',\n",
              "  'text': 'Attention is a very useful technique that helps language models understand the context. In order to understand how attention works, consider the following two sentences:\\n\\nSentence 1: The bank of the river.\\n\\nSentence 2: Money in the bank.',\n",
              "  'url': 'https://docs.cohere.com/docs/transformer-models'}]"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vectorstore.retrieve(\"multi-head attention definition\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bae81baa",
      "metadata": {},
      "source": [
        "# Create a chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c88017f4",
      "metadata": {},
      "source": [
        "![RAG components - Chatbot](../images/llmu/rag/rag-components-chatbot.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e69fbca9",
      "metadata": {
        "id": "e69fbca9"
      },
      "source": [
        "\n",
        "Next, we implement a class to handle the interaction between the user and the chatbot.  It takes an instance of the `Vectorstore` class as input.\n",
        "\n",
        "The `run()` method will be used to run the chatbot application.  It begins with the logic for getting the user message, along with a way for the user to end the conversation.  \n",
        "\n",
        "Based on the user message, the chatbot needs to decide if it needs to consult external information before responding.  If so, the chatbot determines an optimal set of search queries to use for retrieval.  When we call `co.chat()` with `search_queries_only=True`, the Chat endpoint handles this for us automatically.\n",
        "\n",
        "The generated queries can be accessed from the `search_queries` field of the object that is returned.  Then, what happens next depends on how many queries are returned.\n",
        "- If queries are returned, we call the `retrieve()` method of the Vectorstore object for the  retrieval step.  The retrieved document chunks are then passed to the Chat endpoint by adding a `documents` parameter when we call `co.chat()` again.\n",
        "- Otherwise, if no queries are returned, we call the Chat endpoint another time, passing the user message and without needing to add any documents to the call.\n",
        "\n",
        "In either case, we also pass the `conversation_id` parameter, which retains the interactions between the user and the chatbot in the same conversation thread. We also enable the `stream` parameter so we can stream the chatbot response.\n",
        "\n",
        "We then print the chatbot's response.  In the case that the external information was used to generate a response, we also display citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "d2c15a1f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "d2c15a1f",
        "outputId": "8daa9159-338c-45ec-e9ed-830aedcdf0d8"
      },
      "outputs": [],
      "source": [
        "class Chatbot:\n",
        "    def __init__(self, vectorstore: Vectorstore):\n",
        "        \"\"\"\n",
        "        Initializes an instance of the Chatbot class.\n",
        "\n",
        "        Parameters:\n",
        "        vectorstore (Vectorstore): An instance of the Vectorstore class.\n",
        "\n",
        "        \"\"\"\n",
        "        self.vectorstore = vectorstore\n",
        "        self.conversation_id = str(uuid.uuid4())\n",
        "\n",
        "    def run(self):\n",
        "        \"\"\"\n",
        "        Runs the chatbot application.\n",
        "\n",
        "        \"\"\"\n",
        "        while True:\n",
        "            print(f\"\\n{'-'*100}\\n\")\n",
        "            # Get the user message\n",
        "            message = input(\"User: \")\n",
        "\n",
        "            # Typing \"quit\" ends the conversation\n",
        "            if message.lower() == \"quit\":\n",
        "              print(\"Ending chat.\")\n",
        "              break\n",
        "            # else:                       # If using Google Colab, remove this line to avoid printing the same thing twice\n",
        "            #   print(f\"User: {message}\") # If using Google Colab, remove this line to avoid printing the same thing twice\n",
        "\n",
        "            # Generate search queries (if any)\n",
        "            response = co.chat(message=message,\n",
        "                               model=\"command-r\",\n",
        "                               search_queries_only=True)\n",
        "\n",
        "            # If there are search queries, retrieve document chunks and respond\n",
        "            if response.search_queries:\n",
        "                print(\"Retrieving information...\", end=\"\")\n",
        "\n",
        "                # Retrieve document chunks for each query\n",
        "                documents = []\n",
        "                for query in response.search_queries:\n",
        "                    documents.extend(self.vectorstore.retrieve(query.text))\n",
        "\n",
        "                # Use document chunks to respond\n",
        "                response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r-plus\",\n",
        "                    documents=documents,\n",
        "                    conversation_id=self.conversation_id,\n",
        "                )\n",
        "\n",
        "            # If there is no search query, directly respond\n",
        "            else:\n",
        "                response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r-plus\",\n",
        "                    conversation_id=self.conversation_id,\n",
        "                )\n",
        "\n",
        "            # Print the chatbot response, citations, and documents\n",
        "            print(\"\\nChatbot:\")\n",
        "            citations = []\n",
        "            cited_documents = []\n",
        "\n",
        "            # Display response\n",
        "            for event in response:\n",
        "                if event.event_type == \"text-generation\":\n",
        "                    print(event.text, end=\"\")\n",
        "                elif event.event_type == \"citation-generation\":\n",
        "                    citations.extend(event.citations)\n",
        "                elif event.event_type == \"stream-end\":\n",
        "                    cited_documents = event.response.documents\n",
        "\n",
        "            # Display citations and source documents\n",
        "            if citations:\n",
        "              print(\"\\n\\nCITATIONS:\")\n",
        "              for citation in citations:\n",
        "                print(citation)\n",
        "\n",
        "              print(\"\\nDOCUMENTS:\")\n",
        "              for document in cited_documents:\n",
        "                print(document)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "F0X9FZqwOwQ6",
      "metadata": {
        "id": "F0X9FZqwOwQ6"
      },
      "source": [
        "# Run the chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3cb442f7",
      "metadata": {
        "id": "3cb442f7"
      },
      "source": [
        "We can now run the chatbot.  For this, we create the instance of `Chatbot` and run the chatbot by invoking the `run()` method.\n",
        "\n",
        "The format of each citation is:\n",
        "- `start`: The starting point of a span where one or more documents are referenced\n",
        "- `end`: The ending point of a span where one or more documents are referenced\n",
        "- `text`: The text representing this span\n",
        "- `document_ids`: The IDs of the documents being referenced (`doc_0` being the ID of the first document passed to the `documents` creating parameter in the endpoint call, and so on)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "42d3f345",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "42d3f345",
        "outputId": "8b935c8b-b1d4-4913-bdf8-73ba503402b8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "User: Hello, I have a question\n",
            "\n",
            "Chatbot:\n",
            "Of course! I am here to assist you. Please go ahead with your question, and I will do my best to provide a helpful response.\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "User: What’s the difference between word and sentence embeddings\n",
            "Retrieving information...\n",
            "Chatbot:\n",
            "Word embeddings are a way to associate words with lists of numbers (vectors) so that similar words are associated with numbers that are close by, and dissimilar words with numbers that are far away from each other. Sentence embeddings do the same thing but associate a vector to an entire sentence. Similar sentences are assigned similar vectors, and different sentences are assigned different vectors. Each coordinate in the vector identifies some property of the sentence.\n",
            "\n",
            "CITATIONS:\n",
            "start=0 end=15 text='Word embeddings' document_ids=['doc_0', 'doc_1']\n",
            "start=29 end=76 text='associate words with lists of numbers (vectors)' document_ids=['doc_1']\n",
            "start=85 end=144 text='similar words are associated with numbers that are close by' document_ids=['doc_1']\n",
            "start=150 end=214 text='dissimilar words with numbers that are far away from each other.' document_ids=['doc_1']\n",
            "start=215 end=234 text='Sentence embeddings' document_ids=['doc_0', 'doc_1', 'doc_2']\n",
            "start=257 end=298 text='associate a vector to an entire sentence.' document_ids=['doc_0', 'doc_1', 'doc_2']\n",
            "start=299 end=345 text='Similar sentences are assigned similar vectors' document_ids=['doc_0']\n",
            "start=351 end=402 text='different sentences are assigned different vectors.' document_ids=['doc_0']\n",
            "start=403 end=474 text='Each coordinate in the vector identifies some property of the sentence.' document_ids=['doc_0', 'doc_2']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'doc_0', 'text': 'This is where sentence embeddings come into play. A sentence embedding is just like a word embedding, except it associates every sentence with a vector full of numbers, in a coherent way. By coherent, I mean that it satisfies similar properties as a word embedding. For instance, similar sentences are assigned to similar vectors, different sentences are assigned to different vectors, and most importantly, each of the coordinates of the vector identifies some (whether clear or obscure) property of', 'title': 'Text Embeddings', 'url': 'https://docs.cohere.com/docs/text-embeddings'}\n",
            "{'id': 'doc_1', 'text': 'In the previous chapters, you learned about word and sentence embeddings and similarity between words and sentences. In short, a word embedding is a way to associate words with lists of numbers (vectors) in such a way that similar words are associated with numbers that are close by, and dissimilar words with numbers that are far away from each other. A sentence embedding does the same thing, but associating a vector to every sentence. Similarity is a way to measure how similar two words (or', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_2', 'text': 'Sentence embeddings are even more powerful, as they assign a vector of numbers to each sentence, in a way that these numbers also carry important properties of the sentence. One of the Cohere embeddings assigns a vector of length 4096 (i.e., a list of 4096 numbers) to each sentence. Furthermore, multilingual embedding does this for sentences in more than 100 languages. In this way, the sentence “Hello, how are you?” and its corresponding French translation, “Bonjour, comment ça va?” will be', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "User: And what are their similarities\n",
            "Retrieving information...\n",
            "Chatbot:\n",
            "Word and sentence embeddings both use vectors to represent words or sentences in a numerical space. The similarity between embeddings is calculated using a dot product. The similarity between an embedding and itself is 1, and the similarity between an irrelevant word or sentence and any other word or sentence is 0.\n",
            "\n",
            "CITATIONS:\n",
            "start=38 end=45 text='vectors' document_ids=['doc_0', 'doc_3', 'doc_4']\n",
            "start=59 end=77 text='words or sentences' document_ids=['doc_0', 'doc_3', 'doc_4']\n",
            "start=83 end=99 text='numerical space.' document_ids=['doc_0', 'doc_3', 'doc_4']\n",
            "start=104 end=114 text='similarity' document_ids=['doc_5']\n",
            "start=156 end=168 text='dot product.' document_ids=['doc_5']\n",
            "start=173 end=220 text='similarity between an embedding and itself is 1' document_ids=['doc_0', 'doc_4']\n",
            "start=230 end=316 text='similarity between an irrelevant word or sentence and any other word or sentence is 0.' document_ids=['doc_0', 'doc_4']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'doc_0', 'text': 'But let me add some numbers to this reasoning to make it more clear. Imagine that we calculate similarities for the words in each sentence, and we get the following:\\n\\nThis similarity makes sense in the following ways:\\n\\nThe similarity between each word and itself is 1.\\n\\nThe similarity between any irrelevant word (“the”, “of”, etc.) and any other word is 0.\\n\\nThe similarity between “bank” and “river” is 0.11.\\n\\nThe similarity between “bank” and “money” is 0.25.', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_3', 'text': 'assigned very similar numbers, as they have the same semantic meaning.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "{'id': 'doc_4', 'text': 'But let me add some numbers to this reasoning to make it more clear. Imagine that we calculate similarities for the words in each sentence, and we get the following:\\n\\nThis similarity makes sense in the following ways:\\n\\nThe similarity between each word and itself is 1.\\n\\nThe similarity between any irrelevant word (“the”, “of”, etc.) and any other word is 0.\\n\\nThe similarity between “bank” and “river” is 0.11.\\n\\nThe similarity between “bank” and “money” is 0.25.', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_5', 'text': 'Dot Product Similarity\\n\\nLet’s calculate the dot products between the three sentences. The following line of code will do it.\\n\\nAnd the results are:\\n\\nThe similarity between sentences 1 and 2 (0.8188) is much larger than the similarities between the other pairs. This confirms our predictions.\\n\\nJust for consistency, we also calculate the similarities between each sentence and itself, to confirm that a sentence and itself has the highest similarity score.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "User: What do you know about 5G networks\n",
            "Retrieving information...\n",
            "Chatbot:\n",
            "I'm sorry, but I do not have any information about 5G networks. Can I help you with anything else?\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "Ending chat.\n"
          ]
        }
      ],
      "source": [
        "# Create an instance of the Chatbot class\n",
        "chatbot = Chatbot(vectorstore)\n",
        "\n",
        "# Run the chatbot\n",
        "chatbot.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9e69aa7c",
      "metadata": {
        "id": "9e69aa7c"
      },
      "source": [
        "In the conversation above, notice a few observations that reflect the different components of what we built:\n",
        "\n",
        "- **Direct response**: For user messages that don’t require retrieval, such as \"Hello, I have a question\", the chatbot responds directly without requiring retrieval.\n",
        "- **Citation generation**: For responses that do require retrieval (\"What’s the difference between word and sentence embeddings\"), the endpoint returns the response together with the citations.\n",
        "- **State management**: The endpoint maintains the state of the conversation via the `conversation_id` parameter, for example, by being able to correctly respond to a vague user message of \"And what are their similarities\".\n",
        "- **Response synthesis**: The model can decide if none of the retrieved documents provide the necessary information required to answer a user message. For example, when asked the question \"What do you know about 5G networks\", the chatbot goes on and retrieves external information from the index. However, it doesn’t use any of the information in its response as none of them is relevant to the question."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}