neuml--txtai/examples/42_Prompt_driven_search_with_LLMs.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vwELCooy4ljr"
      },
      "source": [
        "# Prompt-driven search with LLMs\n",
        "\n",
        "This notebook revisits the RAG pipeline, which has been covered in a number of previous notebooks. This pipeline is a combination of a similarity instance (embeddings or similarity pipeline) to build a question context and a model that answers questions.\n",
        "\n",
        "The RAG pipeline recently underwent a number of major upgrades to support the following.\n",
        "\n",
        "- Ability to run embeddings searches. Given that content is supported, text can be retrieved from the embeddings instance.\n",
        "- In addition to extractive qa, support text generation models, sequence to sequence models and custom pipelines\n",
        "\n",
        "These changes enable embeddings-guided and prompt-driven search with Large Language Models (LLMs) 🔥🔥🔥"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ew7orE2O441o"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LPQTb25tASIG"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_YnqorRKAbLu"
      },
      "source": [
        "# Create Embeddings and RAG instances\n",
        "\n",
        "An Embeddings instance defines methods to represent text as vectors and build vector indexes for search.\n",
        "\n",
        "The RAG pipeline is a combination of a similarity instance (embeddings or similarity pipeline) to build a question context and a model that answers questions. The model can be a prompt-driven large language model (LLM), an extractive question-answering model or a custom pipeline.\n",
        "\n",
        "Let's run a basic example.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "OUc9gqTyAYnm"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "from txtai import Embeddings, RAG\n",
        "\n",
        "# Create embeddings model with content support\n",
        "embeddings = Embeddings(path=\"sentence-transformers/all-MiniLM-L6-v2\", content=True)\n",
        "\n",
        "# Create the RAG pipeline\n",
        "rag = RAG(embeddings, \"Qwen/Qwen3-4B-Instruct-2507\", template=\"\"\"\n",
        "  Answer the following question using the provided context.\n",
        "\n",
        "  Question:\n",
        "  {question}\n",
        "\n",
        "  Context:\n",
        "  {context}\n",
        "\"\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4X5z3UjnAGe7",
        "outputId": "cacb7e9d-471a-437d-c68d-a8b51f876413"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---- Red Sox - Blue Jays ----\n",
            "{'answer': 'The Blue Jays won the game.'}\n",
            "{'answer': 'The score was 2-1 in favor of the Blue Jays.'}\n",
            "\n",
            "---- Phillies - Braves ----\n",
            "{'answer': 'The Phillies won the game.'}\n",
            "{'answer': 'The score was 5-0 in favor of the Phillies.'}\n",
            "\n",
            "---- Dodgers - Giants ----\n",
            "{'answer': 'The Giants won the game.'}\n",
            "{'answer': 'The score was Giants 5, Dodgers 4.'}\n",
            "\n",
            "---- Flyers - Lightning ----\n",
            "{'answer': 'The Flyers won the game.'}\n",
            "{'answer': 'The score was Flyers 4, Lightning 1.'}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "data = [\"Giants hit 3 HRs to down Dodgers\",\n",
        "        \"Giants 5 Dodgers 4 final\",\n",
        "        \"Dodgers drop Game 2 against the Giants, 5-4\",\n",
        "        \"Blue Jays beat Red Sox final score 2-1\",\n",
        "        \"Red Sox lost to the Blue Jays, 2-1\",\n",
        "        \"Blue Jays at Red Sox is over. Score: 2-1\",\n",
        "        \"Phillies win over the Braves, 5-0\",\n",
        "        \"Phillies 5 Braves 0 final\",\n",
        "        \"Final: Braves lose to the Phillies in the series opener, 5-0\",\n",
        "        \"Lightning goaltender pulled, lose to Flyers 4-1\",\n",
        "        \"Flyers 4 Lightning 1 final\",\n",
        "        \"Flyers win 4-1\"]\n",
        "\n",
        "questions = [\"What team won the game?\", \"What was score?\"]\n",
        "\n",
        "for query in [\"Red Sox - Blue Jays\", \"Phillies - Braves\", \"Dodgers - Giants\", \"Flyers - Lightning\"]:\n",
        "    print(\"----\", query, \"----\")\n",
        "    for answer in rag([f\"{query} {x}\" for x in questions], data):\n",
        "        print(answer)\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7AnPvSeM3N1Z"
      },
      "source": [
        "This code runs a series of questions. First it runs an embeddings filtering query to find the most relevant text. For example, `Red Sox - Blue Jays` finds text related to those teams. Then `What team won the game?` and `What was the score?` are asked.\n",
        "\n",
        "This logic is the same logic found in Notebook 5 - Extractive QA with txtai but uses prompt-based QA vs extractive QA. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aj8GoDk331cS"
      },
      "source": [
        "# Embeddings-guided and Prompt-driven Search\n",
        "\n",
        "Now for the fun stuff. Let's build an embeddings index for the `ag_news` dataset (a set of news stories from the mid 2000s). Then we'll use prompts to ask questions with embeddings results as the context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yL716oEZ43t-",
        "outputId": "23f4b0e7-a60a-4e89-fb57-06966e6612f8"
      },
      "outputs": [
      ],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"ag_news\", split=\"train\")\n",
        "\n",
        "# Create an embeddings index over the dataset\n",
        "embeddings = Embeddings(path=\"sentence-transformers/all-MiniLM-L6-v2\", content=True)\n",
        "embeddings.index(dataset[\"text\"])\n",
        "\n",
        "# Create RAG instance\n",
        "rag = RAG(embeddings, \"Qwen/Qwen3-4B-Instruct-2507\", template=\"\"\"\n",
        "  Answer the following question using the provided context.\n",
        "\n",
        "  Question:\n",
        "  {question}\n",
        "\n",
        "  Context:\n",
        "  {context}\n",
        "\"\"\", output=\"flatten\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ifl8JwLDBL7k"
      },
      "source": [
        "Now let's run a prompt-driven search!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5O1WBJ8153Mo",
        "outputId": "ddf09da2-7d4c-4fd3-b0da-df5631bacd13"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Who won the 2004 presidential election? George W. Bush won the 2004 presidential election.\n",
            "Who did the candidate beat? George W. Bush beat John F. Kerry in the 2004 presidential election.\n"
          ]
        }
      ],
      "source": [
        "question = \"Who won the 2004 presidential election?\"\n",
        "answer = rag(question)\n",
        "print(question, answer)\n",
        "\n",
        "nquestion = \"Who did the candidate beat?\"\n",
        "print(nquestion, rag(f\"{question} {answer}. {nquestion}\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AhViFXH_BZSo"
      },
      "source": [
        "And there are the answers. Let's unpack how this works.\n",
        "\n",
        "The first thing the RAG pipeline does is run an embeddings search to find the most relevant text within the index. A context string is then built using those search results.\n",
        "\n",
        "After that, a prompt is generated, run and the answer printed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JtDcVPdOB0Rv"
      },
      "source": [
        "# Additional examples\n",
        "\n",
        "Before moving on, a couple more example questions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0NNLBwC-83MM",
        "outputId": "4b9fabd7-baf9-4f7d-bb8b-c9c60c5101e4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Who won the World Series in 2004? The Boston Red Sox won the World Series in 2004.\n",
            "What team did the Red Sox beat in the World Series? The Boston Red Sox beat the St. Louis Cardinals in the World Series.\n"
          ]
        }
      ],
      "source": [
        "question = \"Who won the World Series in 2004?\"\n",
        "answer = rag(question)\n",
        "print(question, answer)\n",
        "\n",
        "nquestion = \"What team did the Red Sox beat in the World Series?\"\n",
        "print(nquestion, rag(f\"{question} {answer}. {nquestion}\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        },
        "id": "1P0zqkTW9cZW",
        "outputId": "4f2232db-761b-464f-b47d-6695c84ffb80"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'An interesting fact is that herrings communicate by farting—a quirky and unusual discovery that was honored with an Ig Nobel Prize for its oddball research.'"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "rag(\"Tell me something interesting\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ygFFcwWPGI9p"
      },
      "source": [
        "Whhaaaattt???  Is this a model hallucination?\n",
        "\n",
        "Let's run an embeddings query and see if that text is in the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qZPhLqSxGMbK",
        "outputId": "e3b2909a-1ee8-480f-afa3-201a8e27cb08"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "herrings communicate by farting\n"
          ]
        }
      ],
      "source": [
        "answer = \"herrings communicate by farting\"\n",
        "for x in embeddings.search(\"Tell me something interesting\"):\n",
        "  if answer in x[\"text\"]:\n",
        "    start = x[\"text\"].find(answer)\n",
        "    print(x[\"text\"][start:start + len(answer)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IpgxMc1DZcds"
      },
      "source": [
        "Sure enough it is 😃"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KqfvCXp2B3li"
      },
      "source": [
        "# Wrapping up\n",
        "\n",
        "This notebook covered how to run embeddings-guided and prompt-driven search with LLMs. This functionality is a major step forward towards `Generative Semantic Search` for txtai. More to come, stay tuned!"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "local",
      "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.10.19"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}