examples/63_How_RAG_with_txtai_works.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VGeVB8M41jqW"
      },
      "source": [
        "# How RAG with txtai works\n",
        "\n",
        "[txtai](https://github.com/neuml/txtai) is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\n",
        "\n",
        "Large Language Models (LLMs) have captured the public's attention with their impressive capabilities. The Generative AI era has reached a fever pitch with some predicting the coming rise of superintelligence.\n",
        "\n",
        "LLMs are far from perfect though and we're still a ways away from true AI. The biggest challenge is with hallucinations. Hallucinations is the term for when a LLM generates output that is factually incorrect. The alarming part of this is that on a cursory glance, it actually sounds like factual content. The default behavior of LLMs is to produce plausible answers even when no plausible answer exists. LLMs are not great at saying I don't know.\n",
        "\n",
        "Retrieval Augmented Generation (RAG) helps reduce the risk of hallucinations by limiting the context in which a LLM can generate answers. This is typically done with a search query that hydrates a prompt with a relevant context. RAG has been one of the most practical use cases of the Generative AI era.\n",
        "\n",
        "txtai has a multiple ways to run RAG pipelines as follows.\n",
        "\n",
        "- Embeddings instance and LLM. Run the embeddings search and plug the search results into a LLM prompt.\n",
        "- RAG (aka Extractor) pipeline which automatically adds a search context to LLM prompts.\n",
        "- RAG FastAPI service with YAML\n",
        "\n",
        "This notebook will cover all these methods and shows how RAG with txtai works."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZQrHIw351lwE"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "R0AqRP7v1hdr"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai#egg=txtai[api,pipeline] autoawq"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmPN8RDF1pXd"
      },
      "source": [
        "# Components of a RAG pipeline\n",
        "\n",
        "Before using txtai's RAG pipeline, we'll show how each of the underlying components work together. In this example, we'll load the [txtai Wikipedia embeddings database](https://huggingface.co/NeuML/txtai-wikipedia) and a LLM. From there, we'll run a RAG process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XZ7vPBIs1rGZ"
      },
      "outputs": [],
      "source": [
        "from txtai import Embeddings, LLM\n",
        "\n",
        "# Load Wikipedia Embeddings database\n",
        "embeddings = Embeddings()\n",
        "embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
        "\n",
        "# Create LLM\n",
        "llm = LLM(\"TheBloke/Mistral-7B-OpenOrca-AWQ\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we'll create a prompt template to use for the RAG pipeline. The prompt has a placeholder for the question and context."
      ],
      "metadata": {
        "id": "ufUxz8MXLA8n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Prompt template\n",
        "prompt = \"\"\"<|im_start|>system\n",
        "You are a friendly assistant. You answer questions from users.<|im_end|>\n",
        "<|im_start|>user\n",
        "Answer the following question using only the context below. Only include information\n",
        "specifically discussed.\n",
        "\n",
        "question: {question}\n",
        "context: {context} <|im_end|>\n",
        "<|im_start|>assistant\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "4f8DbjzMKwaE"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "After that, we'll generate the context using an embeddings (aka vector) query. This query finds the top 3 most similar matches to the question **\"How do you make beer 🍺?\"**"
      ],
      "metadata": {
        "id": "x8tp7IoBL27y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "question = \"How do you make beer?\"\n",
        "\n",
        "# Generate context\n",
        "context = \"\\n\".join([x[\"text\"] for x in embeddings.search(question)])\n",
        "print(context)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Uyi_WgpULLFs",
        "outputId": "a4c7ae0c-287c-409a-ab55-c8c353125405"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Brewing is the production of beer by steeping a starch source (commonly cereal grains, the most popular of which is barley) in water and fermenting the resulting sweet liquid with yeast.  It may be done in a brewery by a commercial brewer, at home by a homebrewer, or communally. Brewing has taken place since around the 6th millennium BC, and archaeological evidence suggests that emerging civilizations, including ancient Egypt, China, and Mesopotamia, brewed beer. Since the nineteenth century the brewing industry has been part of most western economies.\n",
            "Beer is produced through steeping a sugar source (commonly Malted cereal grains) in water and then fermenting with yeast. Brewing has taken place since around the 6th millennium BC, and archeological evidence suggests that this technique was used in ancient Egypt. Descriptions of various beer recipes can be found in Sumerian writings, some of the oldest known writing of any sort. Brewing is done in a brewery by a brewer, and the brewing industry is part of most western economies. In 19th century Britain, technological discoveries and improvements such as Burtonisation and the Burton Union system significantly changed beer brewing.\n",
            "Craft beer is a beer that has been made by craft breweries, which typically produce smaller amounts of beer, than larger \"macro\" breweries, and are often independently owned. Such breweries are generally perceived and marketed as emphasising enthusiasm, new flavours, and varied brewing techniques.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we'll take the question and context and put that into the prompt."
      ],
      "metadata": {
        "id": "x_EPbqvFL_Uw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(llm(prompt.format(question=question, context=context)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QCh1IaJ1L1Kt",
        "outputId": "93fbeab4-e477-4834-d01c-366fbf86361f"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "To make beer, you need to steep a starch source, such as malted cereal grains (commonly barley), in water. This process creates a sweet liquid called wort. Then, yeast is added to the wort, which ferments the liquid and produces alcohol and carbon dioxide. The beer is then aged, filtered, and packaged for consumption. This process has been used since around the 6th millennium BC and has been a part of most western economies since the 19th century.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Looking at the generated answer, we can see it's based on the context above. The LLM generates a paragraph of text using the context as input. While this same answer could be directly asked of the LLM, this helps ensure the answer is based on known factual data."
      ],
      "metadata": {
        "id": "CSNbxxIMN_26"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# The RAG Pipeline\n",
        "\n",
        "txtai has a RAG pipeline that makes this even easier. The logic to generate the context and join it context with the prompt is built in. Let's try that."
      ],
      "metadata": {
        "id": "tgLA-61iOP0r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from txtai import RAG\n",
        "\n",
        "# Create RAG pipeline using existing components. LLM parameter can also be a model path.\n",
        "rag = RAG(embeddings, llm, template=prompt)"
      ],
      "metadata": {
        "id": "eWMQUTxFOf_5"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's ask a question similar to the last one. This time we'll ask **\"How do you make wine🍷?\"**"
      ],
      "metadata": {
        "id": "8ePe71D0O0j_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(rag(\"How do you make wine?\", maxlength=2048)[\"answer\"])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3ZaUe2YcOz_B",
        "outputId": "ada69201-9f94-498b-b33a-971b4d2e4fdb"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "To make wine, follow these steps:\n",
            "\n",
            "1. Select the fruit: Choose high-quality grapes or other fruit for wine production.\n",
            "\n",
            "2. Fermentation: Introduce yeast to the fruit, which will consume the sugar present in the juice and convert it into ethanol and carbon dioxide.\n",
            "\n",
            "3. Monitor temperature and oxygen levels: Control the temperature and speed of fermentation, as well as the levels of oxygen present in the must at the start of fermentation.\n",
            "\n",
            "4. Primary fermentation: This stage lasts from 5 to 14 days, during which the yeast consumes the sugar and produces alcohol and carbon dioxide.\n",
            "\n",
            "5. Secondary fermentation (optional): If desired, allow the wine to undergo a secondary fermentation, which can last another 5 to 10 days.\n",
            "\n",
            "6. Fermentation location: Choose the appropriate fermentation vessel, such as stainless steel tanks, open wooden vats, wine barrels, or wine bottles for sparkling wines.\n",
            "\n",
            "7. Bottle and age the wine: Transfer the finished wine into bottles and allow it to age, if desired, to develop flavors and complexity.\n",
            "\n",
            "Remember that wine can be made from various fruits, but grapes are most commonly used, and the term \"wine\" generally refers to grape wine when used without a qualifier.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# RAG API Endpoint\n",
        "\n",
        "Did you know that txtai has a built-in framework for automatically generating FastAPI services? This can be done with a YAML configuration file."
      ],
      "metadata": {
        "id": "L5udDj5dMrjY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile config.yml\n",
        "\n",
        "# Load Wikipedia Embeddings index\n",
        "cloud:\n",
        "  provider: huggingface-hub\n",
        "  container: neuml/txtai-wikipedia\n",
        "\n",
        "# RAG pipeline configuration\n",
        "rag:\n",
        "  path: TheBloke/Mistral-7B-OpenOrca-AWQ\n",
        "  output: flatten\n",
        "  template: |\n",
        "    <|im_start|>system\n",
        "    You are a friendly assistant. You answer questions from users.<|im_end|>\n",
        "    <|im_start|>user\n",
        "    Answer the following question using only the context below. Only include information\n",
        "    specifically discussed.\n",
        "\n",
        "    question: {question}\n",
        "    context: {context} <|im_end|>\n",
        "    <|im_start|>assistant"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NxEQqdzmNA0T",
        "outputId": "268d6027-d3d7-46f1-acc4-2016d3059f5e"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing config.yml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note how the same prompt template and models are set. This time instead of doing that with Python, it's done with a YAML configuration file 🔥\n",
        "\n",
        "Now let's start the API service using this configuration."
      ],
      "metadata": {
        "id": "rAmN_3aWNNyK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!CONFIG=config.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
        "!sleep 90"
      ],
      "metadata": {
        "id": "gzHPH5GKNXOQ"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's run a RAG query using the API service. Keeping with the theme, we'll ask **\"How do you make whisky 🥃?\"**"
      ],
      "metadata": {
        "id": "bqIqEsAjPqsa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!curl \"http://localhost:8000/rag?query=how+do+you+make+whisky&maxlength=2048\""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f14clGi1PeQW",
        "outputId": "0a056edc-e5c2-464d-fef0-c766369a2590"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "To make whisky, follow these steps:\n\n",
            "1. Choose the grains: Select the grains you want to use for your whisky, such as barley, corn, rye, or wheat.\n\n",
            "2. Malt the grains (optional): If using barley, malt the grains by soaking them in water and allowing them to germinate. This process releases enzymes that help break down starches into fermentable sugars.\n\n",
            "3. Mill the grains: Grind the grains to create a coarse flour, which will be mixed with water to create a mash.\n\n",
            "4. Create the mash: Combine the milled grains with hot water in a large vessel, and let it sit for several hours to allow fermentation to occur. The mash should have a temperature of around 65°C (149°F) to encourage the growth of yeast.\n\n",
            "5. Add yeast: Once the mash has cooled to around 30°C (86°F), add yeast to the mixture. The yeast will ferment the sugars in the mash, producing alcohol.\n\n",
            "6. Fermentation: Allow the mixture to ferment for several days, during which the yeast will consume the sugars and produce alcohol and carbon dioxide.\n\n",
            "7. Distillation: Transfer the fermented liquid, called \"wash\" to a copper still. Heat the wash in the still, and the alcohol will vaporize and rise through the still's neck. The vapors are then condensed back into a liquid form, creating a high-proof spirit.\n\n",
            "8. Maturation: Transfer the distilled spirit to wooden casks, typically made of charred white oak. The spirit will mature in the casks for a specified period, usually ranging from 3 to 25 years. During this time, the wood imparts flavors and color to the whisky.\n\n",
            "9. Bottling: Once the whisky has reached the desired maturity, it is bottled and ready for consumption."
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "And as before, we get an answer bound by the search context provided to the LLM. This time it comes from an API service vs a direct Python method."
      ],
      "metadata": {
        "id": "L8xTq1cTVTn7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# RAG API Service with Docker\n",
        "\n",
        "txtai builds Docker images with each release. There are also Docker files available to help configure API services.\n",
        "\n",
        "The Dockerfile below builds an API service using the same config.yml."
      ],
      "metadata": {
        "id": "_UszA8iLQzGr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "```dockerfile\n",
        "# Set base image\n",
        "ARG BASE_IMAGE=neuml/txtai-gpu\n",
        "FROM $BASE_IMAGE\n",
        "\n",
        "# Copy configuration\n",
        "COPY config.yml .\n",
        "\n",
        "# Install latest version of txtai from GitHub\n",
        "RUN \\\n",
        "    apt-get update && \\\n",
        "    apt-get -y --no-install-recommends install git && \\\n",
        "    rm -rf /var/lib/apt/lists && \\\n",
        "    python -m pip install git+https://github.com/neuml/txtai\n",
        "\n",
        "# Run local API instance to cache models in container\n",
        "RUN python -c \"from txtai.api import API; API('config.yml')\"\n",
        "\n",
        "# Start server and listen on all interfaces\n",
        "ENV CONFIG \"config.yml\"\n",
        "ENTRYPOINT [\"uvicorn\", \"--host\", \"0.0.0.0\", \"txtai.api:app\"]\n",
        "```"
      ],
      "metadata": {
        "id": "z5BQu2UPQ-wV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following commands build and start a Docker API service."
      ],
      "metadata": {
        "id": "rhADZQUyVmDL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "```bash\n",
        "docker build -t txtai-wikipedia --build-arg BASE_IMAGE=neuml/txtai-gpu .\n",
        "docker run -d --gpus=all -it -p 8000:8000 txtai-wikipedia\n",
        "```"
      ],
      "metadata": {
        "id": "fMyuCMcDRPzW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This creates the same API service just this time it's through Docker. RAG queries can be run the same way."
      ],
      "metadata": {
        "id": "TF1joce_R4nx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "```bash\n",
        "curl \"http://localhost:8000/rag?query=how+do+you+make+whisky&maxlength=2048\"\n",
        "```"
      ],
      "metadata": {
        "id": "oN6l0G-gSDBb"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPwgCgBc2Er2"
      },
      "source": [
        "# Wrapping up\n",
        "\n",
        "This notebook covered the various ways to run retrieval augmented generation (RAG) with txtai. We hope you find txtai is one of the easiest and most flexible ways to get up and running fast!"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}