{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/10.llms/RAG.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "weEIt-Yps4kN"
   },
   "source": [
    "# Retrieval Augmented Generation\n",
    "\n",
    "In this notebook, we will implement a simple RAG system.\n",
    "\n",
    "Concretely, we will begin by building a document embedding collection. Then for each query, we:\n",
    "1. Embed the query in that same space\n",
    "2. Use FAISS to retrieve the $n$ closest documents.\n",
    "3. Given those retrieved documents, we'll then incorporate them into the context of a prompt for an LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QKSoc5S4jJvm"
   },
   "outputs": [],
   "source": [
    "!pip install sentence-transformers\n",
    "\n",
    "# install faiss for gpu\n",
    "!pip install faiss-gpu-cu12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JJoXYInsigB3"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import operator\n",
    "\n",
    "import faiss\n",
    "import nltk\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "nltk.download(\"punkt\")\n",
    "nltk.download(\"punkt_tab\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ozxD6UWeWwXC"
   },
   "outputs": [],
   "source": [
    "# Run this early to let the model load!\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen3-4B\", device_map=\"cuda\", dtype=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-4B\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6EBE1m49tJAE"
   },
   "source": [
    "We'll use the ACL paper abstracts you worked with before as our set of documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/acl.all.tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"acl.all.tsv\", sep=\"\\t\", names=[\"cite\", \"year\", \"title\", \"abstract\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pGnRaYletfnx"
   },
   "source": [
    "## Building an index\n",
    "\n",
    "We must decide on a document embedding model; even within the SentenceTransformer family, there many pre-trained models that vary by accuracy, size, etc. See [here](https://www.sbert.net/docs/sentence_transformer/pretrained_models.html) for a list of all models.\n",
    "\n",
    "In particular, some models are trained for question answering tasks instead of strict semantic similarity; this means that questions will be placed in a similar region to their relevant answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we'll normalize all vectors so that cosine similarity reduces to a dot product\n",
    "# (enabling the use of inner product as a simliarity metric)\n",
    "def normalize(matrix):\n",
    "    row_norms = np.linalg.norm(matrix, axis=1, keepdims=True)\n",
    "    normalized_rows = matrix / row_norms\n",
    "    return normalized_rows\n",
    "\n",
    "class Index():\n",
    "    def __init__(self, model_name, texts):\n",
    "        self.encoder = SentenceTransformer(model_name)\n",
    "\n",
    "        doc_embeddings = self.encoder.encode(texts)\n",
    "        doc_embeddings = normalize(doc_embeddings)\n",
    "        num_docs, embedding_size = doc_embeddings.shape\n",
    "\n",
    "        # Our dataset is small enough that we can use exact search, so we'll use IndexFlatIP\n",
    "        # (which builds an exact index with doc product as the similarity metric)\n",
    "        self.index = faiss.IndexFlatIP(embedding_size)\n",
    "        self.index.add(doc_embeddings)\n",
    "\n",
    "        # If you want to use faster but approximate search over a larger dataset, use this\n",
    "        # self.index = faiss.IndexFlatIP(embedding_size)\n",
    "        # self.index = faiss.IndexIVFFlat(index, embedding_size, 10, faiss.METRIC_INNER_PRODUCT)\n",
    "        # self.index.train(doc_embeddings)\n",
    "        # self.index.add(doc_embeddings)\n",
    "\n",
    "    def query(self, query, n=3):\n",
    "        query_embedding = self.encoder.encode([query])\n",
    "        query_embedding = normalize(query_embedding)\n",
    "        distances, indices = self.index.search(query_embedding, n)\n",
    "        return distances[0], indices[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXfBj8IOt3H2"
   },
   "source": [
    "RAG depends on having a good retriever since we condition *only* on the documents (and passages of documents) that are retrieved as being relevant. Here, we experiment with two different embedding models:\n",
    "\n",
    "1. `all-mpnet-base-v2` is the default sentence-similarity model in sentence-transformers\n",
    "2. `multi-qa-mpnet-base-dot-v1` is trained on question/answer pairs\n",
    "\n",
    "**Consider**: why might we want to train on question/answer pairs?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5wpYMcaStzrI"
   },
   "outputs": [],
   "source": [
    "mpnet_index = Index(\"sentence-transformers/all-mpnet-base-v2\", df.abstract)\n",
    "multiqa_index = Index(\"sentence-transformers/multi-qa-mpnet-base-dot-v1\", df.abstract)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bGWxcXS0uZdQ"
   },
   "source": [
    "Now let's embed our query in the same representation space and find the documents most similar to it. Which embedding index do you think provides better results?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What was the CONLL 2018 shared task?\"\n",
    "\n",
    "distances, indices = mpnet_index.query(query)\n",
    "for dist, idx in zip(distances, indices):\n",
    "  print(\"%.3f\\t%s (%d)\\t%s\" % (dist, df.title[idx], df.year[idx], df.abstract[idx][:150]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distances, indices = multiqa_index.query(query)\n",
    "for dist, idx in zip(distances, indices):\n",
    "  print(\"%.3f\\t%s (%d)\\t%s\" % (dist, df.title[idx], df.year[idx], df.abstract[idx][:150]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qxtknuWXaTiz"
   },
   "source": [
    "## Generate\n",
    "\n",
    "Now that we've built a retriever, let's now incorporate those retrieved passages into the context of a prompt to answer our initial query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from textwrap import dedent\n",
    "def format_passage(data, idx):\n",
    "    \"\"\"\n",
    "    Generates formatted paper information for a given paper index.\n",
    "    \"\"\"\n",
    "    title = data.title.iloc[idx]\n",
    "    abstract = data.abstract.iloc[idx]\n",
    "    year = data.year.iloc[idx]\n",
    "    cite = data.cite.iloc[idx]\n",
    "\n",
    "    return f\"\"\"\n",
    "    Title: {title}\n",
    "    Year: {year}\n",
    "    Cite-key: {cite}\n",
    "    Abstract: {abstract}\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2mIv2Ma_Wx22"
   },
   "outputs": [],
   "source": [
    "def prompt_model(messages, thinking=False):\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True,\n",
    "        enable_thinking=thinking # Switches between thinking and non-thinking modes. Default is True.\n",
    "    )\n",
    "    \n",
    "    model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)\n",
    "    \n",
    "    # conduct text completion\n",
    "    generated = model.generate(\n",
    "        **model_inputs,\n",
    "        max_new_tokens=500\n",
    "    )\n",
    "\n",
    "    # let's break this down:\n",
    "    #                      | we take the element of the batch (our batch size is 1)\n",
    "    #                      |  |-----------------------------| skip our original input\n",
    "    output_ids = generated[0][len(model_inputs.input_ids[0]):].tolist()\n",
    "\n",
    "    # decode into token space\n",
    "    return tokenizer.decode(output_ids, skip_special_tokens=True).strip(\"\\n\")\n",
    "\n",
    "\n",
    "def generate_without_rag(question, thinking=False):\n",
    "    newline = \"\\n\"\n",
    "    system_prompt = dedent(f\"\"\"\n",
    "        You're a helpful assistant for question answering.\n",
    "    \"\"\").strip()\n",
    "\n",
    "    rag_prompt = dedent(f\"\"\"\n",
    "        Question: {question}\n",
    "    \"\"\").strip()\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": rag_prompt},\n",
    "    ]\n",
    "    return prompt_model(messages)\n",
    "\n",
    "def generate_with_rag(question, index, passages = None, thinking=False, show_rag=False):\n",
    "    newline = \"\\n\"\n",
    "    system_prompt = dedent(f\"\"\"\n",
    "        You're a helpful assistant for question answering. Use the information from the included passages to construct your response.\n",
    "        In your response, only reference the passages with a parenthetical citation of the cite-key. Do not refer to the passages any other way.\n",
    "    \"\"\").strip()\n",
    "\n",
    "    if passages is None:\n",
    "        passages = []\n",
    "        distances, indices = index.query(question)\n",
    "        for dist, idx in zip(distances, indices):\n",
    "            passages.append(format_passage(df, idx))\n",
    "\n",
    "    rag_prompt = dedent(f\"\"\"\n",
    "        Question: {question}\n",
    "    \n",
    "        Relevant passages to the question:\n",
    "\n",
    "        {passages[0]}\n",
    "        \n",
    "        {passages[1]}\n",
    "        \n",
    "        {passages[2]}\n",
    "    \"\"\").strip()\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": rag_prompt},\n",
    "    ]\n",
    "    output = prompt_model(messages)\n",
    "    if show_rag:\n",
    "        output = f\"{output}\\n\\nRAG prompt:{rag_prompt}\"\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's generate responses. First, we try generating a response without any context (relying only on the model's pretraining). Then, we try querying with the `multiqa` and `all` embedding indices.\n",
    "\n",
    "Here is a [list of CONLL shared tasks](https://www.conll.org/previous-tasks) over the years. Are the outputs accurate? Which do you prefer?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What was the CONLL 2018 shared task?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generate_without_rag(query))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5cxzDavXXwlp"
   },
   "outputs": [],
   "source": [
    "print(generate_with_rag(query, multiqa_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generate_with_rag(query, mpnet_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}