retrain.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['pubid', 'question', 'context', 'long_answer', 'final_decision'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, BertForQuestionAnswering\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load the dataset\n",
    "raw_datasets = load_dataset(\"pubmed_qa\", \"pqa_labeled\")\n",
    "\n",
    "raw_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForQuestionAnswering were not initialized from the model checkpoint at adsabs/astroBERT and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# Load the pre-trained model and tokenizer\n",
    "model_name = \"adsabs/astroBERT\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = BertForQuestionAnswering.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"question\"],  examples[\"final_decision\"], padding=\"max_length\", truncation=True, max_length=206)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['pubid', 'question', 'context', 'long_answer', 'final_decision', 'input_ids', 'token_type_ids', 'attention_mask'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)\n",
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = tokenized_datasets.rename_column(\"question\", \"text\")\n",
    "# tokenized_datasets = tokenized_datasets.rename_column(\"final_decision\", \"labels\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_train_dataset = tokenized_datasets[\"train\"].shuffle(seed=42).select(range(1000)) \n",
    "full_train_dataset = tokenized_datasets[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'pubid': 25432938,\n",
       " 'text': \"Did Chile's traffic law reform push police enforcement?\",\n",
       " 'context': {'contexts': [\"The objective of the current study is to determine to what extent the reduction of Chile's traffic fatalities and injuries during 2000-2012 was related to the police traffic enforcement increment registered after the introduction of its 2005 traffic law reform.\",\n",
       "   \"A unique dataset with assembled information from public institutions and analyses based on ordinary least square and robust random effects models was carried out. Dependent variables were traffic fatality and severe injury rates per population and vehicle fleet. Independent variables were: (1) presence of new national traffic law; (2) police officers per population; (3) number of traffic tickets per police officer; and (4) interaction effect of number of traffic tickets per police officer with traffic law reform. Oil prices, alcohol consumption, proportion of male population 15-24 years old, unemployment, road infrastructure investment, years' effects and regions' effects represented control variables.\",\n",
       "   'Empirical estimates from instrumental variables suggest that the enactment of the traffic law reform in interaction with number of traffic tickets per police officer is significantly associated with a decrease of 8% in traffic fatalities and 7% in severe injuries. Piecewise regression model results for the 2007-2012 period suggest that police traffic enforcement reduced traffic fatalities by 59% and severe injuries by 37%.'],\n",
       "  'labels': ['BACKGROUND', 'METHODS', 'RESULTS'],\n",
       "  'meshes': ['Accidents, Traffic',\n",
       "   'Automobile Driving',\n",
       "   'Chile',\n",
       "   'Humans',\n",
       "   'Law Enforcement',\n",
       "   'Models, Statistical',\n",
       "   'Police',\n",
       "   'Risk Factors'],\n",
       "  'reasoning_required_pred': ['y', 'e', 's'],\n",
       "  'reasoning_free_pred': ['y', 'e', 's']},\n",
       " 'long_answer': 'Findings suggest that traffic law reforms in order to have an effect on both traffic fatality and injury rates reduction require changes in police enforcement practices. Last, this case also illustrates how the diffusion of successful road safety practices globally promoted by WHO and World Bank can be an important influence for enhancing national road safety practices.',\n",
       " 'final_decision': 'yes',\n",
       " 'input_ids': [16338,\n",
       "  9931,\n",
       "  2100,\n",
       "  9401,\n",
       "  6471,\n",
       "  25900,\n",
       "  27660,\n",
       "  2627,\n",
       "  2891,\n",
       "  22510,\n",
       "  25427,\n",
       "  25039,\n",
       "  24565,\n",
       "  2152,\n",
       "  20027,\n",
       "  3767,\n",
       "  8039,\n",
       "  16341,\n",
       "  28567,\n",
       "  16341,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340,\n",
       "  16340],\n",
       " 'token_type_ids': [0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  1,\n",
       "  1,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0],\n",
       " 'attention_mask': [1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0,\n",
       "  0]}"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_train_dataset[10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_train_dataset = full_train_dataset.map(\n",
    "    lambda example: {\"label\": 1 if example[\"final_decision\"] == \"yes\" else 0},\n",
    "    remove_columns=[\"context\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['pubid', 'text', 'long_answer', 'final_decision', 'input_ids', 'token_type_ids', 'attention_mask', 'label']\n",
      "Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None)\n"
     ]
    }
   ],
   "source": [
    "print(full_train_dataset.column_names)\n",
    "print(full_train_dataset.features[\"input_ids\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./qa_finetuned\",\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=8,\n",
    "    save_steps=500,\n",
    "    save_total_limit=2,\n",
    ")\n",
    "\n",
    "# training_args = TrainingArguments(\"test_trainer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/375 [01:19<?, ?it/s]\n",
      "  0%|          | 0/375 [01:12<?, ?it/s]\n",
      "  0%|          | 0/375 [00:00<?, ?it/s]"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "BertForQuestionAnswering.forward() got an unexpected keyword argument 'labels'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[119], line 9\u001b[0m\n\u001b[1;32m      2\u001b[0m trainer \u001b[38;5;241m=\u001b[39m Trainer(\n\u001b[1;32m      3\u001b[0m     model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m      4\u001b[0m     args\u001b[38;5;241m=\u001b[39mtraining_args,\n\u001b[1;32m      5\u001b[0m     train_dataset\u001b[38;5;241m=\u001b[39mfull_train_dataset,\n\u001b[1;32m      6\u001b[0m )\n\u001b[1;32m      8\u001b[0m \u001b[38;5;66;03m# Fine-tune the model\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# Save the fine-tuned model\u001b[39;00m\n\u001b[1;32m     12\u001b[0m model\u001b[38;5;241m.\u001b[39msave_pretrained(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m./qa_finetuned\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/transformers/trainer.py:1537\u001b[0m, in \u001b[0;36mTrainer.train\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   1535\u001b[0m         hf_hub_utils\u001b[38;5;241m.\u001b[39menable_progress_bars()\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1537\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1538\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1539\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1540\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1541\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1542\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/transformers/trainer.py:1854\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[1;32m   1851\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcallback_handler\u001b[38;5;241m.\u001b[39mon_step_begin(args, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol)\n\u001b[1;32m   1853\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maccelerator\u001b[38;5;241m.\u001b[39maccumulate(model):\n\u001b[0;32m-> 1854\u001b[0m     tr_loss_step \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1856\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m   1857\u001b[0m     args\u001b[38;5;241m.\u001b[39mlogging_nan_inf_filter\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_torch_tpu_available()\n\u001b[1;32m   1859\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m (torch\u001b[38;5;241m.\u001b[39misnan(tr_loss_step) \u001b[38;5;129;01mor\u001b[39;00m torch\u001b[38;5;241m.\u001b[39misinf(tr_loss_step))\n\u001b[1;32m   1860\u001b[0m ):\n\u001b[1;32m   1861\u001b[0m     \u001b[38;5;66;03m# if loss is nan or inf simply add the average of previous logged losses\u001b[39;00m\n\u001b[1;32m   1862\u001b[0m     tr_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m tr_loss \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mglobal_step \u001b[38;5;241m-\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_globalstep_last_logged)\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/transformers/trainer.py:2728\u001b[0m, in \u001b[0;36mTrainer.training_step\u001b[0;34m(self, model, inputs)\u001b[0m\n\u001b[1;32m   2725\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m loss_mb\u001b[38;5;241m.\u001b[39mreduce_mean()\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[1;32m   2727\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcompute_loss_context_manager():\n\u001b[0;32m-> 2728\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompute_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2730\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mn_gpu \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m   2731\u001b[0m     loss \u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mmean()  \u001b[38;5;66;03m# mean() to average on multi-gpu parallel training\u001b[39;00m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/transformers/trainer.py:2751\u001b[0m, in \u001b[0;36mTrainer.compute_loss\u001b[0;34m(self, model, inputs, return_outputs)\u001b[0m\n\u001b[1;32m   2749\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2750\u001b[0m     labels \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 2751\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2752\u001b[0m \u001b[38;5;66;03m# Save past state if it exists\u001b[39;00m\n\u001b[1;32m   2753\u001b[0m \u001b[38;5;66;03m# TODO: this needs to be fixed and made cleaner later.\u001b[39;00m\n\u001b[1;32m   2754\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mpast_index \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/benchmarking/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mTypeError\u001b[0m: BertForQuestionAnswering.forward() got an unexpected keyword argument 'labels'"
     ]
    }
   ],
   "source": [
    "\n",
    "# Define the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=full_train_dataset,\n",
    ")\n",
    "\n",
    "# Fine-tune the model\n",
    "trainer.train()\n",
    "\n",
    "# Save the fine-tuned model\n",
    "model.save_pretrained(\"./qa_finetuned\")\n",
    "tokenizer.save_pretrained(\"./qa_finetuned\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "benchmarking",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}