train_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from torch.utils.data import random_split, Dataset, DataLoader\n",
    "\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import precision_recall_curve\n",
    "\n",
    "from model import build_transformer\n",
    "\n",
    "torch.set_default_device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_transformer(dropout=0.1,\n",
    "\t\t\t\t\t\tsource_vocab_size=8000+1, target_vocab_size=292, context_length=900 - 3 + 1,\n",
    "\t\t\t\t\t\tdecoder_block_count=6,\n",
    "\t\t\t\t\t\tdecoder_self_attention_head_count=8,\n",
    "\t\t\t\t\t\tdecoder_self_attention_abstraction_coef=0.15,\n",
    "\t\t\t\t\t\tdecoder_feed_forward_abstraction_coef=4,\n",
    "\t\t\t\t\t\tdim=256, epsilon=1e-9)\n",
    "\n",
    "checkpoint = torch.load(\"weights/tr_model_10\", weights_only=True)\n",
    "model.load_state_dict(checkpoint[\"state\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.load('X.pt', weights_only=True).int().to(\"cuda\")\n",
    "L = torch.load('L.pt', weights_only=True).int().to(\"cuda\")\n",
    "Y = torch.load('Y.pt', weights_only=True).half().to(\"cuda\")\n",
    "\n",
    "train_ratio = 0.9\n",
    "train_size = int(len(X) * train_ratio)\n",
    "test_size = len(X) - train_size\n",
    "\n",
    "generator = torch.Generator(device=\"cuda\").manual_seed(42)\n",
    "X_train, X_test = random_split(X, [train_size, test_size], generator)\n",
    "L_train, L_test = random_split(L, [train_size, test_size], generator)\n",
    "Y_train, Y_test = random_split(Y, [train_size, test_size], generator)\n",
    "\n",
    "print(len(X_train), len(X_test))\n",
    "print(len(L_train), len(L_test))\n",
    "print(len(Y_train), len(Y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, X, L, Y):\n",
    "        self.X = X\n",
    "        self.L = L\n",
    "        self.Y = Y\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.X[idx], self.L[idx], self.Y[idx]\n",
    "\n",
    "\n",
    "dataset = CustomDataset(X_train, L_train, Y_train)\n",
    "loader = DataLoader(dataset, batch_size=64, shuffle=True, generator=torch.Generator(device='cuda'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = len(dataset)\n",
    "num_classes = 292\n",
    "context_length = 900 - 3 + 1\n",
    "\n",
    "preds = torch.empty(num_samples, num_classes, dtype=torch.float)\n",
    "real = torch.empty(num_samples, num_classes, dtype=torch.float)\n",
    "\n",
    "masks = torch.ones((context_length, 1, context_length, context_length))\n",
    "for i in range(context_length):\n",
    "\tmasks[i, :, :, i:] = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "\tmodel.eval()\n",
    "\n",
    "\tcount = 0\n",
    "\tbatch_iterator = tqdm(loader, desc=f\"Processing batches\")\n",
    "\tfor x, l, y in batch_iterator:\n",
    "\t\tm = masks[l - 1]\n",
    "\t\tpred = model(x, mask=m)\n",
    "\n",
    "\t\tpreds[count:count + len(x)] = pred\n",
    "\t\treal[count:count + len(x)] = y\n",
    "\n",
    "\t\tcount += len(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the predictions\n",
    "torch.save(preds, \"P.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = preds.cpu().numpy()\n",
    "real = real.cpu().numpy()\n",
    "real = real.astype(np.int32)\n",
    "\n",
    "meilleurs_seuils = []\n",
    "f1_scores = []\n",
    "\n",
    "# Compute the best threshold for each class\n",
    "for c in range(num_classes):\n",
    "    valeurs_c = preds[:, c]\n",
    "    etiquettes_c = real[:, c]\n",
    "    precisions, rappels, seuils = precision_recall_curve(etiquettes_c, valeurs_c)\n",
    "    scores_f1 = 2 * (precisions[:-1] * rappels[:-1]) / (precisions[:-1] + rappels[:-1] + 1e-8)\n",
    "    indice_max = np.argmax(scores_f1)\n",
    "    meilleur_seuil = seuils[indice_max]\n",
    "    meilleurs_seuils.append(meilleur_seuil)\n",
    "    f1_scores.append(scores_f1[indice_max])\n",
    "\n",
    "print(meilleurs_seuils[0:10])\n",
    "print(f1_scores[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_thresholds = np.array(meilleurs_seuils)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the thresholds as a tensor\n",
    "torch.save(torch.tensor(optimal_thresholds), \"T.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute for each class (TP / (TP + FN)), (TP / (TP + FP)), ((VP + VN) / (VP + VN + FP + FN))\n",
    "\n",
    "precision = np.zeros(num_classes)\n",
    "recall = np.zeros(num_classes)\n",
    "accuracy = np.zeros(num_classes)\n",
    "\n",
    "for i in range(num_classes):\n",
    "\tpred_binary = (preds[:, i] >= optimal_thresholds[i]).astype(int)\n",
    "\tTP = np.sum(pred_binary & real[:, i])\n",
    "\tFP = np.sum(pred_binary & (1 - real[:, i]))\n",
    "\tFN = np.sum((1 - pred_binary) & real[:, i])\n",
    "\tTN = np.sum((1 - pred_binary) & (1 - real[:, i]))\n",
    "\n",
    "\tprecision[i] = TP / (TP + FP)\n",
    "\trecall[i] = TP / (TP + FN)\n",
    "\taccuracy[i] = (TP + TN) / (TP + TN + FP + FN)\n",
    "\n",
    "print(\"Average precision:\", np.mean(precision))\n",
    "print(\"Average recall:\", np.mean(recall))\n",
    "print(\"Average accuracy:\", np.mean(accuracy))\n",
    "\n",
    "print(\"Maximal precision:\", np.max(precision))\n",
    "print(\"Maximal recall:\", np.max(recall))\n",
    "print(\"Maximal accuracy:\", np.max(accuracy))\n",
    "\n",
    "print(\"Minimal precision:\", np.min(precision))\n",
    "print(\"Minimal recall:\", np.min(recall))\n",
    "print(\"Minimal accuracy:\", np.min(accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the first class, plot the distribution of the predictions for the positive and negative cases and the chosen threshold\n",
    "\n",
    "n = 42\n",
    "\n",
    "plt.hist(preds[real[:, n] == 1, 0], bins=100, alpha=0.5, label='Positive', density=True)\n",
    "plt.hist(preds[real[:, n] == 0, 0], bins=100, alpha=0.5, label='Negative', density=True)\n",
    "plt.axvline(x=optimal_thresholds[0], color='r', linestyle='dashed', linewidth=2)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "start",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}