eedi-submission.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "def438a9",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:27.835417Z",
     "iopub.status.busy": "2024-10-19T18:11:27.834597Z",
     "iopub.status.idle": "2024-10-19T18:11:44.477225Z",
     "shell.execute_reply": "2024-10-19T18:11:44.476028Z"
    },
    "papermill": {
     "duration": 16.65222,
     "end_time": "2024-10-19T18:11:44.479766",
     "exception": false,
     "start_time": "2024-10-19T18:11:27.827546",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in links: /kaggle/input/lmsys-wheel-files\r\n",
      "Requirement already satisfied: transformers in /opt/conda/lib/python3.10/site-packages (4.45.1)\r\n",
      "Processing /kaggle/input/lmsys-wheel-files/peft-0.11.1-py3-none-any.whl\r\n",
      "Requirement already satisfied: accelerate in /opt/conda/lib/python3.10/site-packages (0.34.2)\r\n",
      "Processing /kaggle/input/lmsys-wheel-files/bitsandbytes-0.43.1-py3-none-manylinux_2_24_x86_64.whl\r\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from transformers) (3.15.1)\r\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.23.2 in /opt/conda/lib/python3.10/site-packages (from transformers) (0.25.1)\r\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from transformers) (1.26.4)\r\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from transformers) (21.3)\r\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages (from transformers) (6.0.2)\r\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from transformers) (2024.5.15)\r\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from transformers) (2.32.3)\r\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /opt/conda/lib/python3.10/site-packages (from transformers) (0.4.5)\r\n",
      "Requirement already satisfied: tokenizers<0.21,>=0.20 in /opt/conda/lib/python3.10/site-packages (from transformers) (0.20.0)\r\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.10/site-packages (from transformers) (4.66.4)\r\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.10/site-packages (from peft) (5.9.3)\r\n",
      "Requirement already satisfied: torch>=1.13.0 in /opt/conda/lib/python3.10/site-packages (from peft) (2.4.0)\r\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (2024.6.1)\r\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (4.12.2)\r\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->transformers) (3.1.2)\r\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft) (1.13.3)\r\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft) (3.3)\r\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft) (3.1.4)\r\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->transformers) (3.3.2)\r\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->transformers) (3.7)\r\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->transformers) (1.26.18)\r\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->transformers) (2024.8.30)\r\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch>=1.13.0->peft) (2.1.5)\r\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/conda/lib/python3.10/site-packages (from sympy->torch>=1.13.0->peft) (1.3.0)\r\n",
      "Installing collected packages: bitsandbytes, peft\r\n",
      "Successfully installed bitsandbytes-0.43.1 peft-0.11.1\r\n"
     ]
    }
   ],
   "source": [
    "!pip install transformers peft accelerate bitsandbytes \\\n",
    "    -U --no-index --find-links /kaggle/input/lmsys-wheel-files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b2da7aa5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:44.493640Z",
     "iopub.status.busy": "2024-10-19T18:11:44.493256Z",
     "iopub.status.idle": "2024-10-19T18:11:45.168030Z",
     "shell.execute_reply": "2024-10-19T18:11:45.167066Z"
    },
    "papermill": {
     "duration": 0.68443,
     "end_time": "2024-10-19T18:11:45.170459",
     "exception": false,
     "start_time": "2024-10-19T18:11:44.486029",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ast\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1abebf31",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:45.182828Z",
     "iopub.status.busy": "2024-10-19T18:11:45.182304Z",
     "iopub.status.idle": "2024-10-19T18:11:45.222785Z",
     "shell.execute_reply": "2024-10-19T18:11:45.222028Z"
    },
    "papermill": {
     "duration": 0.049075,
     "end_time": "2024-10-19T18:11:45.225090",
     "exception": false,
     "start_time": "2024-10-19T18:11:45.176015",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "keep_cols = [\"QuestionId\", \"ConstructName\", \"SubjectName\", \"CorrectAnswer\", \"QuestionText\"]\n",
    "answer_cols = [\"AnswerAText\", \"AnswerBText\", \"AnswerCText\", \"AnswerDText\"]\n",
    "misconception_cols = [\"MisconceptionAId\", \"MisconceptionBId\", \"MisconceptionCId\", \"MisconceptionDId\"]\n",
    " \n",
    "    \n",
    "def wide_to_long(df: pd.DataFrame) -> pd.DataFrame:\n",
    "    # Melt the answer columns\n",
    "    answers_df = pd.melt(\n",
    "        id_vars=keep_cols,\n",
    "        frame=df[keep_cols + answer_cols],\n",
    "        var_name='Answer', value_name='Value'\n",
    "    ).sort_values([\"QuestionId\", \"Answer\"]).reset_index(drop=True)\n",
    "    \n",
    "    # If NOT test set\n",
    "    if misconception_cols[0] in df.columns:\n",
    "        \n",
    "        # Melt the misconception columns\n",
    "        misconceptions_df = pd.melt(\n",
    "            id_vars=keep_cols,\n",
    "            frame=df[keep_cols + misconception_cols],\n",
    "            var_name='Misconception', value_name='MisconceptionId'\n",
    "        ).sort_values([\"QuestionId\", \"Misconception\"]).reset_index(drop=True)\n",
    "\n",
    "        answers_df[['Misconception', 'MisconceptionId']] = misconceptions_df[['Misconception', 'MisconceptionId']]\n",
    "    \n",
    "    return answers_df\n",
    "\n",
    "\n",
    "test_data = pd.read_csv(\"/kaggle/input/eedi-mining-misconceptions-in-mathematics/test.csv\")\n",
    "test = wide_to_long(test_data)\n",
    "test.to_csv(\"test_after_process.csv\", index=False)\n",
    "del test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "656398ca",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:45.237971Z",
     "iopub.status.busy": "2024-10-19T18:11:45.237654Z",
     "iopub.status.idle": "2024-10-19T18:11:45.245107Z",
     "shell.execute_reply": "2024-10-19T18:11:45.244174Z"
    },
    "papermill": {
     "duration": 0.015902,
     "end_time": "2024-10-19T18:11:45.247098",
     "exception": false,
     "start_time": "2024-10-19T18:11:45.231196",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing processor.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile processor.py\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import Dict, List, Optional, Union\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "\n",
    "class plain_processor:\n",
    "    def __init__(\n",
    "        self,\n",
    "        tokenizer: AutoTokenizer,\n",
    "        max_length: int,\n",
    "        template: Optional[str] = \"{ConstructName} {QuestionText} {Answer}\",\n",
    "        add_eos_token: bool = True,\n",
    "    ):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.template = template\n",
    "        self.add_eos_token = add_eos_token\n",
    "        self.max_length = max_length - 1 if add_eos_token else max_length\n",
    "\n",
    "    def preprocess_batch(self, batch_data: Dict[str, List[str]]):\n",
    "        subject_name = batch_data[\"SubjectName\"]\n",
    "        contruct_name = batch_data[\"ConstructName\"]\n",
    "        question_text = batch_data[\"QuestionText\"]\n",
    "        answer_text = batch_data[\"Value\"]\n",
    "        return subject_name, contruct_name, question_text, answer_text\n",
    "    \n",
    "    def format_texts(self, subject_name, contruct_name, question_text, answer_text):\n",
    "        texts = []\n",
    "        for subj, cont, ques, ans in zip(subject_name, contruct_name, question_text, answer_text):\n",
    "            data_dic = {\n",
    "                \"SubjectName\": subj,\n",
    "                \"ConstructName\": cont,\n",
    "                \"QuestionText\": ques,\n",
    "                \"Answer\": ans,\n",
    "            }\n",
    "            text = self.template.format_map(data_dic)\n",
    "            texts.append(text)\n",
    "        return texts\n",
    "    \n",
    "    def __call__(self, batch_data):\n",
    "        batch = self.preprocess_batch(batch_data)\n",
    "        texts = self.format_texts(*batch)\n",
    "        results = defaultdict(list)\n",
    "        outputs = self.tokenizer(\n",
    "            texts,\n",
    "            max_length=self.max_length,\n",
    "            truncation=True,\n",
    "        )\n",
    "        \n",
    "        for input_ids, attention_mask in zip(outputs['input_ids'], outputs['attention_mask']):\n",
    "            input_ids.append(self.tokenizer.eos_token_id)\n",
    "            attention_mask.append(1)\n",
    "\n",
    "        results[\"input_ids\"] = outputs[\"input_ids\"]\n",
    "        results[\"attention_mask\"] = outputs[\"attention_mask\"]\n",
    "\n",
    "        return results\n",
    "    \n",
    "\n",
    "class misconception_processor:\n",
    "    def __init__(\n",
    "        self,\n",
    "        tokenizer: AutoTokenizer,\n",
    "        max_length: int,\n",
    "        add_eos_token: bool = True,\n",
    "    ):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.add_eos_token = add_eos_token\n",
    "        self.max_length = max_length - 1 if add_eos_token else max_length\n",
    "    \n",
    "    def __call__(self, batch_data):\n",
    "        batch_data = batch_data[\"MisconceptionName\"]\n",
    "        results = defaultdict(list)\n",
    "        outputs = self.tokenizer(\n",
    "            batch_data,\n",
    "            max_length=self.max_length,\n",
    "            truncation=True,\n",
    "        )\n",
    "\n",
    "        for input_ids, attention_mask in zip(outputs['input_ids'], outputs['attention_mask']):\n",
    "            input_ids.append(self.tokenizer.eos_token_id)\n",
    "            attention_mask.append(1)\n",
    "\n",
    "        results[\"input_ids\"] = outputs[\"input_ids\"]\n",
    "        results[\"attention_mask\"] = outputs[\"attention_mask\"]\n",
    "\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "853d4b89",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:45.259740Z",
     "iopub.status.busy": "2024-10-19T18:11:45.259434Z",
     "iopub.status.idle": "2024-10-19T18:11:45.268536Z",
     "shell.execute_reply": "2024-10-19T18:11:45.267613Z"
    },
    "papermill": {
     "duration": 0.018023,
     "end_time": "2024-10-19T18:11:45.270529",
     "exception": false,
     "start_time": "2024-10-19T18:11:45.252506",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing inference.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile inference.py\n",
    "\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Dict, List, Optional, Union\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from datasets import Dataset\n",
    "from transformers import (\n",
    "    HfArgumentParser,\n",
    "    AutoTokenizer,\n",
    "    AutoModel,\n",
    "    DataCollatorWithPadding,\n",
    ")\n",
    "from peft import PeftModel\n",
    "from processor import plain_processor, misconception_processor\n",
    "\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "@dataclass\n",
    "class ModelArguments:\n",
    "    model_name_or_path: Optional[str] = field(default=\"BAAI/bge-base-en-v1.5\")\n",
    "    model_max_length: int = field(\n",
    "        default=1024,\n",
    "        metadata={\n",
    "            \"help\": \"Maximum sequence length. Sequences will be right padded (and possibly truncated).\"\n",
    "        },\n",
    "    )\n",
    "    half_precision: bool = field(default=True, metadata={\"help\": \"Whether to use half precision.\"})\n",
    "    add_eos_token: bool = field(default=False)\n",
    "    lora_dir: Optional[str] = field(default=\"/kaggle/input/qwen2-5-1-5b-retrieval/checkpoint-525\")\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataArguments:\n",
    "    train_data_path: str = field(\n",
    "        default=\"data/train_after_process.csv\", metadata={\"help\": \"Path to the training data.\"}\n",
    "    )\n",
    "    misconception_mapping: str = field(\n",
    "        default=\"data/misconception_mapping.csv\", metadata={\"help\": \"Path to the misconception mapping.\"}\n",
    "    )\n",
    "    template: str = field(\n",
    "        default=\"{ConstructName} {QuestionText} {Answer}\", metadata={\"help\": \"Template for the input text.\"}\n",
    "    )\n",
    "    top_k_for_recall: int =  field(default=25, metadata={\"help\": \"Remain top k in recall stage.\"})\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TrainingArguments:\n",
    "    batch_size: int = field(default=8, metadata={\"help\": \"Batch size per GPU for inference.\"})\n",
    "\n",
    "\n",
    "def last_token_pool(\n",
    "    last_hidden_states: Tensor,            \n",
    "    attention_mask: Tensor\n",
    ") -> Tensor:\n",
    "    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])\n",
    "    if left_padding:\n",
    "        return last_hidden_states[:, -1]\n",
    "    else:\n",
    "        sequence_lengths = attention_mask.sum(dim=1) - 1\n",
    "        batch_size = last_hidden_states.shape[0]\n",
    "        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]\n",
    "    \n",
    "\n",
    "@torch.no_grad()\n",
    "@torch.amp.autocast('cuda')\n",
    "def inference(model, dataset, data_collator, batch_size: int = 8):\n",
    "    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=data_collator)\n",
    "    embeddings = []\n",
    "    for batch in tqdm(data_loader):\n",
    "        if \"labels\" in batch.keys():\n",
    "            batch.pop(\"labels\")\n",
    "        batch = {k: v.to(model.device) for k, v in batch.items()}\n",
    "        sentence_embeddings = model(**batch).last_hidden_state\n",
    "        sentence_embeddings = last_token_pool(sentence_embeddings, batch['attention_mask'])\n",
    "        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)\n",
    "        embeddings.append(sentence_embeddings.detach().cpu().numpy())\n",
    "    return np.concatenate(embeddings, axis=0)\n",
    "\n",
    "\n",
    "def main():\n",
    "    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))\n",
    "    model_args, data_args, training_args = parser.parse_args_into_dataclasses()\n",
    "\n",
    "    # prepare tokenizer and model\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\n",
    "        model_args.model_name_or_path,\n",
    "        padding_side=\"right\",\n",
    "        use_fast=True,\n",
    "    )\n",
    "    model = AutoModel.from_pretrained(\n",
    "        model_args.model_name_or_path,\n",
    "        device_map=\"auto\",\n",
    "        torch_dtype=torch.float16 if model_args.half_precision else torch.float32,\n",
    "    )\n",
    "    \n",
    "    model = PeftModel.from_pretrained(model, model_args.lora_dir)\n",
    "    model.eval()\n",
    "\n",
    "    # prepare data\n",
    "    # TODO: Inference for test data.\n",
    "    train_dataset = Dataset.from_csv(data_args.train_data_path)\n",
    "    preprocess = plain_processor(tokenizer, model_args.model_max_length, template=data_args.template)\n",
    "    train_dataset = train_dataset.map(preprocess, batched=True, remove_columns=train_dataset.column_names)\n",
    "\n",
    "    misconception_mapping = Dataset.from_csv(data_args.misconception_mapping)\n",
    "    mis_preprocess = misconception_processor(tokenizer, model_args.model_max_length)\n",
    "    misconception_mapping = misconception_mapping.map(mis_preprocess, batched=True, remove_columns=misconception_mapping.column_names)\n",
    "    \n",
    "    # inference\n",
    "    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "    train_embeddings = inference(\n",
    "        model, train_dataset, batch_size=training_args.batch_size, data_collator=data_collator\n",
    "    )\n",
    "    misconception_embeddings = inference(\n",
    "        model, misconception_mapping, batch_size=training_args.batch_size, data_collator=data_collator\n",
    "    )\n",
    "\n",
    "    # calculate cosine similarity\n",
    "    cos_sim_arr = cosine_similarity(train_embeddings, misconception_embeddings)\n",
    "    sorted_indices = np.argsort(-cos_sim_arr, axis=1)\n",
    "    sorted_indices = sorted_indices[:, :data_args.top_k_for_recall].tolist()\n",
    "    \n",
    "    return sorted_indices\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_sorted_indices = main()\n",
    "    \n",
    "    test_data = pd.read_csv(\"test_after_process.csv\")\n",
    "    test_data[\"Answer_alphabet\"] = test_data[\"Answer\"].str.extract(r'Answer([A-Z])Text$')\n",
    "    test_data[\"QuestionId_Answer\"] = test_data[\"QuestionId\"].astype(\"str\") + \"_\" + test_data[\"Answer_alphabet\"]\n",
    "    test_data[\"MisconceptionId\"] = test_sorted_indices\n",
    "    # filter correct row\n",
    "    test_data = test_data[test_data[\"CorrectAnswer\"] != test_data[\"Answer_alphabet\"]]\n",
    "    test_data = test_data.sort_values(by=[\"QuestionId\", \"Answer_alphabet\"])\n",
    "    \n",
    "    # submission = test_data[[\"QuestionId_Answer\", \"MisconceptionId\"]]\n",
    "    # submission.loc[:, \"MisconceptionId\"] = submission[\"MisconceptionId\"].apply(lambda x: ' '.join(map(str, x[:25])))\n",
    "    # submission.to_csv('submission.csv', index=False)\n",
    "    test_data.to_csv(\"retrieve_result.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "91653fb7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:11:45.284062Z",
     "iopub.status.busy": "2024-10-19T18:11:45.283757Z",
     "iopub.status.idle": "2024-10-19T18:15:13.976212Z",
     "shell.execute_reply": "2024-10-19T18:15:13.975126Z"
    },
    "papermill": {
     "duration": 208.701225,
     "end_time": "2024-10-19T18:15:13.978627",
     "exception": false,
     "start_time": "2024-10-19T18:11:45.277402",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████████████| 4/4 [01:08<00:00, 17.00s/it]\r\n",
      "Generating train split: 12 examples [00:00, 487.43 examples/s]\r\n",
      "Map: 100%|██████████████████████████████| 12/12 [00:00<00:00, 683.90 examples/s]\r\n",
      "Generating train split: 2587 examples [00:00, 271083.63 examples/s]\r\n",
      "Map: 100%|████████████████████████| 2587/2587 [00:00<00:00, 24064.41 examples/s]\r\n",
      "100%|█████████████████████████████████████████████| 2/2 [00:01<00:00,  1.15it/s]\r\n",
      "100%|█████████████████████████████████████████| 324/324 [01:52<00:00,  2.87it/s]\r\n"
     ]
    }
   ],
   "source": [
    "!python inference.py \\\n",
    "    --model_name_or_path /kaggle/input/qwen2-math-7b-instruct \\\n",
    "    --model_max_length 1024 \\\n",
    "    --half_precision False \\\n",
    "    --train_data_path test_after_process.csv \\\n",
    "    --misconception_mapping /kaggle/input/eedi-mining-misconceptions-in-mathematics/misconception_mapping.csv \\\n",
    "    --batch_size 8 \\\n",
    "    --template \"{ConstructName} {QuestionText} {Answer}\" \\\n",
    "    --lora_dir \"/kaggle/input/qwen2-5-1-5b-retrieval/checkpoint-474\" \\\n",
    "    --top_k_for_recall 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2a744977",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:15:14.047369Z",
     "iopub.status.busy": "2024-10-19T18:15:14.046987Z",
     "iopub.status.idle": "2024-10-19T18:15:14.051933Z",
     "shell.execute_reply": "2024-10-19T18:15:14.050792Z"
    },
    "papermill": {
     "duration": 0.042019,
     "end_time": "2024-10-19T18:15:14.054114",
     "exception": false,
     "start_time": "2024-10-19T18:15:14.012095",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# sub = pd.read_csv(\"submission.csv\")\n",
    "# sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "29d2bc1d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:15:14.123690Z",
     "iopub.status.busy": "2024-10-19T18:15:14.122731Z",
     "iopub.status.idle": "2024-10-19T18:15:14.180807Z",
     "shell.execute_reply": "2024-10-19T18:15:14.179689Z"
    },
    "papermill": {
     "duration": 0.095597,
     "end_time": "2024-10-19T18:15:14.183544",
     "exception": false,
     "start_time": "2024-10-19T18:15:14.087947",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "retrieve_result = pd.read_csv(\"retrieve_result.csv\")\n",
    "retrieve_result[\"MisconceptionId\"] = retrieve_result[\"MisconceptionId\"].apply(ast.literal_eval)\n",
    "retrieve_result = retrieve_result.explode(\"MisconceptionId\")\n",
    "misconception_mapping = pd.read_csv(\"/kaggle/input/eedi-mining-misconceptions-in-mathematics/misconception_mapping.csv\")\n",
    "retrieve_result = retrieve_result.merge(misconception_mapping, how=\"left\", left_on=\"MisconceptionId\", right_on=\"MisconceptionId\")\n",
    "retrieve_result.to_csv(\"retrieve_result.csv\", index=False)\n",
    "del retrieve_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eb84367d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:15:14.252492Z",
     "iopub.status.busy": "2024-10-19T18:15:14.252109Z",
     "iopub.status.idle": "2024-10-19T18:15:14.260061Z",
     "shell.execute_reply": "2024-10-19T18:15:14.259058Z"
    },
    "papermill": {
     "duration": 0.044373,
     "end_time": "2024-10-19T18:15:14.262771",
     "exception": false,
     "start_time": "2024-10-19T18:15:14.218398",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing reranker_processor.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile reranker_processor.py\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import Dict, List, Optional, Union\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "\n",
    "class reranker_processor:\n",
    "    def __init__(\n",
    "        self,\n",
    "        tokenizer: AutoTokenizer,\n",
    "        max_length: int,\n",
    "        template: Optional[str] = \"<|im_start|>{ConstructName} {QuestionText} {Answer}\\\\n{DOC}<|im_end|>\",\n",
    "    ):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "        self.template = template\n",
    "\n",
    "    def preprocess_batch(self, batch_data: Dict[str, List[str]]):\n",
    "        subject_name = batch_data[\"SubjectName\"]\n",
    "        contruct_name = batch_data[\"ConstructName\"]\n",
    "        question_text = batch_data[\"QuestionText\"]\n",
    "        answer_text = batch_data[\"Value\"]\n",
    "        misconception = batch_data[\"MisconceptionName\"]\n",
    "        return subject_name, contruct_name, question_text, answer_text, misconception\n",
    "    \n",
    "    def format_texts(self, subject_name, contruct_name, question_text, answer_text, misconception):\n",
    "        texts = []\n",
    "        for subj, cont, ques, ans, misc in zip(subject_name, contruct_name, question_text, answer_text, misconception):\n",
    "            data_dic = {\n",
    "                \"SubjectName\": subj,\n",
    "                \"ConstructName\": cont,\n",
    "                \"QuestionText\": ques,\n",
    "                \"Answer\": ans,\n",
    "                \"DOC\": misc,\n",
    "            }\n",
    "            text = self.template.format_map(data_dic)\n",
    "            texts.append(text)\n",
    "        return texts\n",
    "    \n",
    "    def __call__(self, batch_data):\n",
    "        batch = self.preprocess_batch(batch_data)\n",
    "        texts = self.format_texts(*batch)\n",
    "        results = defaultdict(list)\n",
    "        outputs = self.tokenizer(\n",
    "            texts,\n",
    "            max_length=self.max_length,\n",
    "            truncation=True,\n",
    "        )\n",
    "\n",
    "        results[\"input_ids\"] = outputs[\"input_ids\"]\n",
    "        results[\"attention_mask\"] = outputs[\"attention_mask\"]\n",
    "\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "db50fc43",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:15:14.336586Z",
     "iopub.status.busy": "2024-10-19T18:15:14.336131Z",
     "iopub.status.idle": "2024-10-19T18:15:14.348221Z",
     "shell.execute_reply": "2024-10-19T18:15:14.347057Z"
    },
    "papermill": {
     "duration": 0.050157,
     "end_time": "2024-10-19T18:15:14.350356",
     "exception": false,
     "start_time": "2024-10-19T18:15:14.300199",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing reranker_inference.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile reranker_inference.py\n",
    "\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Dict, List, Optional, Union\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from datasets import Dataset\n",
    "from peft import PeftModel\n",
    "from transformers import (\n",
    "    HfArgumentParser,\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    "    DataCollatorWithPadding,\n",
    ")\n",
    "\n",
    "from reranker_processor import reranker_processor\n",
    "\n",
    "@dataclass\n",
    "class ModelArguments:\n",
    "    model_name_or_path: Optional[str] = field(default=\"BAAI/bge-base-en-v1.5\")\n",
    "    model_max_length: int = field(\n",
    "        default=1024,\n",
    "        metadata={\n",
    "            \"help\": \"Maximum sequence length. Sequences will be right padded (and possibly truncated).\"\n",
    "        },\n",
    "    )\n",
    "    half_precision: bool = field(default=True, metadata={\"help\": \"Whether to use half precision.\"})\n",
    "    add_eos_token: bool = field(default=False)\n",
    "    lora_dir: str = field(default=\"/kaggle/input/qwen2-math-1-5b-it/checkpoint-1887\")\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataArguments:\n",
    "    train_data_path: str = field(\n",
    "        default=\"data/train_after_process.csv\", metadata={\"help\": \"Path to the training data.\"}\n",
    "    )\n",
    "    template: str = field(\n",
    "        default=\"{ConstructName} {QuestionText} {Answer}\", metadata={\"help\": \"Template for the input text.\"}\n",
    "    )\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TrainingArguments:\n",
    "    batch_size: int = field(default=8, metadata={\"help\": \"Batch size per GPU for inference.\"})\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "@torch.amp.autocast('cuda')\n",
    "def inference(model, dataset, data_collator, batch_size: int = 8):\n",
    "    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=data_collator)\n",
    "    probabilities = []\n",
    "    for batch in tqdm(data_loader):\n",
    "        batch = {k: v.to(model.device) for k, v in batch.items()}\n",
    "        outputs = model(**batch).logits\n",
    "        proba = outputs.sigmoid().detach().cpu().numpy()\n",
    "        probabilities.extend(proba[:, 0].astype(np.float32))\n",
    "    return probabilities\n",
    "\n",
    "\n",
    "def main():\n",
    "    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))\n",
    "    model_args, data_args, training_args = parser.parse_args_into_dataclasses()\n",
    "\n",
    "    # prepare tokenizer and model\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\n",
    "        model_args.model_name_or_path,\n",
    "        padding_side=\"right\",\n",
    "        use_fast=True,\n",
    "        add_eos_token=model_args.add_eos_token,\n",
    "    )\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        model_args.model_name_or_path,\n",
    "        num_labels=1,\n",
    "        device_map=\"auto\",\n",
    "        torch_dtype=torch.float16 if model_args.half_precision else torch.float32,\n",
    "    )\n",
    "    model = PeftModel.from_pretrained(model, model_args.lora_dir)\n",
    "    model.eval()\n",
    "    \n",
    "    if \"llama\" in model_args.model_name_or_path.lower():\n",
    "        tokenizer.pad_token = tokenizer.eos_token\n",
    "        model.config.pad_token_id = tokenizer.pad_token_id\n",
    "    \n",
    "    if \"qwen\" in model_args.model_name_or_path.lower():\n",
    "        tokenizer.pad_token = \"<|endoftext|>\"\n",
    "        model.config.pad_token_id = tokenizer.pad_token_id\n",
    "\n",
    "    # prepare data\n",
    "    # TODO: Inference for test data.\n",
    "    train_dataset = Dataset.from_csv(data_args.train_data_path)\n",
    "    preprocess = reranker_processor(tokenizer, model_args.model_max_length, template=data_args.template)\n",
    "    train_dataset = train_dataset.map(preprocess, batched=True, remove_columns=train_dataset.column_names)\n",
    "    \n",
    "    # inference\n",
    "    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "    probabilities = inference(\n",
    "        model, train_dataset, batch_size=training_args.batch_size, data_collator=data_collator\n",
    "    )\n",
    "    \n",
    "    return probabilities\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    probabilities = main()\n",
    "    test_data = pd.read_csv(\"retrieve_result.csv\")\n",
    "    test_data[\"pred_prob\"] = probabilities\n",
    "    test_data = test_data.sort_values(by=[\"QuestionId_Answer\", \"pred_prob\"], ascending=[True, False])\n",
    "    test_data = test_data.groupby(\"QuestionId_Answer\", as_index=False).agg({\"MisconceptionId\": lambda x: list(x)})\n",
    "    \n",
    "    submission = test_data[[\"QuestionId_Answer\", \"MisconceptionId\"]]\n",
    "    submission[\"MisconceptionId\"] = submission[\"MisconceptionId\"].apply(lambda x: ' '.join(map(str, x[:25])))\n",
    "    submission = submission.sort_values(\"QuestionId_Answer\")\n",
    "    submission.to_csv('submission.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e59c10f9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-19T18:15:14.420461Z",
     "iopub.status.busy": "2024-10-19T18:15:14.419559Z",
     "iopub.status.idle": "2024-10-19T18:15:53.963880Z",
     "shell.execute_reply": "2024-10-19T18:15:53.962911Z"
    },
    "papermill": {
     "duration": 39.580773,
     "end_time": "2024-10-19T18:15:53.966235",
     "exception": false,
     "start_time": "2024-10-19T18:15:14.385462",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some weights of Qwen2ForSequenceClassification were not initialized from the model checkpoint at /kaggle/input/qwen2-math-1-5b-instruct/Qwen2-Math-1.5B-Instruct and are newly initialized: ['score.weight']\r\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\r\n",
      "Generating train split: 450 examples [00:00, 37242.98 examples/s]\r\n",
      "Map: 100%|███████████████████████████| 450/450 [00:00<00:00, 5252.51 examples/s]\r\n",
      "100%|███████████████████████████████████████████| 29/29 [00:11<00:00,  2.56it/s]\r\n"
     ]
    }
   ],
   "source": [
    "!python reranker_inference.py \\\n",
    "    --model_name_or_path /kaggle/input/qwen2-math-1-5b-instruct/Qwen2-Math-1.5B-Instruct \\\n",
    "    --model_max_length 1024 \\\n",
    "    --half_precision False \\\n",
    "    --train_data_path retrieve_result.csv \\\n",
    "    --batch_size 16 \\\n",
    "    --template \"<|im_start|>{ConstructName} {QuestionText} {Answer}\\\\n{DOC}<|im_end|>\" \\\n",
    "    --lora_dir /kaggle/input/qwen2-math-1-5b-it/checkpoint-378"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "databundleVersionId": 9738540,
     "sourceId": 82695,
     "sourceType": "competition"
    },
    {
     "datasetId": 5297895,
     "sourceId": 8897601,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 5724358,
     "sourceId": 9669274,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 5724449,
     "sourceId": 9424171,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 5756280,
     "sourceId": 9466816,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 5830057,
     "sourceId": 9565875,
     "sourceType": "datasetVersion"
    },
    {
     "datasetId": 5756511,
     "sourceId": 9667235,
     "sourceType": "datasetVersion"
    }
   ],
   "isGpuEnabled": true,
   "isInternetEnabled": false,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 269.224903,
   "end_time": "2024-10-19T18:15:54.421535",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-10-19T18:11:25.196632",
   "version": "2.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}