inference_deberta_v3_large_RAG.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#数据集链接：\n",
    "#https://www.kaggle.com/competitions/kaggle-llm-science-exam\n",
    "#https://www.kaggle.com/datasets/jjinho/blingfire-018/\n",
    "#https://www.kaggle.com/datasets/mbanaei/datasets-wheel/\n",
    "#https://www.kaggle.com/datasets/tomokihirose/faiss-gpu-173-python310/\n",
    "#https://www.kaggle.com/datasets/yingpengchen/gte-base-pca/\n",
    "#https://www.kaggle.com/datasets/yingpengchen/gte-base-pos/\n",
    "#https://www.kaggle.com/datasets/datafan07/llm-whls/\n",
    "#https://www.kaggle.com/datasets/inversion/sentence-transformers-222/\n",
    "#https://www.kaggle.com/datasets/gmhost/sentencetransformer-hubs/versions/2\n",
    "#https://www.kaggle.com/datasets/gmhost/wiki-2023-index-partition/\n",
    "#https://www.kaggle.com/datasets/jjinho/wikipedia-2023-07-faiss-index/\n",
    "#https://www.kaggle.com/datasets/jjinho/wikipedia-20230701/\n",
    "#https://www.kaggle.com/datasets/gmhost/wikipedia-stem-index/\n",
    "#https://www.kaggle.com/datasets/gmhost/wikipedia-stem-plaintext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T13:46:40.809779Z",
     "iopub.status.busy": "2023-10-04T13:46:40.809167Z",
     "iopub.status.idle": "2023-10-04T13:47:17.856213Z",
     "shell.execute_reply": "2023-10-04T13:47:17.854788Z",
     "shell.execute_reply.started": "2023-10-04T13:46:40.809726Z"
    }
   },
   "outputs": [],
   "source": [
    "!cp /kaggle/input/datasets-wheel/datasets-2.14.4-py3-none-any.whl /kaggle/working\n",
    "!pip install  /kaggle/working/datasets-2.14.4-py3-none-any.whl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T13:47:17.859840Z",
     "iopub.status.busy": "2023-10-04T13:47:17.858739Z",
     "iopub.status.idle": "2023-10-04T13:49:24.126123Z",
     "shell.execute_reply": "2023-10-04T13:49:24.124583Z",
     "shell.execute_reply.started": "2023-10-04T13:47:17.859798Z"
    }
   },
   "outputs": [],
   "source": [
    "# installing offline dependencies\n",
    "!pip install -U /kaggle/input/faiss-gpu-173-python310/faiss_gpu-1.7.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n",
    "!cp -rf /kaggle/input/sentence-transformers-222/sentence-transformers /kaggle/working/sentence-transformers\n",
    "!pip install -U /kaggle/working/sentence-transformers\n",
    "!pip install -U /kaggle/input/blingfire-018/blingfire-0.1.8-py3-none-any.whl\n",
    "\n",
    "!pip install --no-index --no-deps /kaggle/input/llm-whls/transformers-4.31.0-py3-none-any.whl\n",
    "!pip install --no-index --no-deps /kaggle/input/llm-whls/peft-0.4.0-py3-none-any.whl\n",
    "!pip install --no-index --no-deps /kaggle/input/llm-whls/trl-0.5.0-py3-none-any.whl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T13:49:24.129465Z",
     "iopub.status.busy": "2023-10-04T13:49:24.129065Z",
     "iopub.status.idle": "2023-10-04T13:49:24.145330Z",
     "shell.execute_reply": "2023-10-04T13:49:24.144322Z",
     "shell.execute_reply.started": "2023-10-04T13:49:24.129425Z"
    }
   },
   "outputs": [],
   "source": [
    "%%writefile backup.py\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "import os\n",
    "import gc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import re\n",
    "from tqdm.auto import tqdm\n",
    "import blingfire as bf\n",
    "\n",
    "from collections.abc import Iterable\n",
    "\n",
    "import faiss\n",
    "from faiss import write_index, read_index, read_VectorTransform\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "import torch\n",
    "import ctypes\n",
    "libc = ctypes.CDLL(\"libc.so.6\")\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Optional, Union\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from datasets import Dataset\n",
    "from transformers import AutoTokenizer\n",
    "from transformers import AutoModelForMultipleChoice, TrainingArguments, Trainer\n",
    "from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "def process_documents(documents: Iterable[str],\n",
    "                      document_ids: Iterable,\n",
    "                      split_sentences: bool = True,\n",
    "                      filter_len: int = 3,\n",
    "                      disable_progress_bar: bool = False) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Main helper function to process documents from the EMR.\n",
    "\n",
    "    :param documents: Iterable containing documents which are strings\n",
    "    :param document_ids: Iterable containing document unique identifiers\n",
    "    :param document_type: String denoting the document type to be processed\n",
    "    :param document_sections: List of sections for a given document type to process\n",
    "    :param split_sentences: Flag to determine whether to further split sections into sentences\n",
    "    :param filter_len: Minimum character length of a sentence (otherwise filter out)\n",
    "    :param disable_progress_bar: Flag to disable tqdm progress bar\n",
    "    :return: Pandas DataFrame containing the columns `document_id`, `text`, `section`, `offset`\n",
    "    \"\"\"\n",
    "\n",
    "    df = sectionize_documents(documents, document_ids, disable_progress_bar)\n",
    "\n",
    "    if split_sentences:\n",
    "        df = sentencize(df.text.values,\n",
    "                        df.document_id.values,\n",
    "                        df.offset.values,\n",
    "                        filter_len,\n",
    "                        disable_progress_bar)\n",
    "    return df\n",
    "\n",
    "\n",
    "def sectionize_documents(documents: Iterable[str],\n",
    "                         document_ids: Iterable,\n",
    "                         disable_progress_bar: bool = False) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Obtains the sections of the imaging reports and returns only the\n",
    "    selected sections (defaults to FINDINGS, IMPRESSION, and ADDENDUM).\n",
    "\n",
    "    :param documents: Iterable containing documents which are strings\n",
    "    :param document_ids: Iterable containing document unique identifiers\n",
    "    :param disable_progress_bar: Flag to disable tqdm progress bar\n",
    "    :return: Pandas DataFrame containing the columns `document_id`, `text`, `offset`\n",
    "    \"\"\"\n",
    "    processed_documents = []\n",
    "    for document_id, document in tqdm(zip(document_ids, documents), total=len(documents), disable=disable_progress_bar):\n",
    "        row = {}\n",
    "        text, start, end = (document, 0, len(document))\n",
    "        row['document_id'] = document_id\n",
    "        row['text'] = text\n",
    "        row['offset'] = (start, end)\n",
    "\n",
    "        processed_documents.append(row)\n",
    "\n",
    "    _df = pd.DataFrame(processed_documents)\n",
    "    if _df.shape[0] > 0:\n",
    "        return _df.sort_values(['document_id', 'offset']).reset_index(drop=True)\n",
    "    else:\n",
    "        return _df\n",
    "\n",
    "\n",
    "def sentencize(documents: Iterable[str],\n",
    "               document_ids: Iterable,\n",
    "               offsets: Iterable[tuple[int, int]],\n",
    "               filter_len: int = 3,\n",
    "               disable_progress_bar: bool = False) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Split a document into sentences. Can be used with `sectionize_documents`\n",
    "    to further split documents into more manageable pieces. Takes in offsets\n",
    "    to ensure that after splitting, the sentences can be matched to the\n",
    "    location in the original documents.\n",
    "\n",
    "    :param documents: Iterable containing documents which are strings\n",
    "    :param document_ids: Iterable containing document unique identifiers\n",
    "    :param offsets: Iterable tuple of the start and end indices\n",
    "    :param filter_len: Minimum character length of a sentence (otherwise filter out)\n",
    "    :return: Pandas DataFrame containing the columns `document_id`, `text`, `section`, `offset`\n",
    "    \"\"\"\n",
    "\n",
    "    document_sentences = []\n",
    "    for document, document_id, offset in tqdm(zip(documents, document_ids, offsets), total=len(documents),\n",
    "                                              disable=disable_progress_bar):\n",
    "        try:\n",
    "            _, sentence_offsets = bf.text_to_sentences_and_offsets(document)\n",
    "            for o in sentence_offsets:\n",
    "                if o[1] - o[0] > filter_len:\n",
    "                    sentence = document[o[0]:o[1]]\n",
    "                    abs_offsets = (o[0] + offset[0], o[1] + offset[0])\n",
    "                    row = {}\n",
    "                    row['document_id'] = document_id\n",
    "                    row['text'] = sentence\n",
    "                    row['offset'] = abs_offsets\n",
    "                    document_sentences.append(row)\n",
    "        except:\n",
    "            continue\n",
    "    return pd.DataFrame(document_sentences)\n",
    "\n",
    "\n",
    "def get_contexts():\n",
    "    SIM_MODEL = '/kaggle/input/sentencetransformer-hubs/gte-base'\n",
    "    DEVICE = 0\n",
    "    MAX_LENGTH = 384\n",
    "    BATCH_SIZE = 16\n",
    "\n",
    "    WIKI_PATH = \"/kaggle/input/wikipedia-20230701\"\n",
    "    wiki_files = os.listdir(WIKI_PATH)\n",
    "\n",
    "    trn = pd.read_csv(\"/kaggle/input/kaggle-llm-science-exam/test.csv\").drop(\"id\", 1)\n",
    "\n",
    "    model = SentenceTransformer(SIM_MODEL, device='cuda')\n",
    "    model.max_seq_length = MAX_LENGTH\n",
    "    model = model.half()\n",
    "\n",
    "    sentence_index = read_index(\"/kaggle/input/gte-base-pos/wikipedia_gte-base_seq512_title_pos1024_pca.index\")\n",
    "\n",
    "    # prompt_embeddings = model.encode(trn.prompt.values, batch_size=BATCH_SIZE, device=DEVICE, show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True)\n",
    "    prompt_embeddings = model.encode(\n",
    "        trn.apply(lambda row: f\"{row['prompt']}\\n{row['A']}\\n{row['B']}\\n{row['C']}\\n{row['D']}\\n{row['E']}\",\n",
    "                  axis=1).values,\n",
    "        batch_size=BATCH_SIZE, device=DEVICE, show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True)\n",
    "\n",
    "    prompt_embeddings = prompt_embeddings.detach().cpu().numpy()\n",
    "    pca_mat = read_VectorTransform('/kaggle/input/gte-base-pca/gte-base_pca.mat')\n",
    "    prompt_embeddings = pca_mat.apply_py(prompt_embeddings)\n",
    "    _ = gc.collect()\n",
    "\n",
    "    # Get the top 20 pages that are likely to contain the topic of interest\n",
    "    search_score, search_index = sentence_index.search(prompt_embeddings, 20)\n",
    "\n",
    "    # Save memory - delete sentence_index since it is no longer necessary\n",
    "    del sentence_index\n",
    "    del prompt_embeddings\n",
    "    _ = gc.collect()\n",
    "    libc.malloc_trim(0)\n",
    "\n",
    "    df = pd.read_parquet(\"/kaggle/input/wikipedia-20230701/wiki_2023_index.parquet\",\n",
    "                         columns=['id', 'file'])\n",
    "\n",
    "    # Get the article and associated file location using the index\n",
    "    wikipedia_file_data = []\n",
    "\n",
    "    for i, (scr, idx) in tqdm(enumerate(zip(search_score, search_index)), total=len(search_score)):\n",
    "        scr_idx = idx\n",
    "        _df = df.loc[scr_idx].copy()\n",
    "        _df['prompt_id'] = i\n",
    "        wikipedia_file_data.append(_df)\n",
    "    wikipedia_file_data = pd.concat(wikipedia_file_data).reset_index(drop=True)\n",
    "    wikipedia_file_data = wikipedia_file_data[['id', 'prompt_id', 'file']].drop_duplicates().sort_values(\n",
    "        ['file', 'id']).reset_index(drop=True)\n",
    "\n",
    "    # Save memory - delete df since it is no longer necessary\n",
    "    del df\n",
    "    _ = gc.collect()\n",
    "    libc.malloc_trim(0)\n",
    "\n",
    "    # Get the full text data\n",
    "    wiki_text_data = []\n",
    "\n",
    "    for file in tqdm(wikipedia_file_data.file.unique(), total=len(wikipedia_file_data.file.unique())):\n",
    "        _id = [str(i) for i in wikipedia_file_data[wikipedia_file_data['file'] == file]['id'].tolist()]\n",
    "        _df = pd.read_parquet(f\"{WIKI_PATH}/{file}\", columns=['id', 'text', 'title'])\n",
    "\n",
    "        _df_temp = _df[_df['id'].isin(_id)].copy()\n",
    "        del _df\n",
    "        _ = gc.collect()\n",
    "        libc.malloc_trim(0)\n",
    "        wiki_text_data.append(_df_temp)\n",
    "    wiki_text_data = pd.concat(wiki_text_data).drop_duplicates().reset_index(drop=True)\n",
    "    _ = gc.collect()\n",
    "\n",
    "    # Parse documents into sentences\n",
    "    processed_wiki_text_data = process_documents(wiki_text_data.text.values, wiki_text_data.id.values)\n",
    "\n",
    "    # Get embeddings of the wiki text data\n",
    "    wiki_data_embeddings = model.encode(processed_wiki_text_data.text,\n",
    "                                        batch_size=BATCH_SIZE,\n",
    "                                        device=DEVICE,\n",
    "                                        show_progress_bar=True,\n",
    "                                        convert_to_tensor=True,\n",
    "                                        normalize_embeddings=True)  # .half()\n",
    "    wiki_data_embeddings = wiki_data_embeddings.detach().cpu().numpy()\n",
    "\n",
    "    _ = gc.collect()\n",
    "\n",
    "    # Combine all answers\n",
    "    trn['answer_all'] = trn.apply(lambda x: \" \".join([x['A'], x['B'], x['C'], x['D'], x['E']]), axis=1)\n",
    "\n",
    "    # Search using the prompt and answers to guide the search\n",
    "    trn['prompt_answer_stem'] = trn['prompt'] + \" \" + trn['answer_all']\n",
    "\n",
    "    question_embeddings = model.encode(trn.prompt_answer_stem.values, batch_size=BATCH_SIZE, device=DEVICE,\n",
    "                                       show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True)\n",
    "    question_embeddings = question_embeddings.detach().cpu().numpy()\n",
    "\n",
    "    # Parameter to determine how many relevant sentences to include\n",
    "    NUM_SENTENCES_INCLUDE = 10\n",
    "\n",
    "    # List containing just Context\n",
    "    contexts = []\n",
    "\n",
    "    for r in tqdm(trn.itertuples(), total=len(trn)):\n",
    "\n",
    "        prompt_id = r.Index\n",
    "\n",
    "        prompt_indices = processed_wiki_text_data[processed_wiki_text_data['document_id'].isin(\n",
    "            wikipedia_file_data[wikipedia_file_data['prompt_id'] == prompt_id]['id'].values)].index.values\n",
    "\n",
    "        if prompt_indices.shape[0] > 0:\n",
    "            prompt_index = faiss.index_factory(wiki_data_embeddings.shape[1], \"Flat\")\n",
    "            prompt_index.add(wiki_data_embeddings[prompt_indices])\n",
    "\n",
    "            context = \"\"\n",
    "\n",
    "            # Get the top matches\n",
    "            ss, ii = prompt_index.search(question_embeddings, NUM_SENTENCES_INCLUDE)\n",
    "            for _s, _i in zip(ss[prompt_id], ii[prompt_id]):\n",
    "                context += processed_wiki_text_data.loc[prompt_indices]['text'].iloc[_i] + \" \"\n",
    "        contexts.append(context)\n",
    "\n",
    "    trn['context'] = contexts\n",
    "\n",
    "    trn[[\"prompt\", \"context\", \"A\", \"B\", \"C\", \"D\", \"E\"]].to_csv(\"./test_context.csv\", index=False)\n",
    "\n",
    "\n",
    "@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\n",
    "\n",
    "\n",
    "def generate_openbook_output():\n",
    "    #import numpy as np\n",
    "    test_df = pd.read_csv(\"./test_context.csv\")\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'\n",
    "    model_dir = \"/kaggle/input/llm-science-run-context-2\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "    model = AutoModelForMultipleChoice.from_pretrained(model_dir).cuda()\n",
    "    model.eval()\n",
    "\n",
    "    # We'll create a dictionary to convert option names (A, B, C, D, E) into indices and back again\n",
    "    options = 'ABCDE'\n",
    "    indices = list(range(5))\n",
    "\n",
    "    option_to_index = {option: index for option, index in zip(options, indices)}\n",
    "    index_to_option = {index: option for option, index in zip(options, indices)}\n",
    "\n",
    "    def preprocess(example):\n",
    "        first_sentence = [ \"[CLS] \" + example['context'] ] * 5\n",
    "        second_sentences = [\" #### \" + example['prompt'] + \" [SEP] \" + example[option] + \" [SEP]\" for option in 'ABCDE']\n",
    "        tokenized_example = tokenizer(first_sentence, second_sentences, truncation='only_first', \n",
    "                                      max_length=1536, add_special_tokens=False)\n",
    "        tokenized_example['label'] = option_to_index[example['answer']]\n",
    "\n",
    "        return tokenized_example\n",
    "\n",
    "    tokenized_test_dataset = Dataset.from_pandas(test_df[['prompt', 'context', 'A', 'B', 'C', 'D', 'E', 'answer']]).map(preprocess, remove_columns=['prompt', 'context','A', 'B', 'C', 'D', 'E', 'answer'])\n",
    "    tokenized_test_dataset = tokenized_test_dataset.remove_columns([\"__index_level_0__\"])\n",
    "    data_collator = DataCollatorForMultipleChoice(tokenizer=tokenizer)\n",
    "    test_dataloader = DataLoader(tokenized_test_dataset, batch_size=1, shuffle=False, collate_fn=data_collator,num_workers=2, pin_memory=False,)\n",
    "\n",
    "    test_predictions = []\n",
    "    for batch in tqdm(test_dataloader, total=len(test_dataloader)):\n",
    "        for k in batch.keys():\n",
    "            batch[k] = batch[k].cuda()\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "        test_predictions.append(outputs.logits.cpu().detach())\n",
    "\n",
    "    test_predictions = torch.cat(test_predictions)\n",
    "\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_answer_letters[:3]\n",
    "\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_backup.csv', index=False)\n",
    "    \n",
    "get_contexts()\n",
    "generate_openbook_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T13:49:24.147516Z",
     "iopub.status.busy": "2023-10-04T13:49:24.146896Z",
     "iopub.status.idle": "2023-10-04T14:06:08.599371Z",
     "shell.execute_reply": "2023-10-04T14:06:08.597605Z",
     "shell.execute_reply.started": "2023-10-04T13:49:24.147480Z"
    }
   },
   "outputs": [],
   "source": [
    "!python backup.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:06:11.269221Z",
     "iopub.status.busy": "2023-10-04T14:06:11.268829Z",
     "iopub.status.idle": "2023-10-04T14:06:11.300704Z",
     "shell.execute_reply": "2023-10-04T14:06:11.299711Z",
     "shell.execute_reply.started": "2023-10-04T14:06:11.269191Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "backup_model_predictions = pd.read_csv(\"submission_backup.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:06:53.188648Z",
     "iopub.status.busy": "2023-10-04T14:06:53.188159Z",
     "iopub.status.idle": "2023-10-04T14:06:53.197434Z",
     "shell.execute_reply": "2023-10-04T14:06:53.196170Z",
     "shell.execute_reply.started": "2023-10-04T14:06:53.188608Z"
    }
   },
   "outputs": [],
   "source": [
    "%%writefile parsed_context_search.py\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "import os\n",
    "import gc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import re\n",
    "from tqdm.auto import tqdm\n",
    "import blingfire as bf\n",
    "\n",
    "from collections.abc import Iterable\n",
    "\n",
    "import faiss\n",
    "from faiss import write_index, read_index, read_VectorTransform\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "import torch\n",
    "import ctypes\n",
    "libc = ctypes.CDLL(\"libc.so.6\")\n",
    "\n",
    "class CFG:\n",
    "    EMB_MODEL = \"/kaggle/input/sentencetransformer-hubs/gte-base\"\n",
    "    INDEX_PATH = \"/kaggle/input/wikipedia-stem-index/parsed_gte-base.index\"\n",
    "    WIKI_PLAINTEXT_PATH = \"/kaggle/input/wikipedia-stem-plaintext/parsed.parquet\"\n",
    "    \n",
    "    MAX_LENGTH = 512\n",
    "    BATCH_SIZE = 32\n",
    "    MAX_DOC_NUM = 10\n",
    "    \n",
    "    DEBUG = False\n",
    "\n",
    "if CFG.DEBUG:\n",
    "    trn = pd.read_csv(\"/kaggle/input/llm-stem-validationset/STEM_valid.csv\")#.drop(\"context\", 1)\n",
    "else:\n",
    "    trn = pd.read_csv(\"/kaggle/input/kaggle-llm-science-exam/test.csv\")#.drop(\"id\", 1)\n",
    "\n",
    "## Combine all answers\n",
    "trn = trn.fillna('None')\n",
    "trn['answer_all'] = trn.apply(lambda x: \" \".join([str(x['A']), str(x['B']), str(x['C']), str(x['D']), str(x['E'])]), axis=1)\n",
    "## Search using the prompt and answers to guide the search\n",
    "trn['prompt_answer_stem'] = trn['prompt'] + \" \" +trn['prompt'] + \" \" +trn['prompt'] + \" \" + trn['answer_all']\n",
    "\n",
    "model = SentenceTransformer(CFG.EMB_MODEL, device='cuda')\n",
    "model.max_seq_length = CFG.MAX_LENGTH\n",
    "\n",
    "prompt_embeddings = model.encode(trn.prompt_answer_stem.values, batch_size=CFG.BATCH_SIZE, device=0, show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True)\n",
    "prompt_embeddings = prompt_embeddings.detach().cpu().numpy()\n",
    "\n",
    "res = faiss.StandardGpuResources()\n",
    "sentence_index = read_index(CFG.INDEX_PATH)\n",
    "sentence_index_gpu = faiss.index_cpu_to_gpu(res, 0, sentence_index)\n",
    "\n",
    "QUERY_SIZE = 8\n",
    "search_score = []\n",
    "search_index = []\n",
    "total = prompt_embeddings.shape[0]\n",
    "for i in tqdm(range(0, total, QUERY_SIZE)):\n",
    "    ss, si = sentence_index_gpu.search(prompt_embeddings[i:i+QUERY_SIZE], CFG.MAX_DOC_NUM)\n",
    "    search_score.append(ss)\n",
    "    search_index.append(si)\n",
    "\n",
    "search_score = np.concatenate(search_score)\n",
    "search_index = np.concatenate(search_index)\n",
    "\n",
    "## Save memory - delete sentence_index since it is no longer necessary\n",
    "del sentence_index\n",
    "del sentence_index_gpu\n",
    "del prompt_embeddings\n",
    "#del model\n",
    "_ = gc.collect()\n",
    "#torch.cuda.empty_cache()\n",
    "libc.malloc_trim(0)\n",
    "\n",
    "df = pd.read_parquet(CFG.WIKI_PLAINTEXT_PATH,\n",
    "                     columns=['text'])\n",
    "\n",
    "contexts = []\n",
    "\n",
    "for i, (scr, idx) in tqdm(enumerate(zip(search_score, search_index)), total=len(search_score)):\n",
    "    context = \"\"\n",
    "    scr_idx = idx\n",
    "    context_list = df.loc[scr_idx].text.tolist()\n",
    "    context += \" \".join(context_list)\n",
    "    contexts.append(context)\n",
    "\n",
    "## Save memory - delete df since it is no longer necessary\n",
    "del df\n",
    "_ = gc.collect()\n",
    "libc.malloc_trim(0)\n",
    "\n",
    "trn['parsed_context'] = contexts\n",
    "\n",
    "save_cols = [\"prompt\", \"parsed_context\", \"A\", \"B\", \"C\", \"D\", \"E\"]\n",
    "trn[save_cols].to_csv(\"./test_context_part1.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:06:53.656871Z",
     "iopub.status.busy": "2023-10-04T14:06:53.656520Z",
     "iopub.status.idle": "2023-10-04T14:08:40.921345Z",
     "shell.execute_reply": "2023-10-04T14:08:40.919954Z",
     "shell.execute_reply.started": "2023-10-04T14:06:53.656843Z"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "!python parsed_context_search.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:08:44.374396Z",
     "iopub.status.busy": "2023-10-04T14:08:44.374017Z",
     "iopub.status.idle": "2023-10-04T14:08:44.384353Z",
     "shell.execute_reply": "2023-10-04T14:08:44.383304Z",
     "shell.execute_reply.started": "2023-10-04T14:08:44.374365Z"
    }
   },
   "outputs": [],
   "source": [
    "%%writefile cohere_context_search.py\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "import os\n",
    "import gc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import re\n",
    "from tqdm.auto import tqdm\n",
    "import blingfire as bf\n",
    "\n",
    "from collections.abc import Iterable\n",
    "\n",
    "import faiss\n",
    "from faiss import write_index, read_index, read_VectorTransform\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "import torch\n",
    "import ctypes\n",
    "libc = ctypes.CDLL(\"libc.so.6\")\n",
    "\n",
    "class CFG:\n",
    "    EMB_MODEL = \"/kaggle/input/sentencetransformer-hubs/gte-large\"\n",
    "    INDEX_PATH = \"/kaggle/input/wikipedia-stem-pca-index/parsed_gte-large.index\"\n",
    "    WIKI_PLAINTEXT_PATH = \"/kaggle/input/wikipedia-stem-plaintext/parsed.parquet\"\n",
    "    \n",
    "    MAX_LENGTH = 512\n",
    "    BATCH_SIZE = 16\n",
    "    MAX_DOC_NUM = 10\n",
    "    \n",
    "    DEBUG = False\n",
    "\n",
    "trn = pd.read_csv(\"./test_context_part1.csv\")#.drop(\"id\", 1)\n",
    "\n",
    "## Combine all answers\n",
    "trn = trn.fillna('None')\n",
    "trn['answer_all'] = trn.apply(lambda x: \" \".join([str(x['A']), str(x['B']), str(x['C']), str(x['D']), str(x['E'])]), axis=1)\n",
    "## Search using the prompt and answers to guide the search\n",
    "trn['prompt_answer_stem'] = trn['prompt'] + \" \" +trn['prompt'] + \" \" +trn['prompt'] + \" \" + trn['answer_all']\n",
    "\n",
    "model = SentenceTransformer(CFG.EMB_MODEL, device='cuda')\n",
    "model.max_seq_length = CFG.MAX_LENGTH\n",
    "\n",
    "prompt_embeddings = model.encode(trn.prompt_answer_stem.values, batch_size=CFG.BATCH_SIZE, device=0, show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True)\n",
    "prompt_embeddings = prompt_embeddings.detach().cpu().numpy()\n",
    "pca_mat = read_VectorTransform('/kaggle/input/wikipedia-stem-pca-index/parsed_gte-large_pca.mat')\n",
    "prompt_embeddings = pca_mat.apply_py(prompt_embeddings)\n",
    "res = faiss.StandardGpuResources()\n",
    "sentence_index = read_index(CFG.INDEX_PATH)\n",
    "sentence_index_gpu = faiss.index_cpu_to_gpu(res, 0, sentence_index)\n",
    "\n",
    "QUERY_SIZE = 8\n",
    "search_score = []\n",
    "search_index = []\n",
    "total = prompt_embeddings.shape[0]\n",
    "for i in tqdm(range(0, total, QUERY_SIZE)):\n",
    "    ss, si = sentence_index_gpu.search(prompt_embeddings[i:i+QUERY_SIZE], CFG.MAX_DOC_NUM)\n",
    "    search_score.append(ss)\n",
    "    search_index.append(si)\n",
    "\n",
    "search_score = np.concatenate(search_score)\n",
    "search_index = np.concatenate(search_index)\n",
    "\n",
    "## Save memory - delete sentence_index since it is no longer necessary\n",
    "del sentence_index\n",
    "del sentence_index_gpu\n",
    "del prompt_embeddings\n",
    "#del model\n",
    "_ = gc.collect()\n",
    "#torch.cuda.empty_cache()\n",
    "libc.malloc_trim(0)\n",
    "\n",
    "df = pd.read_parquet(CFG.WIKI_PLAINTEXT_PATH,\n",
    "                     columns=['text'])\n",
    "\n",
    "contexts = []\n",
    "\n",
    "for i, (scr, idx) in tqdm(enumerate(zip(search_score, search_index)), total=len(search_score)):\n",
    "    context = \"\"\n",
    "    scr_idx = idx\n",
    "    context_list = df.loc[scr_idx].text.tolist()\n",
    "    context += \" \".join(context_list)\n",
    "    contexts.append(context)\n",
    "\n",
    "## Save memory - delete df since it is no longer necessary\n",
    "del df\n",
    "_ = gc.collect()\n",
    "libc.malloc_trim(0)\n",
    "\n",
    "trn['cohere_context'] = contexts\n",
    "\n",
    "save_cols = [\"prompt\", 'parsed_context', \"cohere_context\", \"A\", \"B\", \"C\", \"D\", \"E\"]\n",
    "trn[save_cols].to_csv(\"./test_context.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:08:45.297601Z",
     "iopub.status.busy": "2023-10-04T14:08:45.296282Z",
     "iopub.status.idle": "2023-10-04T14:11:02.746220Z",
     "shell.execute_reply": "2023-10-04T14:11:02.744985Z",
     "shell.execute_reply.started": "2023-10-04T14:08:45.297558Z"
    }
   },
   "outputs": [],
   "source": [
    "!python cohere_context_search.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:11:21.794047Z",
     "iopub.status.busy": "2023-10-04T14:11:21.790543Z",
     "iopub.status.idle": "2023-10-04T14:18:32.287134Z",
     "shell.execute_reply": "2023-10-04T14:18:32.285085Z",
     "shell.execute_reply.started": "2023-10-04T14:11:21.794002Z"
    }
   },
   "outputs": [],
   "source": [
    "#%%time \n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "from dataclasses import dataclass\n",
    "from typing import Optional, Union\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from datasets import Dataset\n",
    "from transformers import AutoTokenizer\n",
    "from transformers import AutoModelForMultipleChoice, TrainingArguments, Trainer\n",
    "from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "PRETAIN = \"/kaggle/input/llm-deberta-models/hf_gte_base_context/checkpoint-35500/\"\n",
    "max_length = 1536\n",
    "\n",
    "test = pd.read_csv(\"test_context.csv\")\n",
    "test.index = list(range(len(test)))\n",
    "test.id = list(range(len(test)))\n",
    "test['answer'] = 'A'\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(PRETAIN)\n",
    "model = AutoModelForMultipleChoice.from_pretrained(PRETAIN).cuda()\n",
    "model = model.half()\n",
    "model.eval();\n",
    "\n",
    "def predictions_to_map_output(predictions):\n",
    "    sorted_answer_indices = np.argsort(-predictions)\n",
    "    top_answer_indices = sorted_answer_indices[:,:3] # Get the first three answers in each row\n",
    "    top_answers = np.vectorize(index_to_option.get)(top_answer_indices)\n",
    "    return np.apply_along_axis(lambda row: ' '.join(row), 1, top_answers)\n",
    "\n",
    "# We'll create a dictionary to convert option names (A, B, C, D, E) into indices and back again\n",
    "options = 'ABCDE'\n",
    "indices = list(range(5))\n",
    "\n",
    "option_to_index = {option: index for option, index in zip(options, indices)}\n",
    "index_to_option = {index: option for option, index in zip(options, indices)}\n",
    "def preprocess(example):\n",
    "    first_sentence = [ \"[CLS] \" + example['context'] ] * 5\n",
    "    second_sentences = [\" #### \" + example['prompt'] + \" [SEP] \" + example[option] + \" [SEP]\" for option in 'ABCDE']\n",
    "    tokenized_example = tokenizer(first_sentence, second_sentences, truncation='only_first', \n",
    "                                  max_length=max_length, add_special_tokens=False)\n",
    "    tokenized_example['label'] = option_to_index[example['answer']]\n",
    "    \n",
    "    return tokenized_example\n",
    "\n",
    "@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\n",
    "    \n",
    "test_df = test[['prompt', 'parsed_context', 'A', 'B', 'C', 'D', 'E', 'answer']].rename(columns={'parsed_context':\"context\"})\n",
    "tokenized_test_dataset = Dataset.from_pandas(test_df[['prompt', 'context', 'A', 'B', 'C', 'D', 'E', 'answer']]).map(preprocess, remove_columns=['prompt', 'context','A', 'B', 'C', 'D', 'E', 'answer'])\n",
    "tokenized_test_dataset = tokenized_test_dataset.remove_columns([\"__index_level_0__\"])\n",
    "data_collator = DataCollatorForMultipleChoice(tokenizer=tokenizer)\n",
    "test_dataloader = DataLoader(tokenized_test_dataset, batch_size=1, shuffle=False, collate_fn=data_collator,num_workers=2,\n",
    "    pin_memory=False,)\n",
    "test_predictions = []\n",
    "for batch in tqdm(test_dataloader):\n",
    "    for k in batch.keys():\n",
    "        batch[k] = batch[k].cuda()\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**batch)\n",
    "    test_predictions.append(outputs.logits.cpu().detach())\n",
    "\n",
    "test_predictions = torch.cat(test_predictions)\n",
    "predictions1 = test_predictions.numpy()\n",
    "\n",
    "test_df = test[['prompt', 'cohere_context', 'A', 'B', 'C', 'D', 'E', 'answer']].rename(columns={'cohere_context':\"context\"})\n",
    "tokenized_test_dataset = Dataset.from_pandas(test_df[['prompt', 'context', 'A', 'B', 'C', 'D', 'E', 'answer']]).map(preprocess, remove_columns=['prompt', 'context','A', 'B', 'C', 'D', 'E', 'answer'])\n",
    "tokenized_test_dataset = tokenized_test_dataset.remove_columns([\"__index_level_0__\"])\n",
    "data_collator = DataCollatorForMultipleChoice(tokenizer=tokenizer)\n",
    "test_dataloader = DataLoader(tokenized_test_dataset, batch_size=1, shuffle=False, collate_fn=data_collator,num_workers=2,\n",
    "    pin_memory=False,)\n",
    "test_predictions = []\n",
    "for batch in tqdm(test_dataloader):\n",
    "    for k in batch.keys():\n",
    "        batch[k] = batch[k].cuda()\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**batch)\n",
    "    test_predictions.append(outputs.logits.cpu().detach())\n",
    "\n",
    "test_predictions = torch.cat(test_predictions)\n",
    "predictions2 = test_predictions.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:20:26.867559Z",
     "iopub.status.busy": "2023-10-04T14:20:26.867131Z",
     "iopub.status.idle": "2023-10-04T14:20:26.872959Z",
     "shell.execute_reply": "2023-10-04T14:20:26.871850Z",
     "shell.execute_reply.started": "2023-10-04T14:20:26.867525Z"
    }
   },
   "outputs": [],
   "source": [
    "predictions2 = predictions2.astype('float32')\n",
    "predictions1 = predictions1.astype('float32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:20:50.682882Z",
     "iopub.status.busy": "2023-10-04T14:20:50.682488Z",
     "iopub.status.idle": "2023-10-04T14:20:50.742888Z",
     "shell.execute_reply": "2023-10-04T14:20:50.741816Z",
     "shell.execute_reply.started": "2023-10-04T14:20:50.682852Z"
    }
   },
   "outputs": [],
   "source": [
    "predictions = []\n",
    "submit_ids = []\n",
    "\n",
    "for index in tqdm(range(len(test))):\n",
    "    probability1 = torch.softmax(torch.tensor(predictions1[index]), dim=-1)\n",
    "    probability2 = torch.softmax(torch.tensor(predictions2[index]), dim=-1)\n",
    "        \n",
    "    probability_ = (probability1 + probability2)/2\n",
    "\n",
    "    if probability_.max() > 0.4:\n",
    "        predict = np.array(list(\"ABCDE\"))[np.argsort(probability_)][-3:].tolist()[::-1]\n",
    "    else:\n",
    "        predict = backup_model_predictions.iloc[index].prediction.replace(\" \",\"\")\n",
    "    predictions.append(predict)\n",
    "\n",
    "predictions = [\" \".join(i) for i in predictions]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-04T14:21:38.499345Z",
     "iopub.status.busy": "2023-10-04T14:21:38.498981Z",
     "iopub.status.idle": "2023-10-04T14:21:38.529281Z",
     "shell.execute_reply": "2023-10-04T14:21:38.528243Z",
     "shell.execute_reply.started": "2023-10-04T14:21:38.499315Z"
    }
   },
   "outputs": [],
   "source": [
    "submission = pd.DataFrame({'id':range(len(test)),'prediction':predictions})\n",
    "submission.to_csv('submission.csv', index=False)\n",
    "\n",
    "import numpy as np\n",
    "def precision_at_k(r, k):\n",
    "    \"\"\"Precision at k\"\"\"\n",
    "    assert k <= len(r)\n",
    "    assert k != 0\n",
    "    return sum(int(x) for x in r[:k]) / k\n",
    "\n",
    "def MAP_at_3(predictions, true_items):\n",
    "    \"\"\"Score is mean average precision at 3\"\"\"\n",
    "    U = len(predictions)\n",
    "    map_at_3 = 0.0\n",
    "    for u in range(U):\n",
    "        user_preds = predictions[u].split()\n",
    "        user_true = true_items[u]\n",
    "        user_results = [1 if item == user_true else 0 for item in user_preds]\n",
    "        for k in range(min(len(user_preds), 3)):\n",
    "            map_at_3 += precision_at_k(user_results, k+1) * user_results[k]\n",
    "    return map_at_3 / U\n",
    "\n",
    "if len(submission) == 200:\n",
    "    train = pd.read_csv('/kaggle/input/kaggle-llm-science-exam/train.csv')\n",
    "    preds = [pred for pred in submission['prediction']]\n",
    "    print(MAP_at_3(preds, train[\"answer\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}