{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd96b784",
   "metadata": {},
   "source": [
    "# EuroBERT training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95edf3f9-59b6-4af3-8c14-2463256d4ec1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f5cc51",
   "metadata": {},
   "source": [
    "## Load dataset from huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b536d72",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\n",
    "    \"json\",\n",
    "    data_files=\"/home/apapagiannis/EuroVoc/files/*\",\n",
    "    split=\"train\"\n",
    ")\n",
    "\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54d215a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "split_dataset = dataset.train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "train_dataset = split_dataset[\"train\"]\n",
    "eval_dataset = split_dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eefec7f8",
   "metadata": {},
   "source": [
    "## Tokenize text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2236486",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForMaskedLM\n",
    "\n",
    "BERT_MODEL_NAME = \"nlpaueb/legal-bert-base-uncased\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)\n",
    "model = AutoModelForMaskedLM.from_pretrained(BERT_MODEL_NAME)\n",
    "\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.cls_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825d7a3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def tokenize_function(batch):\n",
    "    texts = []\n",
    "\n",
    "    for t in batch[\"text\"]:\n",
    "        if t is None:\n",
    "            texts.append(\"\")\n",
    "        else:\n",
    "            texts.append(str(t))\n",
    "\n",
    "    return tokenizer(\n",
    "        texts,\n",
    "        truncation=True,\n",
    "        max_length=512\n",
    "    )\n",
    "\n",
    "tokenized_train = train_dataset.map(tokenize_function, batched=True, remove_columns=train_dataset.column_names)\n",
    "tokenized_eval = eval_dataset.map(tokenize_function, batched=True, remove_columns=eval_dataset.column_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fc6309a",
   "metadata": {},
   "source": [
    "## Load model and set up training arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccd9ea70",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForMaskedLM\n",
    "\n",
    "model = AutoModelForMaskedLM.from_pretrained(BERT_MODEL_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce5b33a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer,\n",
    "    mlm=True,\n",
    "    mlm_probability=0.15\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "645f2ac6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define LoRA configuration and apply it to the model\n",
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.FEATURE_EXTRACTION,\n",
    "    r=8,\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=[\"query\", \"value\"]  # BERT attention layers\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, lora_config)\n",
    "\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dd7dfbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, EarlyStoppingCallback\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=8,   # INCREASE this with LoRA\n",
    "    per_device_eval_batch_size=8,\n",
    "\n",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=500,\n",
    "\n",
    "    save_steps=1000,\n",
    "    logging_steps=100,\n",
    "\n",
    "    learning_rate=2e-4,\n",
    "    warmup_steps=200,\n",
    "\n",
    "    fp16=True,  # or bf16=True if supported\n",
    "    report_to=\"none\",\n",
    "    \n",
    "    save_total_limit=3,  # Keep only last 3 checkpoints\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"eval_loss\",\n",
    "    greater_is_better=False,\n",
    "\n",
    ")\n",
    "\n",
    "# Early stopping callback\n",
    "# Monitors eval loss - stops if no improvement for 3 evaluations (~1500 steps)\n",
    "early_stopping = EarlyStoppingCallback(\n",
    "    early_stopping_patience=3,  # Stop if no improvement for 3 evals\n",
    "    early_stopping_threshold=0.001,  # Minimum improvement threshold\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531b1862",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_train,\n",
    "    eval_dataset=tokenized_eval,\n",
    "    data_collator=data_collator,\n",
    "    callbacks=[early_stopping],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d41560e-66d7-406f-8be4-3ccbe7950879",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()\n",
    "\n",
    "model.save_pretrained(\"./fine_tuned_model\")\n",
    "tokenizer.save_pretrained(\"./fine_tuned_model\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "eubert",
   "language": "python",
   "name": "eurovoc-venv"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}