train_models_experiment.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 42\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from datasets import load_dataset, DatasetDict, Dataset\n",
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
    "from sklearn.metrics import accuracy_score, cohen_kappa_score, mean_squared_error\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "from torch import nn\n",
    "from tqdm.auto import tqdm\n",
    "from torch.optim import AdamW\n",
    "from transformers import get_cosine_schedule_with_warmup\n",
    "from lightning.pytorch.core.module import LightningModule\n",
    "from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "from lightning.pytorch.callbacks import Callback\n",
    "from lightning.pytorch.trainer import Trainer\n",
    "from lightning.pytorch.loggers.csv_logs import CSVLogger\n",
    "from coral_pytorch.losses import corn_loss\n",
    "from coral_pytorch.dataset import corn_label_from_logits\n",
    "import os\n",
    "\n",
    "\n",
    "RANDOM_SEED = 42\n",
    "REFERENCE_CONCEPT = 4\n",
    "OBJECTIVE = \"regression\"\n",
    "MAX_LENGTH = 512\n",
    "BATCH_SIZE=16\n",
    "WEIGHT_DECAY=0.01\n",
    "GRADIENT_ACC = 1\n",
    "LEARNING_RATE = 1e-4\n",
    "PRECISION = \"16-mixed\"\n",
    "VARIANT = \"large\"\n",
    "TOKENIZER_NAME = f\"neuralmind/bert-{VARIANT}-portuguese-cased\"\n",
    "MODEL_NAME = f\"neuralmind/bert-{VARIANT}-portuguese-cased\"\n",
    "BASE_MODEL = \"bertimbau\"\n",
    "EXPERIMENT_NAME = f\"aes_enem_models-sourceA-{OBJECTIVE}-from-{BASE_MODEL}-{VARIANT}-C{REFERENCE_CONCEPT+1}\"\n",
    "\n",
    "pl.seed_everything(RANDOM_SEED)\n",
    "torch.set_float32_matmul_precision('medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"kamel-usp/aes_enem_dataset\", \"sourceAWithGraders\", cache_dir=\"/tmp/aes_enem\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'id_prompt', 'essay_title', 'essay_text', 'grades', 'essay_year'],\n",
       "        num_rows: 744\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'id_prompt', 'essay_title', 'essay_text', 'grades', 'essay_year'],\n",
       "        num_rows: 195\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['id', 'id_prompt', 'essay_title', 'essay_text', 'grades', 'essay_year'],\n",
       "        num_rows: 216\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "grade_mapping = {\n",
    "    0: 0,\n",
    "    40: 1,\n",
    "    80: 2,\n",
    "    120: 3,\n",
    "    160: 4,\n",
    "    200: 5,\n",
    "}\n",
    "\n",
    "def create_label(row):\n",
    "    grade = row[\"grades\"][REFERENCE_CONCEPT]\n",
    "    return {\"label\": grade_mapping[grade]}\n",
    "\n",
    "dataset = dataset.map(create_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_difference(lists):\n",
    "    # Assuming the first element is the reference for subtraction\n",
    "    reference = lists[0][REFERENCE_CONCEPT]\n",
    "    grader_a = lists[1][REFERENCE_CONCEPT]\n",
    "    grader_b = lists[2][REFERENCE_CONCEPT]\n",
    "\n",
    "    # Calculate absolute differences\n",
    "    diff_ref_a = abs(reference - grader_a)\n",
    "    diff_ref_b = abs(reference - grader_b)\n",
    "    diff_a_b = abs(grader_a - grader_b)\n",
    "\n",
    "    # Check if any difference is greater than 80\n",
    "    return diff_ref_a > 80 or diff_ref_b > 80 or diff_a_b > 80\n",
    "\n",
    "test_df = dataset[\"test\"].to_pandas()\n",
    "new_test_df = pd.merge(\n",
    "    test_df.groupby([\"id_prompt\", \"id\"]).agg({\"grades\": list}).apply(lambda x: compute_difference(x['grades']), axis=1).reset_index(),\n",
    "    test_df,\n",
    "    on=[\"id_prompt\",\"id\"]\n",
    ").rename(columns={0: \"is_hard\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[\"test_easy\"] = Dataset.from_pandas(new_test_df[new_test_df[\"is_hard\"]==False])\n",
    "dataset[\"test_hard\"] = Dataset.from_pandas(new_test_df[new_test_df[\"is_hard\"]==True])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at neuralmind/bert-large-portuguese-cased and are newly initialized: ['classifier.bias', 'classifier.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": [
    "tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME, use_fast=True)\n",
    "def get_model_instance(model_path, objective):\n",
    "    model = None\n",
    "    if objective == \"regression\":\n",
    "        model = AutoModelForSequenceClassification.from_pretrained(\n",
    "                model_path, \n",
    "                cache_dir=\"/tmp/\", \n",
    "                num_labels=1,\n",
    "            )\n",
    "    elif objective == \"classification\" or objective == \"ordinal\":\n",
    "        model = AutoModelForSequenceClassification.from_pretrained(\n",
    "                model_path, \n",
    "                cache_dir=\"/tmp/\", \n",
    "                num_labels=6,\n",
    "            )\n",
    "    return model\n",
    "model = get_model_instance(MODEL_NAME, OBJECTIVE)\n",
    "if model is None:\n",
    "    raise ValueError(\"Please set a Pre defined Objectice\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(dataset):\n",
    "    def tokenize_essays(dataset, tokenizer, max_length=512):\n",
    "        tokenized_text = tokenizer(\n",
    "                dataset[\"essay_text\"],\n",
    "                return_tensors=\"pt\",\n",
    "                truncation=True,\n",
    "                padding=\"max_length\",\n",
    "                max_length=max_length\n",
    "            )\n",
    "        tokenized_text[\"label\"] = dataset[\"label\"]\n",
    "        return tokenized_text\n",
    "    \n",
    "    tokenized_datasets = {\n",
    "        split: tokenize_essays(sub_dataset, tokenizer, MAX_LENGTH)\n",
    "        for split, sub_dataset in dataset.items()\n",
    "    }\n",
    "    dataset_tokenized = DatasetDict({\n",
    "        split: Dataset.from_dict(data)\n",
    "        for split, data in tokenized_datasets.items()\n",
    "    })\n",
    "\n",
    "    return dataset_tokenized\n",
    "\n",
    "dataset_tokenized = prepare_dataset(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train = DataLoader(\n",
    "    dataset_tokenized[\"train\"].with_format(\"torch\"), batch_size=BATCH_SIZE, shuffle=True, num_workers=0\n",
    ")\n",
    "data_val = DataLoader(dataset_tokenized[\"validation\"].with_format(\"torch\"), batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n",
    "data_test = DataLoader(dataset_tokenized[\"test\"].with_format(\"torch\"), batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n",
    "\n",
    "data_test_easy = DataLoader(dataset_tokenized[\"test_easy\"].with_format(\"torch\"), batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n",
    "data_test_hard = DataLoader(dataset_tokenized[\"test_hard\"].with_format(\"torch\"), batch_size=BATCH_SIZE, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0ae1ce3805d4e21bd5445c462781ea1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def predict_classes(output):\n",
    "    if OBJECTIVE == \"regression\":\n",
    "        # Round the tensor to the nearest integer\n",
    "        rounded_tensor = torch.round(output.logits)\n",
    "        # Clamp the values to the range [0, 5]\n",
    "        clamped_tensor = torch.clamp(rounded_tensor, min=0, max=5)\n",
    "        return clamped_tensor.view(-1)\n",
    "    elif OBJECTIVE == \"classification\":\n",
    "        return torch.argmax(output.logits, axis=1)\n",
    "    elif OBJECTIVE == \"ordinal\":\n",
    "        return corn_label_from_logits(output.logits)\n",
    "        \n",
    "def get_predictions_and_labels(model, dataloader):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "    all_predictions = []\n",
    "    all_true_labels = []\n",
    "    i=0\n",
    "    for batch in tqdm(dataloader, desc=\"Obtaining predictions\"):\n",
    "        input_ids = batch[\"input_ids\"].to(device)\n",
    "        attention_mask = batch[\"attention_mask\"].to(device)\n",
    "        labels = batch[\"label\"].to(device)\n",
    "        with torch.no_grad():\n",
    "            output = model(input_ids, attention_mask)\n",
    "            predicted_classes = predict_classes(output) \n",
    "\n",
    "        # If using GPU, need to move the data back to CPU to use numpy.\n",
    "        all_predictions.extend(predicted_classes.cpu().numpy())\n",
    "        all_true_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "    return list(map(lambda x: x * 40, all_predictions)), list(map(lambda x: x * 40, all_true_labels))\n",
    "\n",
    "all_predictions, all_true_labels = get_predictions_and_labels(model, data_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enem_accuracy_score(true_values, predicted_values):\n",
    "    assert len(true_values) == len(predicted_values), \"Mismatched length between true and predicted values.\"\n",
    "\n",
    "    non_divergent_count = sum([1 for t, p in zip(true_values, predicted_values) if abs(t - p) <= 80])\n",
    "    \n",
    "    return non_divergent_count / len(true_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on the validation  set: 0.05\n"
     ]
    }
   ],
   "source": [
    "accuracy = accuracy_score(all_true_labels, all_predictions)\n",
    "print(f\"Accuracy on the validation  set: {accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QWK on the validation set: 0.03\n"
     ]
    }
   ],
   "source": [
    "qwk = cohen_kappa_score(all_true_labels, all_predictions, weights=\"quadratic\", labels=[0,40,80,120,160,200])\n",
    "print(f\"QWK on the validation set: {qwk:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on the validation set: 0.62\n"
     ]
    }
   ],
   "source": [
    "enem_accuracy = enem_accuracy_score(all_true_labels, all_predictions)\n",
    "print(f\"Accuracy on the validation set: {enem_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LightningGradePredictor(LightningModule):\n",
    "    def __init__(self, model: nn.Module, loss_function, learning_rate=0.001, num_warmup_steps=None, num_training_steps=None):\n",
    "        super(LightningGradePredictor, self).__init__()\n",
    "        self.model = model\n",
    "        self.loss = loss_function\n",
    "        self.learning_rate = learning_rate\n",
    "        self.num_warmup_steps = num_warmup_steps\n",
    "        self.num_training_steps = num_training_steps\n",
    "        self.save_hyperparameters(ignore=['model'])\n",
    "        self.training_step_loss = []\n",
    "        self.validation_step_loss = []\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        return self.model(input_ids, attention_mask)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, labels = batch[\"input_ids\"], batch[\"attention_mask\"], batch[\"label\"]\n",
    "        outputs = self(input_ids, attention_mask)\n",
    "        loss = None\n",
    "        if OBJECTIVE == \"regression\":\n",
    "            loss = self.loss(outputs.logits, labels.unsqueeze(1).float())\n",
    "        elif OBJECTIVE == \"classification\":\n",
    "            loss = self.loss(outputs.logits, labels)\n",
    "        elif OBJECTIVE == \"ordinal\":\n",
    "            loss = self.loss(outputs.logits, labels, num_classes=self.model.num_labels)\n",
    "        self.log('train_loss', loss)\n",
    "        self.training_step_loss.append(loss)\n",
    "        optimizer = self.optimizers()\n",
    "        if isinstance(optimizer, list):\n",
    "            optimizer = optimizer[0]\n",
    "        lr_current = optimizer.optimizer.param_groups[0]['lr']\n",
    "        self.log('current_lr', lr_current)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, labels = batch[\"input_ids\"], batch[\"attention_mask\"], batch[\"label\"]\n",
    "        outputs = self(input_ids, attention_mask)\n",
    "        if OBJECTIVE == \"regression\":\n",
    "            loss = self.loss(outputs.logits, labels.unsqueeze(1).float())\n",
    "        elif OBJECTIVE == \"classification\":\n",
    "            loss = self.loss(outputs.logits, labels)\n",
    "        elif OBJECTIVE == \"ordinal\":\n",
    "            loss = self.loss(outputs.logits, labels, num_classes=self.model.num_labels)\n",
    "        self.log('val_loss', loss)\n",
    "        self.validation_step_loss.append(loss)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, labels = batch[\"input_ids\"], batch[\"attention_mask\"], batch[\"label\"]\n",
    "        outputs = self(input_ids, attention_mask)\n",
    "        if OBJECTIVE == \"regression\":\n",
    "            loss = self.loss(outputs.logits,  labels.unsqueeze(1).float())\n",
    "        elif OBJECTIVE == \"classification\":\n",
    "            loss = self.loss(outputs.logits, labels)\n",
    "        elif OBJECTIVE == \"ordinal\":\n",
    "            loss = self.loss(outputs.logits, labels, num_classes=self.model.num_labels)\n",
    "        self.log('test_loss', loss)\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = AdamW(self.parameters(), lr=self.learning_rate, weight_decay=WEIGHT_DECAY)\n",
    "\n",
    "        scheduler = get_cosine_schedule_with_warmup(\n",
    "            optimizer,\n",
    "            num_warmup_steps=self.num_warmup_steps,\n",
    "            num_training_steps=self.num_training_steps\n",
    "        )\n",
    "\n",
    "        return [optimizer], [{\"scheduler\": scheduler, \"interval\": \"step\"}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EpochEndCallback(Callback):\n",
    "    def __init__(self):\n",
    "        self.metrics_df = pd.DataFrame(columns=\n",
    "                                       ['Epoch', 'Train Loss', 'Validation Loss',\n",
    "                                         'Train QWK', 'Validation QWK',\n",
    "                                         'Train RMSE', 'Validation RMSE'\n",
    "                                        ])\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    def on_train_epoch_end(self, trainer, pl_module):\n",
    "        current_epoch = trainer.current_epoch\n",
    "        # Metrics for training data\n",
    "        train_dataloader = trainer.train_dataloader\n",
    "        val_dataloader = trainer.val_dataloaders\n",
    "        \n",
    "        epoch_train_loss_mean = torch.stack(pl_module.training_step_loss).mean()\n",
    "        epoch_val_loss_mean = torch.stack(pl_module.validation_step_loss).mean()\n",
    "        pl_module.training_step_loss.clear()\n",
    "        pl_module.validation_step_loss.clear()\n",
    "        model.eval()\n",
    "        train_predictions, train_true_labels = get_predictions_and_labels(model, train_dataloader)\n",
    "        val_predictions, val_true_labels = get_predictions_and_labels(model, val_dataloader)\n",
    "        train_qwk = cohen_kappa_score(train_true_labels, train_predictions, weights=\"quadratic\", labels=[0,40,80,120,160,200])\n",
    "        val_qwk = cohen_kappa_score(val_true_labels, val_predictions, weights=\"quadratic\", labels=[0,40,80,120,160,200])\n",
    "        pl_module.log('val_qwk', val_qwk)\n",
    "\n",
    "        train_rmse = mean_squared_error(train_true_labels, train_predictions)\n",
    "        val_rmse = mean_squared_error(val_true_labels, val_predictions)\n",
    "        train_rmse = np.sqrt(train_rmse)\n",
    "        val_rmse = np.sqrt(val_rmse)\n",
    "        pl_module.log('val_rmse', val_rmse)\n",
    "\n",
    "        new_row = {'Epoch': current_epoch, \n",
    "                    'Train Loss': epoch_train_loss_mean.cpu().detach().numpy(),\n",
    "                    'Validation Loss': epoch_val_loss_mean.cpu().detach().numpy(),\n",
    "                    'Train QWK': train_qwk,\n",
    "                    'Validation QWK': val_qwk,\n",
    "                    'Train RMSE': train_rmse,\n",
    "                    'Validation RMSE': val_rmse\n",
    "                    }\n",
    "        new_row = pd.Series(new_row).to_frame().T\n",
    "        self.metrics_df = pd.concat([self.metrics_df, new_row])\n",
    "        display(self.metrics_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "early_stop_callback = EarlyStopping(monitor=\"val_qwk\", patience=3, verbose=True, mode=\"max\")\n",
    "checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor=\"val_qwk\", mode=\"max\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit Automatic Mixed Precision (AMP)\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "def get_loss():\n",
    "    if OBJECTIVE == \"regression\":\n",
    "        return torch.nn.functional.mse_loss\n",
    "    elif OBJECTIVE == \"classification\":\n",
    "        return torch.nn.functional.cross_entropy\n",
    "    elif OBJECTIVE == \"ordinal\":\n",
    "        return corn_loss\n",
    "       \n",
    "loss_function = get_loss() \n",
    "NUM_EPOCHS = 20\n",
    "num_training_steps = NUM_EPOCHS * len(data_train)\n",
    "warmup_steps = int(num_training_steps * 0.1)\n",
    "steps_per_epoch = num_training_steps//NUM_EPOCHS\n",
    "model.train()\n",
    "trainer = Trainer(\n",
    "    max_epochs=NUM_EPOCHS, \n",
    "    log_every_n_steps=steps_per_epoch, \n",
    "    logger=CSVLogger(\"model_logs\", name=EXPERIMENT_NAME), \n",
    "    callbacks=[EpochEndCallback(), early_stop_callback, checkpoint_callback],\n",
    "    precision=PRECISION,\n",
    "    accumulate_grad_batches=GRADIENT_ACC\n",
    "    )\n",
    "ligthning_model = LightningGradePredictor(model, \n",
    "                    loss_function=loss_function,\n",
    "                    num_training_steps=num_training_steps,\n",
    "                    num_warmup_steps=warmup_steps,\n",
    "                    learning_rate = LEARNING_RATE\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing logger folder: model_logs/aes_enem_models-sourceA-regression-from-bertimbau-large-C4\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name  | Type                          | Params\n",
      "--------------------------------------------------------\n",
      "0 | model | BertForSequenceClassification | 334 M \n",
      "--------------------------------------------------------\n",
      "334 M     Trainable params\n",
      "0         Non-trainable params\n",
      "334 M     Total params\n",
      "1,337.590 Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "322b0c3de8ee4e46840493322aa9aaa0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/andrebarbosa/miniconda3/envs/aes_enem/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=47` in the `DataLoader` to improve performance.\n",
      "/home/andrebarbosa/miniconda3/envs/aes_enem/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=47` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "55a58c90ddac4729b2c0f489b4c4a298",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/andrebarbosa/miniconda3/envs/aes_enem/lib/python3.11/site-packages/torch/optim/lr_scheduler.py:136: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  warnings.warn(\"Detected call of `lr_scheduler.step()` before `optimizer.step()`. \"\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aabdc07194954941a994dd49773b9d2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "420f8925b9fa4966a6e0b230373d3e92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d1e665cd04d422e9af5a6709b0a806d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0   3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_qwk improved. New best score: 0.265\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6578e8643eea4263afd14d2f5308300a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2927928729c4deaadfa0497f7b9219b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "49ab06cb3b1f4de9a904f634dd3d381a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0   3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1  1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_qwk improved by 0.078 >= min_delta = 0.0. New best score: 0.343\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fee84e043cb4c79801c9d46f7c5ab54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7741511abfc9488ebbb6dc6c24809234",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbc8515414f4496d8d384df8ad3d8f90",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_qwk improved by 0.048 >= min_delta = 0.0. New best score: 0.391\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e9139af1b414c2a8f12dcdd14afe60b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef0a1eea7f6d423fbdf89b1b5fab2f1b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2549cd7fcbf842458db582ba17e20a16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a9909547d994adcb01d65f6950984ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c21f8b7395d45d391fec235230ae04d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34a5f2758e0b44f290f8e44afab52526",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4</td>\n",
       "      <td>0.85529745</td>\n",
       "      <td>1.4771056</td>\n",
       "      <td>0.570872</td>\n",
       "      <td>0.333728</td>\n",
       "      <td>40.982556</td>\n",
       "      <td>49.115249</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "0     4  0.85529745       1.4771056  0.570872       0.333728  40.982556   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  \n",
       "0       49.115249  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa3b6bf10d384454aab2de8fdcf54a1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a32820afad6749889acfa396abd020a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00b512633d2f45e78b88717f5d38853a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4</td>\n",
       "      <td>0.85529745</td>\n",
       "      <td>1.4771056</td>\n",
       "      <td>0.570872</td>\n",
       "      <td>0.333728</td>\n",
       "      <td>40.982556</td>\n",
       "      <td>49.115249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0.89050436</td>\n",
       "      <td>0.93175906</td>\n",
       "      <td>0.578909</td>\n",
       "      <td>0.435526</td>\n",
       "      <td>35.680783</td>\n",
       "      <td>40.306518</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "0     4  0.85529745       1.4771056  0.570872       0.333728  40.982556   \n",
       "0     5  0.89050436      0.93175906  0.578909       0.435526  35.680783   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  \n",
       "0       49.115249  \n",
       "0       40.306518  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_qwk improved by 0.044 >= min_delta = 0.0. New best score: 0.436\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d93047370e7d4d27a8956475f288c8b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bef17e440864ae39f5b7867ef9b58b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1e34caf4f85412fb787f739a66354f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4</td>\n",
       "      <td>0.85529745</td>\n",
       "      <td>1.4771056</td>\n",
       "      <td>0.570872</td>\n",
       "      <td>0.333728</td>\n",
       "      <td>40.982556</td>\n",
       "      <td>49.115249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0.89050436</td>\n",
       "      <td>0.93175906</td>\n",
       "      <td>0.578909</td>\n",
       "      <td>0.435526</td>\n",
       "      <td>35.680783</td>\n",
       "      <td>40.306518</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>0.83689</td>\n",
       "      <td>0.9720299</td>\n",
       "      <td>0.579828</td>\n",
       "      <td>0.399404</td>\n",
       "      <td>35.861142</td>\n",
       "      <td>42.196123</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "0     4  0.85529745       1.4771056  0.570872       0.333728  40.982556   \n",
       "0     5  0.89050436      0.93175906  0.578909       0.435526  35.680783   \n",
       "0     6     0.83689       0.9720299  0.579828       0.399404  35.861142   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  \n",
       "0       49.115249  \n",
       "0       40.306518  \n",
       "0       42.196123  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a74cf8622924a008d8f47722de05831",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "809aa9f48599458bb3272b080e5b67da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "689099e054234e8e98df26b868ac661e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4</td>\n",
       "      <td>0.85529745</td>\n",
       "      <td>1.4771056</td>\n",
       "      <td>0.570872</td>\n",
       "      <td>0.333728</td>\n",
       "      <td>40.982556</td>\n",
       "      <td>49.115249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0.89050436</td>\n",
       "      <td>0.93175906</td>\n",
       "      <td>0.578909</td>\n",
       "      <td>0.435526</td>\n",
       "      <td>35.680783</td>\n",
       "      <td>40.306518</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>0.83689</td>\n",
       "      <td>0.9720299</td>\n",
       "      <td>0.579828</td>\n",
       "      <td>0.399404</td>\n",
       "      <td>35.861142</td>\n",
       "      <td>42.196123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7</td>\n",
       "      <td>0.725893</td>\n",
       "      <td>0.9039988</td>\n",
       "      <td>0.594298</td>\n",
       "      <td>0.284216</td>\n",
       "      <td>34.172239</td>\n",
       "      <td>43.536015</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "0     4  0.85529745       1.4771056  0.570872       0.333728  40.982556   \n",
       "0     5  0.89050436      0.93175906  0.578909       0.435526  35.680783   \n",
       "0     6     0.83689       0.9720299  0.579828       0.399404  35.861142   \n",
       "0     7    0.725893       0.9039988  0.594298       0.284216  34.172239   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  \n",
       "0       49.115249  \n",
       "0       40.306518  \n",
       "0       42.196123  \n",
       "0       43.536015  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dacd510c3fe44e8286951a3c90e416c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9ddcc2e557f438c9de911baec096ab4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/47 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69814e923124439aa2a8139063936e70",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Train Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Train QWK</th>\n",
       "      <th>Validation QWK</th>\n",
       "      <th>Train RMSE</th>\n",
       "      <th>Validation RMSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3.124451</td>\n",
       "      <td>1.691738</td>\n",
       "      <td>0.272835</td>\n",
       "      <td>0.264975</td>\n",
       "      <td>43.204938</td>\n",
       "      <td>43.347138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1.0187497</td>\n",
       "      <td>0.9265555</td>\n",
       "      <td>0.539333</td>\n",
       "      <td>0.343312</td>\n",
       "      <td>36.100219</td>\n",
       "      <td>42.098785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0.82056665</td>\n",
       "      <td>1.0139564</td>\n",
       "      <td>0.602232</td>\n",
       "      <td>0.391133</td>\n",
       "      <td>36.778909</td>\n",
       "      <td>42.583252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>0.7671431</td>\n",
       "      <td>0.90879524</td>\n",
       "      <td>0.577217</td>\n",
       "      <td>0.308752</td>\n",
       "      <td>34.57888</td>\n",
       "      <td>41.80541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4</td>\n",
       "      <td>0.85529745</td>\n",
       "      <td>1.4771056</td>\n",
       "      <td>0.570872</td>\n",
       "      <td>0.333728</td>\n",
       "      <td>40.982556</td>\n",
       "      <td>49.115249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0.89050436</td>\n",
       "      <td>0.93175906</td>\n",
       "      <td>0.578909</td>\n",
       "      <td>0.435526</td>\n",
       "      <td>35.680783</td>\n",
       "      <td>40.306518</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>0.83689</td>\n",
       "      <td>0.9720299</td>\n",
       "      <td>0.579828</td>\n",
       "      <td>0.399404</td>\n",
       "      <td>35.861142</td>\n",
       "      <td>42.196123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7</td>\n",
       "      <td>0.725893</td>\n",
       "      <td>0.9039988</td>\n",
       "      <td>0.594298</td>\n",
       "      <td>0.284216</td>\n",
       "      <td>34.172239</td>\n",
       "      <td>43.536015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>8</td>\n",
       "      <td>0.70747334</td>\n",
       "      <td>1.3059199</td>\n",
       "      <td>0.538355</td>\n",
       "      <td>0.405419</td>\n",
       "      <td>40.613574</td>\n",
       "      <td>44.097793</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Epoch  Train Loss Validation Loss Train QWK Validation QWK Train RMSE  \\\n",
       "0     0    3.124451        1.691738  0.272835       0.264975  43.204938   \n",
       "0     1   1.0187497       0.9265555  0.539333       0.343312  36.100219   \n",
       "0     2  0.82056665       1.0139564  0.602232       0.391133  36.778909   \n",
       "0     3   0.7671431      0.90879524  0.577217       0.308752   34.57888   \n",
       "0     4  0.85529745       1.4771056  0.570872       0.333728  40.982556   \n",
       "0     5  0.89050436      0.93175906  0.578909       0.435526  35.680783   \n",
       "0     6     0.83689       0.9720299  0.579828       0.399404  35.861142   \n",
       "0     7    0.725893       0.9039988  0.594298       0.284216  34.172239   \n",
       "0     8  0.70747334       1.3059199  0.538355       0.405419  40.613574   \n",
       "\n",
       "  Validation RMSE  \n",
       "0       43.347138  \n",
       "0       42.098785  \n",
       "0       42.583252  \n",
       "0        41.80541  \n",
       "0       49.115249  \n",
       "0       40.306518  \n",
       "0       42.196123  \n",
       "0       43.536015  \n",
       "0       44.097793  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Monitored metric val_qwk did not improve in the last 3 records. Best score: 0.436. Signaling Trainer to stop.\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(model=ligthning_model, train_dataloaders=data_train, val_dataloaders=data_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6fb1470d48b343e18fca4beba039e4d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RMSE on the validation set:  40.3065\n"
     ]
    }
   ],
   "source": [
    "best_model = LightningGradePredictor.load_from_checkpoint(checkpoint_callback.best_model_path, \n",
    "                                                  model=ligthning_model.model, \n",
    "                                                  loss_function=ligthning_model.loss)\n",
    "best_model.model.eval()\n",
    "all_predictions, all_true_labels = get_predictions_and_labels(best_model.model, data_val)\n",
    "rmse_val = mean_squared_error(all_true_labels, all_predictions, squared=False)\n",
    "print(f\"RMSE on the validation set: {rmse_val: .4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(model, dataset, test_group):\n",
    "    all_predictions, all_true_labels = get_predictions_and_labels(model, dataset)\n",
    "    accuracy = accuracy_score(all_true_labels, all_predictions)\n",
    "    qwk = cohen_kappa_score(all_true_labels, all_predictions, weights=\"quadratic\", labels=[0,40,80,120,160,200]) \n",
    "    rmse = mean_squared_error(all_true_labels, all_predictions, squared=False)\n",
    "    horizontal_discrepancy = enem_accuracy_score(all_true_labels, all_predictions)\n",
    "    result = {\n",
    "        'Experiment Reference': EXPERIMENT_NAME,\n",
    "        'Test Group': test_group,\n",
    "        'Competence': REFERENCE_CONCEPT,\n",
    "        'Accuracy': [accuracy],\n",
    "        'RMSE': [rmse],\n",
    "        'QWK': [qwk],\n",
    "        'HDIV': [1- horizontal_discrepancy]\n",
    "    }\n",
    "    return pd.DataFrame(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "881d56ac6f7c4790afd5d6a68264bfa7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Experiment Reference</th>\n",
       "      <th>Test Group</th>\n",
       "      <th>Competence</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>RMSE</th>\n",
       "      <th>QWK</th>\n",
       "      <th>HDIV</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>aes_enem_models-sourceA-regression-from-bertim...</td>\n",
       "      <td>validation</td>\n",
       "      <td>3</td>\n",
       "      <td>0.420513</td>\n",
       "      <td>40.306518</td>\n",
       "      <td>0.435526</td>\n",
       "      <td>0.010256</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                Experiment Reference  Test Group  Competence  \\\n",
       "0  aes_enem_models-sourceA-regression-from-bertim...  validation           3   \n",
       "\n",
       "   Accuracy       RMSE       QWK      HDIV  \n",
       "0  0.420513  40.306518  0.435526  0.010256  "
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_metrics(best_model.model, data_val, 'validation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_report(model, data_test, data_test_easy, data_test_hard):\n",
    "    result = pd.concat([\n",
    "        compute_metrics(model, data_test, \"full\"),\n",
    "        compute_metrics(model, data_test_easy, \"easy\"),\n",
    "        compute_metrics(model, data_test_hard, \"hard\")\n",
    "    ])\n",
    "    directory = \"experiment_reports\"\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "    file_path = f\"{directory}/{EXPERIMENT_NAME}.csv\"\n",
    "    file_exists = os.path.isfile(file_path) and os.path.getsize(file_path) > 0\n",
    "    if file_exists:\n",
    "        result.to_csv(file_path, mode='a', header=False, index=False)\n",
    "    else:\n",
    "        result.to_csv(file_path, mode='w', header=True, index=False)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9232f33c65d1474595fe5b6943c0615c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b84f008a6cf245e899e7f6fa095dd2fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d90e809495d74b21add878ad45e46f35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Experiment Reference</th>\n",
       "      <th>Test Group</th>\n",
       "      <th>Competence</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>RMSE</th>\n",
       "      <th>QWK</th>\n",
       "      <th>HDIV</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>aes_enem_models-sourceA-regression-from-bertim...</td>\n",
       "      <td>full</td>\n",
       "      <td>3</td>\n",
       "      <td>0.439815</td>\n",
       "      <td>43.033148</td>\n",
       "      <td>0.343641</td>\n",
       "      <td>0.041667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>aes_enem_models-sourceA-regression-from-bertim...</td>\n",
       "      <td>easy</td>\n",
       "      <td>3</td>\n",
       "      <td>0.478788</td>\n",
       "      <td>32.659863</td>\n",
       "      <td>0.414441</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>aes_enem_models-sourceA-regression-from-bertim...</td>\n",
       "      <td>hard</td>\n",
       "      <td>3</td>\n",
       "      <td>0.313725</td>\n",
       "      <td>66.273350</td>\n",
       "      <td>-0.016225</td>\n",
       "      <td>0.176471</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                Experiment Reference Test Group  Competence  \\\n",
       "0  aes_enem_models-sourceA-regression-from-bertim...       full           3   \n",
       "0  aes_enem_models-sourceA-regression-from-bertim...       easy           3   \n",
       "0  aes_enem_models-sourceA-regression-from-bertim...       hard           3   \n",
       "\n",
       "   Accuracy       RMSE       QWK      HDIV  \n",
       "0  0.439815  43.033148  0.343641  0.041667  \n",
       "0  0.478788  32.659863  0.414441  0.000000  \n",
       "0  0.313725  66.273350 -0.016225  0.176471  "
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generate_report(best_model.model, data_test, data_test_easy, data_test_hard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'model_checkpoints/hugging_face/finetuning/aes_enem_models-sourceA-regression-from-bertimbau-large-C4'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bert_save_path = f\"model_checkpoints/hugging_face/finetuning/{EXPERIMENT_NAME}\"\n",
    "bert_save_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model.model.save_pretrained(bert_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b60de8b1f5847a58f4fdc733e1cad52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Obtaining predictions:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Experiment Reference</th>\n",
       "      <th>Test Group</th>\n",
       "      <th>Competence</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>RMSE</th>\n",
       "      <th>QWK</th>\n",
       "      <th>HDIV</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>aes_enem_models-sourceA-regression-from-bertim...</td>\n",
       "      <td>full</td>\n",
       "      <td>3</td>\n",
       "      <td>0.439815</td>\n",
       "      <td>43.033148</td>\n",
       "      <td>0.343641</td>\n",
       "      <td>0.041667</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                Experiment Reference Test Group  Competence  \\\n",
       "0  aes_enem_models-sourceA-regression-from-bertim...       full           3   \n",
       "\n",
       "   Accuracy       RMSE       QWK      HDIV  \n",
       "0  0.439815  43.033148  0.343641  0.041667  "
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = get_model_instance(bert_save_path, OBJECTIVE) \n",
    "compute_metrics(model, data_test, \"full\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aes_enem",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}