Meta Original RAG Model Implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f5cd44ee",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Package Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "57354445",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2ddfbdc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# downloading packages for running the notebook\n",
    "import sys\n",
    "import subprocess\n",
    "\n",
    "subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-r', 'requirements.txt', '--quiet'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94d923f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import faiss\n",
    "import torch\n",
    "from pathlib import Path\n",
    "from functools import partial\n",
    "from typing import List, Optional\n",
    "from tempfile import TemporaryDirectory\n",
    "from dataclasses import dataclass, field\n",
    "from datasets import Value, Features, Sequence, load_dataset\n",
    "from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast, HfArgumentParser, RagRetriever, RagSequenceForGeneration, RagTokenizer\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "429ddabb",
   "metadata": {},
   "source": [
    "## Text Preprocessing Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "96ec8848",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_text(text: str, n=100, character=\" \") -> List[str]:\n",
    "    \"\"\"\n",
    "        Chunking the document for Indexing.\n",
    "        Split the text every n-th occurrence of character. \n",
    "    \"\"\"\n",
    "    \n",
    "    text = text.split(character)\n",
    "    \n",
    "    return [character.join(text[i : i + n]).strip() for i in range(0, len(text), n)]\n",
    "\n",
    "\n",
    "def split_documents(documents: dict) -> dict:\n",
    "    \"\"\"\n",
    "        Split documents into passages\n",
    "    \"\"\"\n",
    "    \n",
    "    titles, texts = [], []\n",
    "    for title, text in zip(documents[\"title\"], documents[\"text\"]):\n",
    "        if text is not None:\n",
    "            for passage in split_text(text):\n",
    "                titles.append(title if title is not None else \"\")\n",
    "                texts.append(passage)\n",
    "    \n",
    "    return {\"title\": titles, \"text\": texts}\n",
    "\n",
    "\n",
    "def embed(documents: dict, ctx_encoder: DPRContextEncoder, ctx_tokenizer: DPRContextEncoderTokenizerFast) -> dict:\n",
    "    \"\"\"\n",
    "        Compute the DPR embeddings of document passages\n",
    "    \"\"\"\n",
    "    \n",
    "    input_ids = ctx_tokenizer(documents[\"title\"], documents[\"text\"], truncation=True, padding=\"longest\", return_tensors=\"pt\")[\"input_ids\"]\n",
    "    embeddings = ctx_encoder(input_ids.to(device=device), return_dict=True).pooler_output\n",
    "    \n",
    "    return {\"embeddings\": embeddings.detach().cpu().numpy()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a39b2d2f",
   "metadata": {},
   "source": [
    "Note: The following code block is used to cleanup the dataset to create question answer pair for RAG setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "35f3cb61",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('dataset.txt', 'r') as f:\n",
    "    data = f.read()\n",
    "    \n",
    "data = data.replace('\\n\\n','\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b34b044c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.split('---')\n",
    "\n",
    "for i in range(0, len(data)):\n",
    "    if i==4:\n",
    "        data[i] = data[i].replace('\\n**', '\\n###').replace('**','')\n",
    "    elif i==3:\n",
    "        data[i] = data[i].replace('**','')\n",
    "    else:\n",
    "        data[i] = data[i].replace('**','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c369b8a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ques_ans = dict()\n",
    "for i in range(0, len(data)):\n",
    "    temp = data[i]\n",
    "    temp = temp.split('\\n###')\n",
    "    \n",
    "    for j in range(1, len(temp)):\n",
    "        tp = temp[j].split('\\n')\n",
    "        ques_ans[tp[0]] = \" \".join(tp[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "06f32b30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "qa_dict = dict()\n",
    "qa_dict['title'] = []\n",
    "qa_dict['text'] = []\n",
    "\n",
    "for key, value in ques_ans.items():\n",
    "    qa_dict['title'].append(key)\n",
    "    qa_dict['text'].append(value)\n",
    "    \n",
    "qa_df = pd.DataFrame.from_dict(qa_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "efce0928",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['title', 'text'],\n",
       "    num_rows: 50\n",
       "})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "# You can load a Dataset object this way\n",
    "dataset = Dataset.from_pandas(qa_df)\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "654bc4a1",
   "metadata": {},
   "source": [
    "Note: To use the DPR Context Encoder the Question/Answer Pairs need to be kept as Title/Text, as the model has been built on the same perspective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5186951b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at facebook/dpr-ctx_encoder-single-nq-base were not used when initializing DPRContextEncoder: ['ctx_encoder.bert_model.pooler.dense.bias', 'ctx_encoder.bert_model.pooler.dense.weight']\n",
      "- This IS expected if you are initializing DPRContextEncoder from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DPRContextEncoder from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'DPRQuestionEncoderTokenizer'. \n",
      "The class this function is called from is 'DPRContextEncoderTokenizerFast'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# And compute the embeddings\n",
    "ctx_encoder = DPRContextEncoder.from_pretrained(\"facebook/dpr-ctx_encoder-single-nq-base\").to(device=device)\n",
    "ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained(\"facebook/dpr-ctx_encoder-single-nq-base\")\n",
    "\n",
    "new_features = Features(\n",
    "    {\"text\": Value(\"string\"), \"title\": Value(\"string\"), \"embeddings\": Sequence(Value(\"float64\"))}\n",
    ")  # optional, save as float32 instead of float64 to save space\n",
    "\n",
    "dataset = dataset.map(\n",
    "    partial(embed, ctx_encoder=ctx_encoder, ctx_tokenizer=ctx_tokenizer),\n",
    "    batched=True,\n",
    "    batch_size=8,\n",
    "    features=new_features,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "02b1fea4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10375ed2cc6a48ef8d6d94a36204b2fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text', 'title', 'embeddings'],\n",
       "    num_rows: 50\n",
       "})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's use the Faiss implementation of HNSW for fast approximate nearest neighbor search\n",
    "\n",
    "index = faiss.IndexHNSWFlat(768, 128, faiss.METRIC_INNER_PRODUCT)\n",
    "dataset.add_faiss_index(\"embeddings\", custom_index=index)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f49bcf",
   "metadata": {},
   "source": [
    "Refer to this HF Repo - https://huggingface.co/facebook/rag-token-nq and Paper - https://arxiv.org/pdf/2005.11401.pdf for more context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "19bf9d0b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'DPRQuestionEncoderTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'DPRQuestionEncoderTokenizerFast'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'BartTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'BartTokenizerFast'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01e90635616b46de93d56910ff45230d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading config.json:   0%|          | 0.00/4.60k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7cba1417930e44d8a8a2ee2162e74573",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/2.06G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'DPRQuestionEncoderTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'DPRQuestionEncoderTokenizerFast'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'BartTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'RagTokenizer'. \n",
      "The class this function is called from is 'BartTokenizerFast'.\n"
     ]
    }
   ],
   "source": [
    "retriever = RagRetriever.from_pretrained(\n",
    "    \"facebook/rag-token-nq\", index_name=\"custom\", indexed_dataset=dataset\n",
    ")\n",
    "model = RagSequenceForGeneration.from_pretrained(\"facebook/rag-token-nq\", retriever=retriever, force_download=True)\n",
    "tokenizer = RagTokenizer.from_pretrained(\"facebook/rag-token-nq\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "edf4f596",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What are the documents required to apply for the new pan\n",
      "a citizenship renunciation letter\n"
     ]
    }
   ],
   "source": [
    "question = \"What are the documents required to apply for the new pan\"\n",
    "input_ids = tokenizer.question_encoder(question, return_tensors=\"pt\")[\"input_ids\"]\n",
    "generated = model.generate(input_ids)\n",
    "generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]\n",
    "\n",
    "print(question)\n",
    "print(generated_string.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fed1927f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WHow long does it usually take to receive the PAN card after applying?\n",
      "3 weeks\n"
     ]
    }
   ],
   "source": [
    "question = \"WHow long does it usually take to receive the PAN card after applying?\"\n",
    "input_ids = tokenizer.question_encoder(question, return_tensors=\"pt\")[\"input_ids\"]\n",
    "generated = model.generate(input_ids)\n",
    "generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]\n",
    "\n",
    "print(question)\n",
    "print(generated_string.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "133e583b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What is the cost/fees of a PAN card?\n",
      "us $ 2,500\n"
     ]
    }
   ],
   "source": [
    "question = \"What is the cost/fees of a PAN card?\"\n",
    "input_ids = tokenizer.question_encoder(question, return_tensors=\"pt\")[\"input_ids\"]\n",
    "generated = model.generate(input_ids)\n",
    "generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]\n",
    "\n",
    "print(question)\n",
    "print(generated_string.strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3a11838",
   "metadata": {},
   "source": [
    "## Not Proceeding Further as the results are not good enough !"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}