2_QA_with_deberta_wikipedia_rag.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "48dbd313"
      },
      "source": [
        "# DeBERTa + Wikipedia RAG\n",
        "### The input data for this notebook can be obtained via [this](https://www.kaggle.com/competitions/kaggle-llm-science-exam) Kaggle competition\n"
      ],
      "id": "48dbd313"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7602c53d",
        "outputId": "d5431794-c9ba-4312-d6e1-f721b806e4de"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Processing /kaggle/input/faiss-gpu-173-python310/faiss_gpu-1.7.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\r\n",
            "Installing collected packages: faiss-gpu\r\n",
            "Successfully installed faiss-gpu-1.7.2\r\n",
            "Processing /kaggle/input/datasets-214/datasets-2.14.5-py3-none-any.whl\r\n",
            "Installing collected packages: datasets\r\n",
            "  Attempting uninstall: datasets\r\n",
            "    Found existing installation: datasets 2.1.0\r\n",
            "    Uninstalling datasets-2.1.0:\r\n",
            "      Successfully uninstalled datasets-2.1.0\r\n",
            "Successfully installed datasets-2.14.5\r\n"
          ]
        }
      ],
      "source": [
        "# Installing offline dependencies\n",
        "!pip install -U --no-deps /kaggle/input/faiss-gpu-173-python310/faiss_gpu-1.7.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n",
        "!pip install -U --no-deps /kaggle/input/datasets-214/datasets-2.14.5-py3-none-any.whl"
      ],
      "id": "7602c53d"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Create FAISS index\n",
        "### This index contains embeddings of wiki text"
      ],
      "metadata": {
        "id": "aEbSCtVUqj8S"
      },
      "id": "aEbSCtVUqj8S"
    },
    {
      "cell_type": "code",
      "source": [
        "import faiss\n",
        "from tqdm.auto import tqdm\n",
        "from pathlib import Path\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "import torch\n",
        "import collections\n",
        "from torch.utils.data import DataLoader\n",
        "from transformers import AutoTokenizer, AutoModel, AutoModelForMultipleChoice\n",
        "from datasets import load_from_disk\n",
        "from torch.nn.utils.rnn import pack_sequence\n",
        "\n",
        "wikipedia_path = Path(\"/kaggle/input/270k-llm-wiki\") # Wiki data URL: https://www.kaggle.com/datasets/eashish/270k-llm-wiki\n",
        "embedding_size = 384\n",
        "batch_size = 128\n",
        "max_length = 512\n",
        "checkpoint = 'BAAI/bge-base-en-v1.5'\n",
        "embedding_size = 768\n",
        "\n",
        "# Load tokenizer and model\n",
        "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
        "model = AutoModel.from_pretrained(checkpoint).cuda().half()\n",
        "\n",
        "def transform(batch):\n",
        "\n",
        "    if 'BAAI' in checkpoint:\n",
        "        batch[\"text\"] = [\"Represent this sentence for searching relevant passages: \" + x for x in batch[\"text\"]]\n",
        "    elif checkpoint == \"intfloat/e5-small-v2\":\n",
        "        batch[\"text\"] = [\"passage: \"+ x for x in batch[\"text\"]]\n",
        "\n",
        "    tokens = tokenizer(batch[\"text\"], truncation=True, padding='max_length', return_tensors=\"pt\", max_length=max_length)\n",
        "    return tokens.to(\"cuda\")\n",
        "\n",
        "# Create faiss index, it will use the same index as wikipedia_index (not the \"id\", but the row index)\n",
        "faiss_index = faiss.IndexFlatL2(embedding_size)\n",
        "\n",
        "# Create dataset and dataloader\n",
        "dataset = load_from_disk(wikipedia_path)\n",
        "dataset.set_transform(transform)\n",
        "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "# Compute embeddings\n",
        "outputs = np.zeros((len(dataset), embedding_size), dtype=np.float16)\n",
        "with torch.inference_mode():\n",
        "        for i, batch in tqdm(enumerate(dataloader), leave=False, total=len(dataloader)):\n",
        "            embeddings = model(**batch).pooler_output\n",
        "            embeddings = F.normalize(embeddings, p=2, dim=1)\n",
        "            outputs[batch_size*i:batch_size*(i+1)] = embeddings.detach().cpu().numpy()\n",
        "\n",
        "# Add embeddings to faiss index (it will use the same index as wiki_2023_index.parquet)\n",
        "faiss_index.add(outputs.astype(np.float32))\n",
        "faiss.write_index(faiss_index, str(wikipedia_path/ f\"faiss_index_{checkpoint.split('/')[-1]}.index\"))"
      ],
      "metadata": {
        "id": "j1EoFmNtqhih"
      },
      "id": "j1EoFmNtqhih",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9a45a651",
        "outputId": "785c3722-44f0-4bdb-acb8-255a3be68fb8"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5\n",
            "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
          ]
        }
      ],
      "source": [
        "import gc\n",
        "import logging\n",
        "from time import time\n",
        "from concurrent.futures import ThreadPoolExecutor\n",
        "import ctypes\n",
        "from functools import partial\n",
        "import pandas as pd\n",
        "\n",
        "# For RAG\n",
        "import faiss\n",
        "from torch.utils.data import DataLoader\n",
        "from datasets import Dataset\n",
        "\n",
        "NUM_TITLES = 5\n",
        "MAX_SEQ_LEN = 512\n",
        "MODEL_PATH = \"/kaggle/input/bge-small-faiss/\" # Faiss model saved in above cell\n",
        "\n",
        "# For LLM\n",
        "from transformers import AutoConfig, AutoModelForCausalLM\n",
        "from accelerate import init_empty_weights\n",
        "from accelerate.utils.modeling import set_module_tensor_to_device\n",
        "from safetensors.torch import load_file\n",
        "\n",
        "N_BATCHES = 5\n",
        "MAX_CONTEXT = 2750\n",
        "MAX_LENGTH = 4096"
      ],
      "id": "9a45a651"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "24c3bf19"
      },
      "outputs": [],
      "source": [
        "# Function to clean RAM & vRAM\n",
        "def clean_memory():\n",
        "    gc.collect()\n",
        "    ctypes.CDLL(\"libc.so.6\").malloc_trim(0)\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "# Load data\n",
        "df = pd.read_csv(\"/kaggle/input/kaggle-llm-science-exam/test.csv\", index_col=\"id\")\n",
        "IS_TEST_SET = len(df) != 200"
      ],
      "id": "24c3bf19"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e321c117"
      },
      "source": [
        "## 1. Wikipedia Retrieval Augmented Generation (RAG)\n",
        "\n",
        "We use the [bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) to embed the Wikipedia dataset."
      ],
      "id": "e321c117"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3e668c30"
      },
      "outputs": [],
      "source": [
        "# New SentenceTransformer class similar to the one used in @Mgöksu notebook but relying on the transformers library only\n",
        "\n",
        "class SentenceTransformer:\n",
        "    def __init__(self, checkpoint, device=\"cuda:0\"):\n",
        "        self.device = device\n",
        "        self.checkpoint = checkpoint\n",
        "        self.model = AutoModel.from_pretrained(checkpoint).to(self.device).half()\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
        "\n",
        "    def transform(self, batch):\n",
        "        tokens = self.tokenizer(batch[\"text\"], truncation=True, padding=True, return_tensors=\"pt\", max_length=MAX_SEQ_LEN)\n",
        "        return tokens.to(self.device)\n",
        "\n",
        "    def get_dataloader(self, sentences, batch_size=32):\n",
        "        sentences = [\"Represent this sentence for searching relevant passages: \" + x for x in sentences]\n",
        "        dataset = Dataset.from_dict({\"text\": sentences})\n",
        "        dataset.set_transform(self.transform)\n",
        "        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
        "        return dataloader\n",
        "\n",
        "    def encode(self, sentences, show_progress_bar=False, batch_size=32):\n",
        "        dataloader = self.get_dataloader(sentences, batch_size=batch_size)\n",
        "        pbar = tqdm(dataloader) if show_progress_bar else dataloader\n",
        "\n",
        "        embeddings = []\n",
        "        for batch in pbar:\n",
        "            with torch.no_grad():\n",
        "                e = self.model(**batch).pooler_output\n",
        "                e = F.normalize(e, p=2, dim=1)\n",
        "                embeddings.append(e.detach().cpu().numpy())\n",
        "        embeddings = np.concatenate(embeddings, axis=0)\n",
        "        return embeddings"
      ],
      "id": "3e668c30"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6352136c"
      },
      "outputs": [],
      "source": [
        "if IS_TEST_SET:\n",
        "    # Load embedding model\n",
        "    start = time()\n",
        "    print(f\"Starting prompt embedding, t={time() - start :.1f}s\")\n",
        "    model = SentenceTransformer(MODEL_PATH, device=\"cuda:0\")\n",
        "\n",
        "    # Get embeddings of prompts\n",
        "    f = lambda row : \" \".join([row[\"prompt\"], row[\"A\"], row[\"B\"], row[\"C\"], row[\"D\"], row[\"E\"]])\n",
        "    inputs = df.apply(f, axis=1).values # better results than prompt only\n",
        "    prompt_embeddings = model.encode(inputs, show_progress_bar=False)\n",
        "\n",
        "    # Search closest sentences in the wikipedia index\n",
        "    print(f\"Loading faiss index, t={time() - start :.1f}s\")\n",
        "    faiss_index = faiss.read_index(MODEL_PATH + '/faiss.index')\n",
        "    # faiss_index = faiss.index_cpu_to_all_gpus(faiss_index) # causes OOM, and not that long on CPU\n",
        "\n",
        "    print(f\"Starting text search, t={time() - start :.1f}s\")\n",
        "    search_index = faiss_index.search(np.float32(prompt_embeddings), NUM_TITLES)[1]\n",
        "\n",
        "    print(f\"Starting context extraction, t={time() - start :.1f}s\")\n",
        "    dataset = load_from_disk(\"/kaggle/input/all-paraphs-parsed-expanded\")\n",
        "    for i in range(len(df)):\n",
        "        df.loc[i, \"context\"] = \"-\" + \"\\n-\".join([dataset[int(j)][\"text\"] for j in search_index[i]])\n",
        "\n",
        "    # Free memory\n",
        "    faiss_index.reset()\n",
        "    del faiss_index, prompt_embeddings, model, dataset\n",
        "    clean_memory()\n",
        "    print(f\"Context added, t={time() - start :.1f}s\")"
      ],
      "id": "6352136c"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Multiple ways to perform answer selection from the context obtained above via RAG.\n",
        "\n",
        "1. ### We can use multiple models for inferencing directly on the context obtained above.\n",
        "2. ### Fine tune a model offline on the train data as well as external data and use it for inference.\n",
        "3. ### Fine tune via PEFT techniques."
      ],
      "metadata": {
        "id": "ONlK5cAEEsvB"
      },
      "id": "ONlK5cAEEsvB"
    },
    {
      "cell_type": "code",
      "source": [
        "test_df = df\n",
        "test_df.index = list(range(len(test_df)))\n",
        "test_df['id'] = list(range(len(test_df)))\n",
        "test_df[\"prompt\"] = test_df[\"context\"].apply(lambda x: x[:1750]) + \" #### \" +  test_df[\"prompt\"]\n",
        "test_df['answer'] = 'A'"
      ],
      "metadata": {
        "id": "u_Pl5aH8Er8a"
      },
      "id": "u_Pl5aH8Er8a",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "df_train = pd.read_csv('/kaggle/input/kaggle-llm-science-exam/train.csv')\n",
        "df_train = df_train.drop(columns=\"id\")\n",
        "\n",
        "df_train = pd.concat([\n",
        "    df_train,\n",
        "    pd.read_csv('/kaggle/input/additional-train-data-for-llm-science-exam/extra_train_set.csv'),\n",
        "    pd.read_csv('/kaggle/input/llm-mcq-dataset/100_examples.csv')\n",
        "])\n",
        "df_train.drop_duplicates().reset_index(inplace=True, drop=True)"
      ],
      "metadata": {
        "id": "87p7YNV3F4Hh"
      },
      "id": "87p7YNV3F4Hh",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "option_to_index = {option: idx for idx, option in enumerate('ABCDE')}\n",
        "index_to_option = {v: k for k,v in option_to_index.items()}\n",
        "\n",
        "def preprocess(example):\n",
        "    first_sentence = [example['prompt']] * 5\n",
        "    second_sentences = [example[option] for option in 'ABCDE']\n",
        "    tokenized_example = tokenizer(first_sentence, second_sentences, truncation=False)\n",
        "    tokenized_example['label'] = option_to_index[example['answer']]\n",
        "\n",
        "    return tokenized_example"
      ],
      "metadata": {
        "id": "91__XaKuF68h"
      },
      "id": "91__XaKuF68h",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@dataclass\n",
        "class DataCollatorForMultipleChoice:\n",
        "    tokenizer: PreTrainedTokenizerBase\n",
        "    padding: Union[bool, str, PaddingStrategy] = True\n",
        "    max_length: Optional[int] = None\n",
        "    pad_to_multiple_of: Optional[int] = None\n",
        "\n",
        "    def __call__(self, features):\n",
        "        label_name = 'label' if 'label' in features[0].keys() else 'labels'\n",
        "        labels = [feature.pop(label_name) for feature in features]\n",
        "        batch_size = len(features)\n",
        "        num_choices = len(features[0]['input_ids'])\n",
        "        flattened_features = [\n",
        "            [{k: v[i] for k, v in feature.items()} for i in range(num_choices)] for feature in features\n",
        "        ]\n",
        "        flattened_features = sum(flattened_features, [])\n",
        "\n",
        "        batch = self.tokenizer.pad(\n",
        "            flattened_features,\n",
        "            padding=self.padding,\n",
        "            max_length=self.max_length,\n",
        "            pad_to_multiple_of=self.pad_to_multiple_of,\n",
        "            return_tensors='pt',\n",
        "        )\n",
        "        batch = {k: v.view(batch_size, num_choices, -1) for k, v in batch.items()}\n",
        "        batch['labels'] = torch.tensor(labels, dtype=torch.int64)\n",
        "        return batch"
      ],
      "metadata": {
        "id": "767ATSX5F6_y"
      },
      "id": "767ATSX5F6_y",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Let's fine tune DeBERTa model on the train data"
      ],
      "metadata": {
        "id": "PMYfouy-GLzZ"
      },
      "id": "PMYfouy-GLzZ"
    },
    {
      "cell_type": "code",
      "source": [
        "model_dir = '/kaggle/input/deberta-v3-large-hf-weights'\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
        "\n",
        "dataset = Dataset.from_pandas(df_train)\n",
        "tokenized_dataset = dataset.map(preprocess, remove_columns=['prompt', 'A', 'B', 'C', 'D', 'E', 'answer'])\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "    warmup_ratio=0.8,\n",
        "    learning_rate=5e-6,\n",
        "    per_device_train_batch_size=2,\n",
        "    per_device_eval_batch_size=4,\n",
        "    num_train_epochs=3,\n",
        "    report_to='none',\n",
        "    output_dir='.',\n",
        ")\n",
        "\n",
        "model = AutoModelForMultipleChoice.from_pretrained(model_dir).cuda()\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    tokenizer=tokenizer,\n",
        "    data_collator=DataCollatorForMultipleChoice(tokenizer=tokenizer),\n",
        "    train_dataset=tokenized_dataset,\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "EqiOQGU3GJlr"
      },
      "id": "EqiOQGU3GJlr",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenized_test_dataset = Dataset.from_pandas(test_df.drop(columns=['id'])).map(preprocess, remove_columns=['prompt', 'A', 'B', 'C', 'D', 'E'])\n",
        "tokenized_test_dataset = tokenized_test_dataset.remove_columns([\"__index_level_0__\"])"
      ],
      "metadata": {
        "id": "Z8Q2Wp0iGoKZ"
      },
      "id": "Z8Q2Wp0iGoKZ",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get predictions on the test set"
      ],
      "metadata": {
        "id": "Amj8IcSXHZKY"
      },
      "id": "Amj8IcSXHZKY"
    },
    {
      "cell_type": "code",
      "source": [
        "test_predictions = trainer.predict(tokenized_test_dataset).predictions\n",
        "predictions_as_ids = np.argsort(-test_predictions, 1)\n",
        "\n",
        "predictions_as_answer_letters = np.array(list('ABCDE'))[predictions_as_ids]\n",
        "predictions_as_string = test_df['prediction'] = [\n",
        "    ' '.join(row) for row in predictions_as_answer_letters[:, :3]\n",
        "]\n",
        "\n",
        "submission = test_df[['id', 'prediction']]\n",
        "submission.to_csv('submission.csv', index=False)"
      ],
      "metadata": {
        "id": "gQ3eWSg4Gvkl"
      },
      "id": "gQ3eWSg4Gvkl",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "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.4"
    },
    "papermill": {
      "default_parameters": {},
      "duration": 78.671009,
      "end_time": "2023-10-01T06:34:15.091793",
      "environment_variables": {},
      "exception": null,
      "input_path": "__notebook__.ipynb",
      "output_path": "__notebook__.ipynb",
      "parameters": {},
      "start_time": "2023-10-01T06:32:56.420784",
      "version": "2.4.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}