notebooks/rag_langchain_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also follow along on Google Colab!\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/MadryLab/context-cite/blob/main/notebooks/rag_langchain_example.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running `ContextCite` with a RAG using LangChain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we'll show a quick example of how to use `ContextCite` with a RAG chain using the `langchain` library. **If running in Colab, be sure to change your to a GPU runtime!** Thanks to Bagatur Askaryan for helpful feedback!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU context-cite langchain-community langchain-openai langchain-core langchain-text-splitters faiss-gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-05-05 18:56:13--  https://raw.githubusercontent.com/MadryLab/context-cite/main/assets/solar_eclipse.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 24030 (23K) [text/plain]\n",
      "Saving to: ‘solar_eclipse.txt’\n",
      "\n",
      "solar_eclipse.txt   100%[===================>]  23.47K  --.-KB/s    in 0s      \n",
      "\n",
      "2024-05-05 18:56:13 (158 MB/s) - ‘solar_eclipse.txt’ saved [24030/24030]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://raw.githubusercontent.com/MadryLab/context-cite/main/assets/solar_eclipse.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to\n",
      "[nltk_data]     /mnt/xfs/home/bencw/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch as ch\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "from langchain_core.runnables import RunnablePassthrough, chain\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline\n",
    "\n",
    "from context_cite import ContextCiter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = None # Add your OpenAI key here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a langchain RAG chain that does not involve `ContextCite`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple RAG chain (without ContextCite)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll load a model and tokenizer (which we'll use later on with ContextCite too)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n",
    "PROMPT_TEMPLATE = \"Context: {context}\\n\\nQuery: {query}\"\n",
    "GENERATE_KWARGS = {\"max_new_tokens\": 512, \"do_sample\": False}\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=ch.float16).cuda()\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer, device=model.device, **GENERATE_KWARGS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll create a RAG chain using a local `txt` file (a Wikipedia article about the Transformer architecture) as our \"database\" to keep things simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Created a chunk of size 1508, which is longer than the specified 1000\n",
      "Created a chunk of size 1399, which is longer than the specified 1000\n",
      "Created a chunk of size 1273, which is longer than the specified 1000\n",
      "Created a chunk of size 1490, which is longer than the specified 1000\n",
      "Created a chunk of size 1275, which is longer than the specified 1000\n",
      "Created a chunk of size 1385, which is longer than the specified 1000\n",
      "Created a chunk of size 1012, which is longer than the specified 1000\n"
     ]
    }
   ],
   "source": [
    "loader = TextLoader(\"solar_eclipse.txt\")\n",
    "documents = loader.load()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "embeddings = OpenAIEmbeddings()\n",
    "db = FAISS.from_documents(texts, embeddings)\n",
    "retriever = db.as_retriever()\n",
    "messages = [{\"role\": \"user\", \"content\": PROMPT_TEMPLATE}]\n",
    "chat_prompt_template = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "prompt = PromptTemplate.from_template(chat_prompt_template)\n",
    "llm = HuggingFacePipeline(pipeline=pipe)\n",
    "\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join([d.page_content for d in docs])\n",
    "\n",
    "chain = (\n",
    "    {\"context\": retriever | format_docs, \"query\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The longest duration of totality for this solar eclipse was 4 minutes and 28 seconds near the Mexican town of Nazas, Durango.\n"
     ]
    }
   ],
   "source": [
    "query = \"Where was the longest duration of totality for this solar eclipse?\"\n",
    "output = chain.invoke(query)\n",
    "response = output.split(f\"<|assistant|>\\n\")[-1]\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding in `ContextCite`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll add ContextCite by wrapping the `prompt` and `llm` in a `ContextCiter` Runnable class. This class will take care of formatting the context and query, as well as run generation with our LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables.base import Runnable, Input, Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContextCiteRunnable(Runnable):\n",
    "    def __init__(self, model, tokenizer):\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "    def invoke(self, context_and_query: Input, _: Input) -> Output:\n",
    "        context = context_and_query[\"context\"]\n",
    "        query = context_and_query[\"query\"]\n",
    "        cc = ContextCiter(model, tokenizer, context, query)\n",
    "        return cc.response, cc.get_attributions(as_dataframe=True, top_k=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "cc_runnable = ContextCiteRunnable(model, tokenizer)\n",
    "cc_chain = (\n",
    "    {\"context\": retriever | format_docs, \"query\": RunnablePassthrough()}\n",
    "    | cc_runnable\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attributed: The longest duration of totality for this solar eclipse was 4 minutes and 28 seconds near the Mexican town of Nazas, Durango.</s>\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48bfef9bac0b43d181423c92f6087a4a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/64 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response, attributions = cc_chain.invoke(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_7e042_row0_col0 {\n",
       "  background-color: rgb(80, 180, 80);\n",
       "}\n",
       "#T_7e042_row1_col0 {\n",
       "  background-color: rgb(245.40682464937174, 250.88863913544503, 245.40682464937174);\n",
       "}\n",
       "#T_7e042_row2_col0 {\n",
       "  background-color: rgb(249.8669654789871, 252.80012806242306, 249.8669654789871);\n",
       "}\n",
       "#T_7e042_row3_col0 {\n",
       "  background-color: rgb(251.8032105934135, 253.6299473971772, 251.8032105934135);\n",
       "}\n",
       "#T_7e042_row4_col0 {\n",
       "  background-color: rgb(252.53265986266737, 253.94256851257174, 252.53265986266737);\n",
       "}\n",
       "#T_7e042_row5_col0 {\n",
       "  background-color: rgb(252.59873423131407, 253.9708860991346, 252.59873423131407);\n",
       "}\n",
       "#T_7e042_row6_col0 {\n",
       "  background-color: rgb(253.54974004116752, 254.37846001764322, 253.54974004116752);\n",
       "}\n",
       "#T_7e042_row7_col0 {\n",
       "  background-color: rgb(255, 255, 255);\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_7e042\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_7e042_level0_col0\" class=\"col_heading level0 col0\" >Score</th>\n",
       "      <th id=\"T_7e042_level0_col1\" class=\"col_heading level0 col1\" >Source</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "      <td id=\"T_7e042_row0_col0\" class=\"data row0 col0\" >45.271</td>\n",
       "      <td id=\"T_7e042_row0_col1\" class=\"data row0 col1\" >With a magnitude of 1.0566, the eclipse's longest duration of totality was 4 minutes and 28 seconds near the Mexican town of Nazas, Durango.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "      <td id=\"T_7e042_row1_col0\" class=\"data row1 col0\" >2.482</td>\n",
       "      <td id=\"T_7e042_row1_col1\" class=\"data row1 col1\" >This gave the eclipse a wider path of totality and more maximum time in totality (4 min 28 s) compared to the total eclipse in 2017 (2 min 40 s), which had a magnitude of 1.0306.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "      <td id=\"T_7e042_row2_col0\" class=\"data row2 col0\" >1.328</td>\n",
       "      <td id=\"T_7e042_row2_col1\" class=\"data row2 col1\" >Later, the total solar eclipse was visible from North America, starting from the west coast of Mexico then ascending in a northeasterly direction through Mexico, the United States, and Canada, before ending in the Atlantic Ocean about 700 kilometers southwest of Ireland.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "      <td id=\"T_7e042_row3_col0\" class=\"data row3 col0\" >0.827</td>\n",
       "      <td id=\"T_7e042_row3_col1\" class=\"data row3 col1\" >TOP: Solar prominences as seen from Third Connecticut Lake, New Hampshire - MIDDLE: Solar activity 08 April 2024 imaged by NASA Solar Dynamics Observatory AIA 304 telescope.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "      <td id=\"T_7e042_row4_col0\" class=\"data row4 col0\" >0.638</td>\n",
       "      <td id=\"T_7e042_row4_col1\" class=\"data row4 col1\" >A solar eclipse occurs when the Moon passes between Earth and the Sun, thereby obscuring the Sun.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row5\" class=\"row_heading level0 row5\" >5</th>\n",
       "      <td id=\"T_7e042_row5_col0\" class=\"data row5 col0\" >0.621</td>\n",
       "      <td id=\"T_7e042_row5_col1\" class=\"data row5 col1\" >- BOTTOM: National Solar Observatory GONG telescope movie of solar activity in H-Alpha for the day of the April 8, 2024 eclipse, showing how prominences hardly changed during the eclipse.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row6\" class=\"row_heading level0 row6\" >6</th>\n",
       "      <td id=\"T_7e042_row6_col0\" class=\"data row6 col0\" >0.375</td>\n",
       "      <td id=\"T_7e042_row6_col1\" class=\"data row6 col1\" >Animation of the eclipse path (including the path of totality)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7e042_level0_row7\" class=\"row_heading level0 row7\" >7</th>\n",
       "      <td id=\"T_7e042_row7_col0\" class=\"data row7 col0\" >0.000</td>\n",
       "      <td id=\"T_7e042_row7_col1\" class=\"data row7 col1\" >[2][3]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fd82abc2350>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attributions"
   ]
  }
 ],
 "metadata": {
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}