{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0",
   "metadata": {},
   "source": [
    "# Fine-tune a BERT sentence classifier\n",
    "\n",
    "Install the checkout's `[notebooks]` extra and run `data.ipynb` first.\n",
    "The default is an offline CPU demonstration: a tiny randomly initialised\n",
    "BERT with a local vocabulary, trained for one epoch. It verifies the\n",
    "pipeline and does **not** reproduce SciBERT or its historical scores.\n",
    "\n",
    "Set `GREAT_AI_NOTEBOOK_FULL=1` before starting Jupyter to fine-tune\n",
    "`allenai/scibert_scivocab_uncased` instead. This downloads pretrained weights\n",
    "and takes substantially more memory and time; a GPU is recommended.\n",
    "Run the data and deployment notebooks in the same mode."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2346798f",
   "metadata": {},
   "source": [
    "## Local setup\n",
    "\n",
    "Use Python 3.12 or 3.13 and install this checkout with `uv pip install --editable '.[notebooks]'`.\n",
    "See **Run notebooks locally** in the repository README for the kernel, execution order,\n",
    "and full-data options. Run related notebooks from the same working directory\n",
    "to share the default `.cache/` model store and `tracing_database.json` database.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824837fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from docs.notebook_support import FULL_RUN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cell-3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dependencies are installed once with uv pip install --editable \".[notebooks]\"."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-4",
   "metadata": {},
   "source": [
    "Load the dataset saved locally by `data.ipynb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cell-5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;39mFetching cached versions of summary-train-dataset-small\u001b[0m\n",
      "\u001b[38;5;39mLatest version of summary-train-dataset-small is 0 (from versions: 0)\u001b[0m\n",
      "\u001b[38;5;39mFile summary-train-dataset-small-0 found in cache\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from great_ai.large_file import LargeFileLocal\n",
    "import json\n",
    "\n",
    "with LargeFileLocal(\"summary-train-dataset-small\", encoding=\"utf-8\") as f:\n",
    "    X, y = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-6",
   "metadata": {},
   "source": [
    "Finetune SciBERT, for more info about this step, check out [HuggingFace](https://huggingface.co/docs/transformers/training).\n",
    "If you're only here for `great-ai`, feel free to skip the next cell.\n",
    "\n",
    "We use an 80/20 training/validation split. Early stopping and checkpoint selection both use validation binary F1 for positive label `1`. There is no independent test split in this example. The split and training seed is fixed for future runs, and the selected checkpoint's evaluation and run metadata are saved under `models/`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W1028 12:48:58.270000 53390 torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c74823d6364447a96609229fbbc5512",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/60 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': '0.6939', 'grad_norm': '0.5046', 'learning_rate': '1.667e-05', 'epoch': '1'}\n",
      "{'eval_loss': '0.6939', 'eval_f1': '0', 'eval_runtime': '0.0123', 'eval_samples_per_second': '974.9', 'eval_steps_per_second': '81.24', 'epoch': '1'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "23492278f66e46f09e9eee0680c21b92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[transformers] early stopping required metric_for_best_model, but did not find eval_f1 so early stopping is disabled\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': '0.0882', 'train_samples_per_second': '544.2', 'train_steps_per_second': '34.01', 'train_loss': '0.6939', 'epoch': '1'}\n",
      "{'validation_loss': '0.6939', 'validation_f1': '0', 'validation_runtime': '0.0025', 'validation_samples_per_second': '4804', 'validation_steps_per_second': '400.3', 'epoch': '1'}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'validation_loss': 0.693895161151886,\n",
       " 'validation_f1': 0.0,\n",
       " 'validation_runtime': 0.0025,\n",
       " 'validation_samples_per_second': 4804.013,\n",
       " 'validation_steps_per_second': 400.334,\n",
       " 'epoch': 1.0}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import (\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    EarlyStoppingCallback,\n",
    "    set_seed,\n",
    ")\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from datasets import Dataset\n",
    "from sklearn.metrics import f1_score\n",
    "import torch\n",
    "\n",
    "MODEL = \"allenai/scibert_scivocab_uncased\"\n",
    "BATCH_SIZE = 8 if FULL_RUN else 16\n",
    "SPLIT_SEED = 42\n",
    "\n",
    "set_seed(SPLIT_SEED)\n",
    "\n",
    "if FULL_RUN:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)\n",
    "else:\n",
    "    import re\n",
    "    from transformers import BertConfig, BertForSequenceClassification, BertTokenizer\n",
    "\n",
    "    torch.set_num_threads(2)\n",
    "    vocabulary = [\"[PAD]\", \"[UNK]\", \"[CLS]\", \"[SEP]\", \"[MASK]\"]\n",
    "    vocabulary += sorted(set(re.findall(r\"[a-z]+|[0-9]+|[^\\w\\s]\", \" \".join(X).lower())))\n",
    "    vocab_path = Path(\"vocab.txt\")\n",
    "    vocab_path.write_text(\"\\n\".join(vocabulary) + \"\\n\", encoding=\"utf-8\")\n",
    "    tokenizer = BertTokenizer(vocab=str(vocab_path), do_lower_case=True)\n",
    "    model = BertForSequenceClassification(\n",
    "        BertConfig(\n",
    "            vocab_size=len(tokenizer),\n",
    "            hidden_size=32,\n",
    "            num_hidden_layers=1,\n",
    "            num_attention_heads=2,\n",
    "            intermediate_size=64,\n",
    "            max_position_embeddings=512,\n",
    "            num_labels=2,\n",
    "        )\n",
    "    )\n",
    "\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "dataset = (\n",
    "    Dataset.from_dict({\"text\": X, \"label\": y})\n",
    "    .map(lambda v: tokenizer(v[\"text\"], truncation=True, max_length=128), batched=True)\n",
    "    .remove_columns(\"text\")\n",
    "    .train_test_split(test_size=0.2, shuffle=True, seed=SPLIT_SEED)\n",
    ")\n",
    "\n",
    "# The split API calls the held-out partition \"test\"; we use it for validation.\n",
    "dataset[\"validation\"] = dataset.pop(\"test\")\n",
    "\n",
    "\n",
    "def compute_metrics(p):\n",
    "    predictions = (\n",
    "        p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "    )\n",
    "    return {\n",
    "        \"f1\": f1_score(\n",
    "            p.label_ids,\n",
    "            np.argmax(predictions, axis=1),\n",
    "            average=\"binary\",\n",
    "            pos_label=1,\n",
    "            zero_division=0,\n",
    "        )\n",
    "    }\n",
    "\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"models\",\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    per_device_eval_batch_size=BATCH_SIZE,\n",
    "    save_total_limit=5,\n",
    "    num_train_epochs=50 if FULL_RUN else 1,\n",
    "    save_strategy=\"epoch\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    weight_decay=0.01,\n",
    "    metric_for_best_model=\"f1\",\n",
    "    load_best_model_at_end=True,\n",
    "    seed=SPLIT_SEED,\n",
    "    use_cpu=not FULL_RUN,\n",
    "    report_to=\"none\",\n",
    "    dataloader_pin_memory=FULL_RUN,\n",
    "    disable_tqdm=not FULL_RUN,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    processing_class=tokenizer,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    eval_dataset=dataset[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],\n",
    ")\n",
    "result = trainer.train()\n",
    "\n",
    "# This evaluates the selected checkpoint on validation, not an untouched test set.\n",
    "validation_metrics = trainer.evaluate(metric_key_prefix=\"validation\")\n",
    "trainer.save_metrics(\"validation\", validation_metrics)\n",
    "Path(\"models/validation-run.json\").write_text(\n",
    "    json.dumps(\n",
    "        {\n",
    "            \"mode\": \"scibert\" if FULL_RUN else \"tiny-random-bert-demo\",\n",
    "            \"split_seed\": SPLIT_SEED,\n",
    "            \"training_seed\": SPLIT_SEED,\n",
    "            \"validation_fraction\": 0.2,\n",
    "            \"train_samples\": len(dataset[\"train\"]),\n",
    "            \"validation_samples\": len(dataset[\"validation\"]),\n",
    "            \"metric\": \"binary_f1\",\n",
    "            \"positive_label\": 1,\n",
    "            \"best_checkpoint\": trainer.state.best_model_checkpoint,\n",
    "            \"best_validation_f1\": trainer.state.best_metric,\n",
    "        },\n",
    "        indent=2,\n",
    "    )\n",
    "    + \"\\n\",\n",
    "    encoding=\"utf-8\",\n",
    ")\n",
    "validation_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-8",
   "metadata": {},
   "source": [
    "The [saved historical run](https://great-ai.schmelczer.dev/examples/scibert/validation-results.json) selected a checkpoint with **validation binary F1 of 0.8966** (recorded value: 0.896551724137931). This validation split was also used for early stopping and checkpoint selection, so the number is not an independent test result or macro F1. The historical split seed was not recorded; the seeded code above defines future runs and does not reproduce that split. Its results are saved to `models/validation_results.json` and `models/validation-run.json`.\n",
    "\n",
    "To estimate performance after model selection, evaluate once on a separate untouched test split. We make no test-set performance claim here. Next, store the selected model for deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "634b947ab9ee4f20bef78f780d5441ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n",
      "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n",
      "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n",
      "\u001b[38;5;39mModel scibert-highlights uploaded with version 0\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'scibert-highlights:0'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from great_ai import save_model\n",
    "\n",
    "# Trainer has loaded the selected validation checkpoint. Save its tokenizer\n",
    "# too, so deployment never needs to download a different vocabulary.\n",
    "trainer.save_model(\"pretrained\")\n",
    "tokenizer.save_pretrained(\"pretrained\")\n",
    "save_model(Path(\"pretrained\"), key=\"scibert-highlights\", keep_last_n=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-10",
   "metadata": {},
   "source": [
    "Next: [Part 3](https://great-ai.schmelczer.dev/examples/scibert/deploy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.13.5)",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}