src/2_train_models/train_models.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the 3 models\n",
    "\n",
    "## 0. Load the annotated dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import transformers\n",
    "from datasets import Dataset, DatasetDict\n",
    "from sklearn.model_selection import train_test_split\n",
    "import evaluate\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dict = DatasetDict.load_from_disk(\"./data/annotation_generated_from_xlsx/annotation.dataset\")\n",
    "dataset_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_augmented_dict = DatasetDict.load_from_disk(\"./data/annotation_generated_from_xlsx/annotation_chatgpt_augmented.dataset\")\n",
    "\n",
    "dataset_augmented_dict"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Fine-tune Bert-like model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_average_method = \"micro\"\n",
    "nb_epoch = 10\n",
    "label_correspondance = {\n",
    "    0: \"0\",\n",
    "    1: \"covariate\",\n",
    "}\n",
    "metric = evaluate.load(\"seqeval\")\n",
    "\n",
    "def label_id(desired_value):\n",
    "    # reverse Correspondance between label value and their index\n",
    "    for key, value in label_correspondance.items():\n",
    "        if value == desired_value:\n",
    "            return key\n",
    "\n",
    "def align_labels_with_tokens(labels, word_ids):\n",
    "    new_labels = []\n",
    "    current_word = None\n",
    "    for word_id in word_ids:\n",
    "        if word_id != current_word:\n",
    "            # Start of a new word!\n",
    "            current_word = word_id\n",
    "            label = -100 if word_id is None else labels[word_id]\n",
    "            new_labels.append(label)\n",
    "        elif word_id is None:\n",
    "            # Special token\n",
    "            new_labels.append(-100)\n",
    "        else:\n",
    "            # Same word as previous token\n",
    "            label = labels[word_id]\n",
    "            # If the label is B-XXX we change it to I-XXX\n",
    "            # if label % 2 == 1:\n",
    "                # label += 1\n",
    "            new_labels.append(label)\n",
    "    return new_labels\n",
    "\n",
    "def compute_metrics(eval_preds):\n",
    "    logits, labels = eval_preds\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "\n",
    "    # Remove ignored index (special tokens) and convert to labels\n",
    "    true_labels = [[label_correspondance[l] for l in label if l != -100] for label in labels]\n",
    "    true_predictions = [\n",
    "        [label_correspondance[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)\n",
    "    ]\n",
    "    all_metrics = metric.compute(predictions=true_predictions, references=true_labels)\n",
    "    return {\n",
    "        \"precision\": all_metrics[\"overall_precision\"],\n",
    "        \"recall\": all_metrics[\"overall_recall\"],\n",
    "        \"f1\": all_metrics[\"overall_f1\"],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from mlflow import pytorch\n",
    "\n",
    "def train_bert_like_model(pretrained_model, dataset_dict, dataset_name):\n",
    "\n",
    "    from transformers import AutoTokenizer, RobertaTokenizerFast\n",
    "    try:\n",
    "        tokenizer = RobertaTokenizerFast.from_pretrained(pretrained_model, \n",
    "                                                        add_prefix_space=True,\n",
    "                                                        truncating = True,\n",
    "                                                        model_max_length=512\n",
    "                                                        )\n",
    "    except:\n",
    "        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, \n",
    "                                                        add_prefix_space=True,\n",
    "                                                        truncating = True,\n",
    "                                                        model_max_length=512\n",
    "                                                        )\n",
    "\n",
    "    def tokenize_and_align_labels(examples):\n",
    "        tokenized_inputs = tokenizer(\n",
    "            examples[\"tokens\"], truncation=True, is_split_into_words=True,\n",
    "        )\n",
    "        all_labels = examples[\"ner_tags\"]\n",
    "        new_labels = []\n",
    "        for i, labels in enumerate(all_labels):\n",
    "            # word_ids = tokenized_inputs.word_ids(i)\n",
    "            # print(f\"i: {i} | labels: {labels}\")\n",
    "            word_ids = tokenized_inputs.word_ids(i)\n",
    "            new_labels.append(align_labels_with_tokens(labels, word_ids))\n",
    "\n",
    "        tokenized_inputs[\"labels\"] = new_labels\n",
    "        return tokenized_inputs\n",
    "\n",
    "    tokenized_datasets = dataset_dict.map(\n",
    "        tokenize_and_align_labels,\n",
    "        batched=True,\n",
    "        remove_columns=dataset_dict[\"train\"].column_names,\n",
    "    )\n",
    "    # print(tokenized_datasets)\n",
    "\n",
    "    metric = evaluate.load(\"seqeval\")\n",
    "    id2label = label_correspondance\n",
    "    label2id = {v: k for k, v in id2label.items()}\n",
    "\n",
    "    from transformers import AutoModelForTokenClassification\n",
    "\n",
    "    model = AutoModelForTokenClassification.from_pretrained(\n",
    "        pretrained_model,\n",
    "        id2label=id2label,\n",
    "        label2id=label2id,\n",
    "        \n",
    "    )\n",
    "\n",
    "    from transformers import TrainingArguments\n",
    "\n",
    "    args = TrainingArguments(\n",
    "        f\"mood_covariate_from_{pretrained_model}_{dataset_name}\",\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        save_strategy=\"no\",\n",
    "        learning_rate=2e-5,\n",
    "        num_train_epochs=nb_epoch,\n",
    "        # output_dir=\"./models\"\n",
    "        # load_best_model_at_end=True,\n",
    "        # save_strategy=\"epoch\",mood_covariate_from_{pretrained_model}\n",
    "        #weight_decay=0.01,\n",
    "    )\n",
    "\n",
    "    from transformers import DataCollatorForTokenClassification\n",
    "\n",
    "    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)\n",
    "\n",
    "    from transformers import Trainer\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=args,\n",
    "        train_dataset=tokenized_datasets[\"train\"],\n",
    "        eval_dataset=tokenized_datasets[\"test\"],\n",
    "        data_collator=data_collator,\n",
    "        compute_metrics=compute_metrics,\n",
    "        tokenizer=tokenizer,\n",
    "    )\n",
    "    with mlflow.start_run(run_name=pretrained_model) as run:\n",
    "        trainer.train()\n",
    "\n",
    "    trainer.save_model(f\"./models/mood_covariate_from_{pretrained_model}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.set_experiment(\"Bert-Like models\")\n",
    "\n",
    "\n",
    "pretrained_model = \"roberta-base\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"base\")\n",
    "\n",
    "pretrained_model = \"microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"base\")\n",
    "\n",
    "pretrained_model = \"FacebookAI/xlm-roberta-base\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.set_experiment(\"Hybride\")\n",
    "\n",
    "\n",
    "pretrained_model = \"roberta-base\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"gpt3-5-augmented\")\n",
    "\n",
    "pretrained_model = \"microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"gpt3-5-augmented\")\n",
    "\n",
    "pretrained_model = \"FacebookAI/xlm-roberta-base\"\n",
    "train_bert_like_model(pretrained_model, dataset_dict, \"gpt3-5-augmented\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSequenceClassification\n",
    "from transformers import pipeline\n",
    "\n",
    "model_path = f\"./models/mood_covariate_from_{pretrained_model}\"\n",
    "\n",
    "# Load the model using the pipeline for Named Entity Recognition (NER)\n",
    "ner_classifier = pipeline(\"ner\", model=model_path, aggregation_strategy=\"simple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = ner_classifier(\" in Figs. S3 â€“S6. Absolute humidity was found to be signiï¬cantly linked to epidemic onset dates at the spatial scale ( p= 0.029), but not at the other scales. The associated coefï¬cient was negative ( -0.4763). Mobility ï¬‚ows were not found to be signiï¬cantly linked to epidemic onset dates (p= 0.57 with the corrected model, p= 0.73 with the uncorrected model). In the corrected model, the coefï¬cient \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for r in res:\n",
    "    print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mood_luca_review_avian_flu",
   "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.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}