train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Hyperonym Barba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains instructions on how to train a Hyperonym Barba model with public and private NLI datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_size = \"large\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install TensorFlow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.11.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install Hugging Face libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!pip install transformers==4.28.1 datasets==2.11.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset, concatenate_datasets, Features, Value, ClassLabel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set number of processes to use for parallel operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_proc = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A typical NLI model generally has three output labels, namely `entailment`, `neutral` and `contradiction`.\n",
    "\n",
    "To support various datasets, Barba uses only two labels, `entailment` and `not_entailment`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = Features(\n",
    "    {\n",
    "        \"hypothesis\": Value(dtype=\"string\"),\n",
    "        \"premise\": Value(dtype=\"string\"),\n",
    "        \"label\": ClassLabel(names=[\"entailment\", \"not_entailment\"]),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function for removing redundant columns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def strip_columns(dataset):\n",
    "    columns = dataset[list(dataset)[0]].column_names\n",
    "    columns = [col for col in columns if col not in features]\n",
    "    return dataset.remove_columns(columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function for squashing `neutral` and `contradiction` into a single label:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def squash_labels(dataset):\n",
    "    def fn(example):\n",
    "        if example[\"label\"] == 2:\n",
    "            example[\"label\"] = 1\n",
    "        return example\n",
    "    return dataset.map(fn, features=features, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All-in-one function for normalizing datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(dataset):\n",
    "    dataset = strip_columns(dataset)\n",
    "    dataset = squash_labels(dataset)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load public datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### SNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The SNLI corpus (version 1.0) is a collection of 570k human-written English sentence pairs manually labeled for balanced classification with the labels entailment, contradiction, and neutral, supporting the task of natural language inference (NLI), also known as recognizing textual entailment (RTE)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snli = load_dataset(\"snli\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snli = normalize(snli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Multi-Genre Natural Language Inference Corpus is a crowdsourced collection of sentence pairs with textual entailment annotations. Given a premise sentence and a hypothesis sentence, the task is to predict whether the premise entails the hypothesis (entailment), contradicts the hypothesis (contradiction), or neither (neutral). The premise sentences are gathered from ten different sources, including transcribed speech, fiction, and government reports. The authors of the benchmark use the standard test set, for which they obtained private labels from the RTE authors, and evaluate on both the matched (in-domain) and mismatched (cross-domain) section. They also uses and recommend the SNLI corpus as 550k examples of auxiliary training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli = load_dataset(\"glue\", \"mnli\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli = normalize(mnli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MRPC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Microsoft Research Paraphrase Corpus (Dolan & Brockett, 2005) is a corpus of sentence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mrpc = load_dataset(\"glue\", \"mrpc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mrpc = mrpc.rename_column(\"sentence1\", \"premise\")\n",
    "mrpc = mrpc.rename_column(\"sentence2\", \"hypothesis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mrpc = normalize(mrpc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### RTE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Recognizing Textual Entailment (RTE) datasets come from a series of annual textual entailment challenges. The authors of the benchmark combined the data from RTE1 (Dagan et al., 2006), RTE2 (Bar Haim et al., 2006), RTE3 (Giampiccolo et al., 2007), and RTE5 (Bentivogli et al., 2009). Examples are constructed based on news and Wikipedia text. The authors of the benchmark convert all datasets to a two-class split, where for three-class datasets they collapse neutral and contradiction into not entailment, for consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rte = load_dataset(\"glue\", \"rte\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rte = rte.rename_column(\"sentence1\", \"premise\")\n",
    "rte = rte.rename_column(\"sentence2\", \"hypothesis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rte = normalize(rte)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### WNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Winograd Schema Challenge (Levesque et al., 2011) is a reading comprehension task in which a system must read a sentence with a pronoun and select the referent of that pronoun from a list of choices. The examples are manually constructed to foil simple statistical methods: Each one is contingent on contextual information provided by a single word or phrase in the sentence. To convert the problem into sentence pair classification, the authors of the benchmark construct sentence pairs by replacing the ambiguous pronoun with each possible referent. The task is to predict if the sentence with the pronoun substituted is entailed by the original sentence. They use a small evaluation set consisting of new examples derived from fiction books that was shared privately by the authors of the original corpus. While the included training set is balanced between two classes, the test set is imbalanced between them (65% not entailment). Also, due to a data quirk, the development set is adversarial: hypotheses are sometimes shared between training and development examples, so if a model memorizes the training examples, they will predict the wrong label on corresponding development set example. As with QNLI, each example is evaluated separately, so there is not a systematic correspondence between a model's score on this task and its score on the unconverted original task. The authors of the benchmark call converted dataset WNLI (Winograd NLI)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wnli = load_dataset(\"glue\", \"wnli\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WNLI uses 0 as `not_entailment` and 1 as `entailment`, so we need to adjust the labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wnli_fix(example):\n",
    "    if example[\"label\"] == 1:\n",
    "        example[\"label\"] = 0\n",
    "    elif example[\"label\"] == 0:\n",
    "        example[\"label\"] = 1\n",
    "    return example\n",
    "\n",
    "\n",
    "wnli = wnli.map(wnli_fix, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wnli = wnli.rename_column(\"sentence1\", \"premise\")\n",
    "wnli = wnli.rename_column(\"sentence2\", \"hypothesis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wnli = normalize(wnli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### OCNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OCNLI stands for Original Chinese Natural Language Inference. It is corpus for Chinese Natural Language Inference, collected following closely the procedures of MNLI, but with enhanced strategies aiming for more challenging inference pairs. We want to emphasize we did not use human/machine translation in creating the dataset, and thus our Chinese texts are original and not translated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocnli = load_dataset(\"clue\", \"ocnli\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OCNLI uses 0 as `neutral` and 1 as `entailment`, so we need to adjust the labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ocnli_fix(example):\n",
    "    if example[\"label\"] == 1:\n",
    "        example[\"label\"] = 0\n",
    "    elif example[\"label\"] == 0:\n",
    "        example[\"label\"] = 1\n",
    "    return example\n",
    "\n",
    "\n",
    "ocnli = ocnli.map(ocnli_fix, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocnli = ocnli.rename_column(\"sentence1\", \"premise\")\n",
    "ocnli = ocnli.rename_column(\"sentence2\", \"hypothesis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ocnli = normalize(ocnli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### JNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "JNLI is a Japanese version of the NLI (Natural Language Inference) dataset. NLI is a task to recognize the inference relation that a premise sentence has to a hypothesis sentence. The inference relations are entailment, contradiction, and neutral."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "jnli = load_dataset(\"shunk031/JGLUE\", \"JNLI\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jnli = jnli.rename_column(\"sentence1\", \"premise\")\n",
    "jnli = jnli.rename_column(\"sentence2\", \"hypothesis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jnli = normalize(jnli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### JaNLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The JaNLI (Japanese Adversarial NLI) dataset, inspired by the English HANS dataset, is designed to necessitate an understanding of Japanese linguistic phenomena and to illuminate the vulnerabilities of models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "janli = load_dataset(\"hpprc/janli\", \"base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "janli = normalize(janli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### KLUE-NLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "KLUE is a collection of 8 tasks to evaluate natural language understanding capability of Korean language models. We delibrately select the 8 tasks, which are Topic Classification, Semantic Textual Similarity, Natural Language Inference, Named Entity Recognition, Relation Extraction, Dependency Parsing, Machine Reading Comprehension, and Dialogue State Tracking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "klue_nli = load_dataset(\"klue\", \"nli\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "klue_nli = normalize(klue_nli)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Group public datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public_train_datasets = [\n",
    "    snli[\"train\"],\n",
    "    mnli[\"train\"],\n",
    "    mrpc[\"train\"],\n",
    "    rte[\"train\"],\n",
    "    wnli[\"train\"],\n",
    "    ocnli[\"train\"],\n",
    "    jnli[\"train\"],\n",
    "    janli[\"train\"],\n",
    "    klue_nli[\"train\"],\n",
    "]\n",
    "public_validation_datasets = [\n",
    "    snli[\"validation\"],\n",
    "    mnli[\"validation_matched\"],\n",
    "    mrpc[\"validation\"],\n",
    "    rte[\"validation\"],\n",
    "    wnli[\"validation\"],\n",
    "    ocnli[\"validation\"],\n",
    "    jnli[\"validation\"],\n",
    "    klue_nli[\"validation\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load private datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try to load private datasets in the `datasets` directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "private_train_datasets = []\n",
    "private_validation_datasets = []\n",
    "if os.path.isdir(\"datasets\"):\n",
    "    try:\n",
    "        private_dataset = load_dataset(\"./datasets\")\n",
    "        private_dataset = normalize(private_dataset)\n",
    "        if \"train\" in private_dataset:\n",
    "            private_train_datasets.append(private_dataset[\"train\"])\n",
    "        if \"validation\" in private_dataset:\n",
    "            private_validation_datasets.append(private_dataset[\"validation\"])\n",
    "    except FileNotFoundError:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(private_train_datasets)\n",
    "print(private_validation_datasets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Concatenate datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = concatenate_datasets(public_train_datasets + private_train_datasets)\n",
    "validation_dataset = concatenate_datasets(public_validation_datasets + private_validation_datasets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_label(dataset):\n",
    "    def fn(example):\n",
    "        if example[\"label\"] < 0 or example[\"label\"] > 1:\n",
    "            return False\n",
    "        if len(example[\"hypothesis\"]) == 0:\n",
    "            return False\n",
    "        if len(example[\"premise\"]) == 0:\n",
    "            return False\n",
    "        return True\n",
    "    return dataset.filter(fn, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = filter_label(train_dataset)\n",
    "validation_dataset = filter_label(validation_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyze datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute class weight:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class_weight = {0: 0.0, 1: 0.0}\n",
    "for example in train_dataset:\n",
    "    class_weight[example[\"label\"]] += 1.0\n",
    "for i in range(2):\n",
    "    class_weight[i] = (1 / class_weight[i]) * len(train_dataset) * 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(class_weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenize datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load pretrained tokenizer for XLM-RoBERTa:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(f\"xlm-roberta-{model_size}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test tokenization using examples from the [original implementation](https://github.com/facebookresearch/XLM#ii-cross-lingual-language-model-pretraining-xlm):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer(\"Hello world!\"))  # [0, 35378,  8999, 38, 2]\n",
    "print(tokenizer(\"你好，世界\"))  # [0, 6, 124084, 4, 3221, 2]\n",
    "print(tokenizer(\"a\", \"b\", padding=\"max_length\"))  # [0, 10, 2, 2, 876, 2, 1, 1, 1, ..., 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(dataset):\n",
    "    def fn(examples):\n",
    "        return tokenizer(examples[\"hypothesis\"], examples[\"premise\"], truncation=\"do_not_truncate\")\n",
    "    return dataset.map(fn, batched=True, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset = tokenize(train_dataset)\n",
    "validation_dataset = tokenize(validation_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_input(dataset):\n",
    "    def fn(example):\n",
    "        return len(example[\"input_ids\"]) <= tokenizer.model_max_length\n",
    "    return dataset.filter(fn, num_proc=num_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = filter_input(train_dataset)\n",
    "validation_dataset = filter_input(validation_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set hyperparameters based on [XNLI tasks for XLM-RoBERTa](https://github.com/facebookresearch/fairseq/issues/1367#issuecomment-555609917):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 5e-6\n",
    "batch_size = 16 * strategy.num_replicas_in_sync\n",
    "num_epochs = 30\n",
    "patience = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TFAutoModelForSequenceClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    model = TFAutoModelForSequenceClassification.from_pretrained(f\"xlm-roberta-{model_size}\", num_labels=2)\n",
    "    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate), metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert datasets into TensorFlow format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_train_dataset = model.prepare_tf_dataset(train_dataset, shuffle=True, batch_size=batch_size, tokenizer=tokenizer)\n",
    "tf_validation_dataset = model.prepare_tf_dataset(validation_dataset, shuffle=False, batch_size=batch_size, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create callback for early stopping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "callback = tf.keras.callbacks.EarlyStopping(monitor=\"val_accuracy\", patience=patience, restore_best_weights=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fine-tune the pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    tf_train_dataset,\n",
    "    validation_data=tf_validation_dataset,\n",
    "    epochs=num_epochs,\n",
    "    class_weight=class_weight,\n",
    "    callbacks=[callback],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%mkdir -p models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(f\"models/barba-{model_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_pretrained(f\"models/barba-{model_size}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}