training notebooks

notebooks/train_fasttext.ipynb

@@ -0,0 +1,161 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f83ef6b7",
+   "source": "# Train: fastText Baseline\n\nTrain a fastText supervised classifier on entity-sentiment marker-mode text.\n\n**Structure**\n1. Load & Prepare Data\n2. Training\n3. Evaluation\n4. Training & Evaluation — Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "1225a248",
+   "source": "import os\nimport sys\nfrom collections import Counter\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pandas as pd\nfrom sklearn.metrics import f1_score, classification_report\n\nsys.path.insert(0, os.path.abspath(\"..\"))\n\nfrom src.models.fasttext import (\n    prepare_data,\n    train,\n    evaluate,\n    evaluate_entity_level,\n    predict_samples,\n    MODE,\n    LABEL_PREFIX,\n    _write_fasttext_file,\n)\nfrom src.models.dataset import deduplicate_positions, flatten_to_examples, load_data, split_data\nfrom src.schemas.labels import SENTIMENT_LABELS",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2d3f288",
+   "source": "## 1. Load & Prepare Data",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "d7e4bf2f",
+   "source": "DATA_PATH = \"../data/data_augmented_256.jsonl\"\nOUTPUT_DIR = \"../models/fasttext\"\nSEED = 42\nVAL_SPLIT = 0.1\nTEST_SPLIT = 0.1\n\nLR = 0.5\nEPOCH = 25\nWORD_NGRAMS = 2\nDIM = 100",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "ec7421a5",
+   "source": "samples = load_data(DATA_PATH)\nprint(f\"Loaded {len(samples)} samples\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "f405ae7f",
+   "source": "examples = flatten_to_examples(samples, mode=MODE)\n\nn_ents = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in examples})\nprint(f\"{len(samples)} samples -> {n_ents} entities -> {len(examples)} examples\")\n\ndist = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in examples)\nprint(f\"Label distribution: {dict(dist)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "f1cff37c",
+   "source": "train_exs, val_exs, test_exs = split_data(\n    examples, VAL_SPLIT, TEST_SPLIT, seed=SEED\n)\nprint(f\"Train: {len(train_exs)} | Val: {len(val_exs)} | Test: {len(test_exs)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "381fe572",
+   "source": "## 2. Training",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "cb664641",
+   "source": "model = train(\n    train_exs,\n    val_exs,\n    output_dir=OUTPUT_DIR,\n    lr=LR,\n    epoch=EPOCH,\n    word_ngrams=WORD_NGRAMS,\n    dim=DIM,\n)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5fc12804",
+   "source": "## 3. Evaluation\n\n### Per-position evaluation",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "281f15be",
+   "source": "macro_f1 = evaluate(model, test_exs, split_name=\"test\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "776cc993",
+   "source": "### Entity-level evaluation (highest-confidence aggregation)\n\nWhen an entity has multiple positions, each produces a separate prediction.\nHere we aggregate by keeping the prediction with the highest confidence per entity,\nmatching the production inference behavior.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "a3658e91",
+   "source": "macro_f1_entity = evaluate_entity_level(model, test_exs, split_name=\"test\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "53f6d155",
+   "source": "## 4. Training & Evaluation — Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "67d96471",
+   "source": "deduped_samples = deduplicate_positions(samples)\ndedup_examples = flatten_to_examples(deduped_samples, mode=MODE)\n\nn_ents_dedup = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in dedup_examples})\nprint(f\"{len(deduped_samples)} samples -> {n_ents_dedup} entities -> {len(dedup_examples)} examples (deduplicated)\")\n\ndist_dedup = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in dedup_examples)\nprint(f\"Label distribution: {dict(dist_dedup)}\")\n\ntrain_dedup, val_dedup, test_dedup = split_data(\n    dedup_examples, VAL_SPLIT, TEST_SPLIT, seed=SEED\n)\nprint(f\"Train: {len(train_dedup)} | Val: {len(val_dedup)} | Test: {len(test_dedup)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "44bd734f",
+   "source": "model_dedup = train(\n    train_dedup,\n    val_dedup,\n    output_dir=os.path.join(OUTPUT_DIR, \"dedup\"),\n    lr=LR,\n    epoch=EPOCH,\n    word_ngrams=WORD_NGRAMS,\n    dim=DIM,\n)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "471dd000",
+   "source": "### Dedup test evaluation",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "b0b5232d",
+   "source": "macro_f1_dedup = evaluate(model_dedup, test_dedup, split_name=\"test (deduplicated)\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b9288376",
+   "source": "### Comparison: All Positions vs Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "60ecd7da",
+   "source": "comparison = pd.DataFrame({\n    \"variant\": [\n        \"all positions (per-position)\",\n        \"all positions (entity-level)\",\n        \"deduplicated (per-position)\",\n    ],\n    \"train_examples\": [len(train_exs), len(train_exs), len(train_dedup)],\n    \"test_macro_f1\": [\n        round(macro_f1, 4),\n        round(macro_f1_entity, 4),\n        round(macro_f1_dedup, 4),\n    ],\n})\ncomparison",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/train_marker.ipynb

@@ -0,0 +1,417 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a1",
+   "metadata": {},
+   "source": [
+    "# Train: Marker Mode\n",
+    "\n",
+    "Fine-tune DistilBERT with `[E]...[/E]` entity-wrapping (single-sequence input).\n",
+    "\n",
+    "**Structure**\n",
+    "1. Load Config & Data\n",
+    "2. Tokenizer & Model\n",
+    "3. Training\n",
+    "4. Evaluation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "a2",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-18T23:17:44.162106Z",
+     "start_time": "2026-04-18T23:17:41.253237Z"
+    }
+   },
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "import random\n",
+    "from collections import Counter\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "torch.backends.cuda.enable_cudnn_sdp(False)\n",
+    "from sklearn.metrics import classification_report, f1_score\n",
+    "from transformers import (\n",
+    "    AutoModelForSequenceClassification,\n",
+    "    AutoTokenizer,\n",
+    "    EarlyStoppingCallback,\n",
+    ")\n",
+    "\n",
+    "sys.path.insert(0, os.path.abspath(\"..\"))\n",
+    "\n",
+    "from src.schemas.config import TrainingConfig\n",
+    "from src.schemas.labels import MODES, SENTIMENT_LABELS\n",
+    "from src.models.dataset import (\n",
+    "    EntitySentimentDataset,\n",
+    "    deduplicate_positions,\n",
+    "    flatten_to_examples,\n",
+    "    load_data,\n",
+    "    split_data,\n",
+    ")\n",
+    "from src.models.distillbert import (\n",
+    "    WeightedLossTrainer,\n",
+    "    compute_class_weights,\n",
+    "    make_compute_metrics,\n",
+    ")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3",
+   "metadata": {},
+   "source": [
+    "## 1. Load Config & Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "a4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-18T23:17:44.232796Z",
+     "start_time": "2026-04-18T23:17:44.227568Z"
+    }
+   },
+   "source": [
+    "cfg = TrainingConfig.from_json(\"../data/config_marker.json\")\n",
+    "mode_cfg = MODES[cfg.mode]\n",
+    "\n",
+    "random.seed(cfg.seed)\n",
+    "np.random.seed(cfg.seed)\n",
+    "torch.manual_seed(cfg.seed)\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}  |  Mode: {cfg.mode}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Device: cpu  |  Mode: marker\n"
+     ]
+    }
+   ],
+   "execution_count": 2
+  },
+  {
+   "cell_type": "code",
+   "id": "a5",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-18T23:17:52.714033Z",
+     "start_time": "2026-04-18T23:17:52.556959Z"
+    }
+   },
+   "source": "samples = load_data(os.path.join(\"..\", cfg.data_path))\nexamples = flatten_to_examples(samples, mode=cfg.mode)\n\nn_ents = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in examples})\nprint(f\"{len(samples)} samples -> {n_ents} entities -> {len(examples)} examples\")\n\ndist = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in examples)\nprint(f\"Label distribution: {dict(dist)}\")",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "id": "a6",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-18T23:17:57.689676Z",
+     "start_time": "2026-04-18T23:17:57.680885Z"
+    }
+   },
+   "source": [
+    "train_exs, val_exs, test_exs = split_data(\n",
+    "    examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n",
+    ")\n",
+    "print(f\"Train: {len(train_exs)} | Val: {len(val_exs)} | Test: {len(test_exs)}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train: 21313 | Val: 2796 | Test: 2667\n"
+     ]
+    }
+   ],
+   "execution_count": 5
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ac38d1e",
+   "source": "### Deduplicated dataset (one position per entity)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "31e27ba1",
+   "source": "deduped_samples = deduplicate_positions(samples)\ndedup_examples = flatten_to_examples(deduped_samples, mode=cfg.mode)\n\nn_ents_dedup = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in dedup_examples})\nprint(f\"{len(deduped_samples)} samples -> {n_ents_dedup} entities -> {len(dedup_examples)} examples (deduplicated)\")\n\ndist_dedup = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in dedup_examples)\nprint(f\"Label distribution: {dict(dist_dedup)}\")\n\ntrain_dedup, val_dedup, test_dedup = split_data(\n    dedup_examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n)\nprint(f\"Train: {len(train_dedup)} | Val: {len(val_dedup)} | Test: {len(test_dedup)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a7",
+   "metadata": {},
+   "source": [
+    "## 2. Tokenizer & Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n",
+    "tokenizer.add_special_tokens(\n",
+    "    {\"additional_special_tokens\": [mode_cfg.entity_start, mode_cfg.entity_end]}\n",
+    ")\n",
+    "\n",
+    "model = AutoModelForSequenceClassification.from_pretrained(\n",
+    "    cfg.model_name,\n",
+    "    num_labels=mode_cfg.labels.num_labels,\n",
+    "    id2label=mode_cfg.labels.id2label,\n",
+    "    label2id=mode_cfg.labels.label2id,\n",
+    ")\n",
+    "model.resize_token_embeddings(len(tokenizer))\n",
+    "print(f\"Model vocab size: {len(tokenizer)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "33b27fd2",
+   "source": "special_tokens = [\n    t.content for t in tokenizer.added_tokens_decoder.values()\n    if t.special and t.content not in {\n        tokenizer.cls_token, tokenizer.sep_token,\n        tokenizer.pad_token, tokenizer.unk_token, tokenizer.mask_token,\n    }\n]\nprint(\"Special tokens:\", special_tokens)\nprint()\nfor token in special_tokens:\n    token_id = tokenizer.convert_tokens_to_ids(token)\n    print(f\"  {token!r} -> token_id={token_id}\")\n\nsample_text = f\"{mode_cfg.entity_start} Apple {mode_cfg.entity_end} released a new product.\"\nencoded = tokenizer(sample_text, add_special_tokens=False)\ndecoded_tokens = tokenizer.convert_ids_to_tokens(encoded[\"input_ids\"])\nprint(f\"\\nSample: {sample_text!r}\")\nprint(f\"Tokens: {decoded_tokens}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac72d17e",
+   "source": "### Truncation Analysis",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "13f4e56e",
+   "source": "lengths = []\nfor ex in examples:\n    enc = tokenizer(ex[\"seg_a\"], truncation=False)\n    lengths.append(len(enc[\"input_ids\"]))\n\nlengths = np.array(lengths)\nn_truncated = (lengths > cfg.max_len).sum()\nprint(f\"Max length: {cfg.max_len}\")\nprint(f\"Total examples: {len(lengths)}\")\nprint(f\"Truncated: {n_truncated} ({100 * n_truncated / len(lengths):.1f}%)\")\nprint(f\"Token length — min: {lengths.min()}, median: {int(np.median(lengths))}, \"\n      f\"mean: {lengths.mean():.0f}, max: {lengths.max()}\")\n\nfig, ax = plt.subplots(figsize=(10, 3))\nax.hist(lengths, bins=50, edgecolor=\"black\", linewidth=0.5)\nax.axvline(cfg.max_len, color=\"red\", linestyle=\"--\", label=f\"max_len={cfg.max_len}\")\nax.set_xlabel(\"Token count\")\nax.set_ylabel(\"Examples\")\nax.set_title(\"Input Length Distribution (marker)\")\nax.legend()\nplt.tight_layout()\nplt.show()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a9",
+   "metadata": {},
+   "source": [
+    "## 3. Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_ds = EntitySentimentDataset(train_exs, tokenizer, cfg.max_len)\n",
+    "val_ds = EntitySentimentDataset(val_exs, tokenizer, cfg.max_len)\n",
+    "\n",
+    "training_args = cfg.to_training_arguments()\n",
+    "training_args.output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "training_args.logging_dir = os.path.join(\"..\", cfg.output_dir, \"logs\")\n",
+    "\n",
+    "class_weights = compute_class_weights(train_exs, n_classes=mode_cfg.labels.num_labels)\n",
+    "print(f\"Class weights: {class_weights.tolist()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = WeightedLossTrainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_ds,\n",
+    "    eval_dataset=val_ds,\n",
+    "    compute_metrics=make_compute_metrics(cfg.mode),\n",
+    "    callbacks=[EarlyStoppingCallback(\n",
+    "        early_stopping_patience=cfg.early_stopping_patience\n",
+    "    )],\n",
+    "    class_weights=class_weights,\n",
+    ")\n",
+    "\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "trainer.save_model(output_dir)\n",
+    "tokenizer.save_pretrained(output_dir)\n",
+    "\n",
+    "with open(os.path.join(output_dir, \"mode.txt\"), \"w\") as f:\n",
+    "    f.write(cfg.mode)\n",
+    "\n",
+    "print(f\"Model saved to '{output_dir}'\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e7ed82a",
+   "source": "### Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "ad91c2dd",
+   "source": "import pandas as pd\n\neval_logs = [l for l in trainer.state.log_history if \"eval_loss\" in l]\n\nrows = []\nfor l in eval_logs:\n    rows.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1\": round(l[\"eval_macro_f1\"], 4),\n        \"f1_negative\": round(l.get(\"eval_f1_negative\", 0), 4),\n        \"f1_neutral\": round(l.get(\"eval_f1_neutral\", 0), 4),\n        \"f1_positive\": round(l.get(\"eval_f1_positive\", 0), 4),\n    })\n\ndf_metrics = pd.DataFrame(rows)\ndf_metrics",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a13",
+   "metadata": {},
+   "source": [
+    "## 4. Evaluation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_ds = EntitySentimentDataset(test_exs, tokenizer, cfg.max_len)\n",
+    "pred_output = trainer.predict(test_ds)\n",
+    "\n",
+    "test_preds = np.argmax(pred_output.predictions, axis=-1).tolist()\n",
+    "test_labels = pred_output.label_ids.tolist()\n",
+    "macro_f1 = f1_score(test_labels, test_preds, average=\"macro\")\n",
+    "\n",
+    "print(f\"Test macro-F1: {macro_f1:.4f}\")\n",
+    "print()\n",
+    "print(classification_report(\n",
+    "    test_labels, test_preds,\n",
+    "    target_names=list(SENTIMENT_LABELS.classes),\n",
+    "))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bf46fe96",
+   "source": "### Entity-level Evaluation (highest-confidence aggregation)\n\nWhen an entity has multiple positions, each produces a separate prediction.\nHere we aggregate by keeping the prediction with the highest confidence per entity,\nmatching the production inference behavior.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "951b95c0",
+   "source": "import torch.nn.functional as F\n\nprobs = F.softmax(torch.tensor(pred_output.predictions), dim=-1).numpy()\nconfidences = probs.max(axis=-1)\npreds_all = np.argmax(probs, axis=-1)\n\nsentiments = list(SENTIMENT_LABELS.classes)\nentity_best: dict[tuple, tuple[int, int, float]] = {}\nfor ex, pred_id, conf, label in zip(test_exs, preds_all, confidences, test_labels):\n    key = (ex[\"sample_id\"], ex[\"entity_id\"])\n    if key not in entity_best or conf > entity_best[key][2]:\n        entity_best[key] = (int(pred_id), int(label), float(conf))\n\nent_preds = [v[0] for v in entity_best.values()]\nent_labels = [v[1] for v in entity_best.values()]\n\nmacro_f1_entity = f1_score(ent_labels, ent_preds, average=\"macro\")\nprint(f\"Entity-level test macro-F1: {macro_f1_entity:.4f}\")\nprint(f\"({len(entity_best)} unique entities from {len(test_exs)} position-level examples)\")\nprint()\nprint(classification_report(\n    ent_labels, ent_preds,\n    target_names=sentiments,\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c5aa45e7",
+   "source": "## 5. Training & Evaluation — Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "2d95ad6b",
+   "source": "model_dedup = AutoModelForSequenceClassification.from_pretrained(\n    cfg.model_name,\n    num_labels=mode_cfg.labels.num_labels,\n    id2label=mode_cfg.labels.id2label,\n    label2id=mode_cfg.labels.label2id,\n)\nmodel_dedup.resize_token_embeddings(len(tokenizer))\n\ntrain_ds_dedup = EntitySentimentDataset(train_dedup, tokenizer, cfg.max_len)\nval_ds_dedup = EntitySentimentDataset(val_dedup, tokenizer, cfg.max_len)\n\ntraining_args_dedup = cfg.to_training_arguments()\ntraining_args_dedup.output_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntraining_args_dedup.logging_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\", \"logs\")\n\nclass_weights_dedup = compute_class_weights(train_dedup, n_classes=mode_cfg.labels.num_labels)\nprint(f\"Class weights (dedup): {class_weights_dedup.tolist()}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "e3283201",
+   "source": "trainer_dedup = WeightedLossTrainer(\n    model=model_dedup,\n    args=training_args_dedup,\n    train_dataset=train_ds_dedup,\n    eval_dataset=val_ds_dedup,\n    compute_metrics=make_compute_metrics(cfg.mode),\n    callbacks=[EarlyStoppingCallback(\n        early_stopping_patience=cfg.early_stopping_patience\n    )],\n    class_weights=class_weights_dedup,\n)\n\ntrainer_dedup.train()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "98b3e53d",
+   "source": "output_dir_dedup = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntrainer_dedup.save_model(output_dir_dedup)\ntokenizer.save_pretrained(output_dir_dedup)\n\nwith open(os.path.join(output_dir_dedup, \"mode.txt\"), \"w\") as f:\n    f.write(cfg.mode)\n\nprint(f\"Dedup model saved to '{output_dir_dedup}'\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c774bc8b",
+   "source": "### Dedup Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "880964ef",
+   "source": "eval_logs_d = [l for l in trainer_dedup.state.log_history if \"eval_loss\" in l]\n\nrows_d = []\nfor l in eval_logs_d:\n    rows_d.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1\": round(l[\"eval_macro_f1\"], 4),\n        \"f1_negative\": round(l.get(\"eval_f1_negative\", 0), 4),\n        \"f1_neutral\": round(l.get(\"eval_f1_neutral\", 0), 4),\n        \"f1_positive\": round(l.get(\"eval_f1_positive\", 0), 4),\n    })\n\ndf_metrics_dedup = pd.DataFrame(rows_d)\ndf_metrics_dedup",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c72ca8f2",
+   "source": "### Dedup Test Evaluation",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "c14d2477",
+   "source": "test_ds_dedup = EntitySentimentDataset(test_dedup, tokenizer, cfg.max_len)\npred_output_dedup = trainer_dedup.predict(test_ds_dedup)\n\ntest_preds_dedup = np.argmax(pred_output_dedup.predictions, axis=-1).tolist()\ntest_labels_dedup = pred_output_dedup.label_ids.tolist()\nmacro_f1_dedup = f1_score(test_labels_dedup, test_preds_dedup, average=\"macro\")\n\nprint(f\"Test macro-F1 (dedup): {macro_f1_dedup:.4f}\")\nprint()\nprint(classification_report(\n    test_labels_dedup, test_preds_dedup,\n    target_names=list(SENTIMENT_LABELS.classes),\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9c10beee",
+   "source": "### Comparison: All Positions vs Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "381b8bde",
+   "source": "comparison = pd.DataFrame({\n    \"variant\": [\"all positions (per-position)\", \"all positions (entity-level)\", \"deduplicated\"],\n    \"train_examples\": [len(train_exs), len(train_exs), len(train_dedup)],\n    \"test_examples\": [len(test_exs), f\"{len(test_exs)} -> {len(entity_best)} entities\", len(test_dedup)],\n    \"test_macro_f1\": [round(macro_f1, 4), round(macro_f1_entity, 4), round(macro_f1_dedup, 4)],\n})\ncomparison",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/train_qa_b.ipynb

@@ -0,0 +1,426 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c1",
+   "metadata": {},
+   "source": [
+    "# Train: QA-B Mode\n",
+    "\n",
+    "Fine-tune DistilBERT with sentence-pair QA-B formulation (Sun et al. 2019).\n",
+    "\n",
+    "- Seg A: entity-centred context window\n",
+    "- Seg B: \"The polarity of the {entity_type} {entity} is {sentiment} .\"\n",
+    "- Label: binary (no=0, yes=1) per hypothesis\n",
+    "- 3-class prediction: argmax over P(yes) across (negative, neutral, positive) triplet\n",
+    "\n",
+    "**Structure**\n",
+    "1. Load Config & Data\n",
+    "2. Tokenizer & Model\n",
+    "3. Training\n",
+    "4. Evaluation (binary + reconstructed 3-class)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "c2",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-19T00:10:30.098089Z",
+     "start_time": "2026-04-19T00:10:27.289423Z"
+    }
+   },
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "import random\n",
+    "from collections import Counter\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "torch.backends.cuda.enable_cudnn_sdp(False)\n",
+    "from sklearn.metrics import classification_report, f1_score\n",
+    "from transformers import (\n",
+    "    AutoModelForSequenceClassification,\n",
+    "    AutoTokenizer,\n",
+    "    EarlyStoppingCallback,\n",
+    ")\n",
+    "\n",
+    "sys.path.insert(0, os.path.abspath(\"..\"))\n",
+    "\n",
+    "from src.schemas.config import TrainingConfig\n",
+    "from src.schemas.labels import MODES, SENTIMENT_LABELS\n",
+    "from src.models.dataset import (\n",
+    "    EntitySentimentDataset,\n",
+    "    deduplicate_positions,\n",
+    "    flatten_to_examples,\n",
+    "    load_data,\n",
+    "    split_data,\n",
+    ")\n",
+    "from src.models.distillbert import (\n",
+    "    WeightedLossTrainer,\n",
+    "    compute_class_weights,\n",
+    "    evaluate_qa_b_test,\n",
+    "    make_compute_metrics,\n",
+    ")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c3",
+   "metadata": {},
+   "source": [
+    "## 1. Load Config & Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "c4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-19T00:10:30.106930Z",
+     "start_time": "2026-04-19T00:10:30.101538Z"
+    }
+   },
+   "source": [
+    "cfg = TrainingConfig.from_json(\"../data/config_qa_b.json\")\n",
+    "mode_cfg = MODES[cfg.mode]\n",
+    "\n",
+    "random.seed(cfg.seed)\n",
+    "np.random.seed(cfg.seed)\n",
+    "torch.manual_seed(cfg.seed)\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}  |  Mode: {cfg.mode}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Device: cpu  |  Mode: qa_b\n"
+     ]
+    }
+   ],
+   "execution_count": 2
+  },
+  {
+   "cell_type": "code",
+   "id": "c5",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-19T00:10:34.294830Z",
+     "start_time": "2026-04-19T00:10:34.114481Z"
+    }
+   },
+   "source": "samples = load_data(os.path.join(\"..\", cfg.data_path))\nexamples = flatten_to_examples(samples, mode=cfg.mode)\n\nn_ents = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in examples})\nprint(f\"{len(samples)} samples -> {n_ents} entities -> {len(examples)} examples\")\n\ndist = Counter(e[\"label\"] for e in examples)\nprint(f\"Binary labels: yes={dist[1]}  no={dist[0]}\")",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1629 samples -> 10550 entities -> 80328 examples\n",
+      "Binary labels: yes=26776  no=53552\n"
+     ]
+    }
+   ],
+   "execution_count": 3
+  },
+  {
+   "cell_type": "code",
+   "id": "c6",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-04-19T00:10:34.443682Z",
+     "start_time": "2026-04-19T00:10:34.432031Z"
+    }
+   },
+   "source": [
+    "train_exs, val_exs, test_exs = split_data(\n",
+    "    examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n",
+    ")\n",
+    "print(f\"Train: {len(train_exs)} | Val: {len(val_exs)} | Test: {len(test_exs)}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train: 63939 | Val: 8388 | Test: 8001\n"
+     ]
+    }
+   ],
+   "execution_count": 4
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fc2c6912",
+   "source": "### Deduplicated dataset (one position per entity)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "08b48902",
+   "source": "deduped_samples = deduplicate_positions(samples)\ndedup_examples = flatten_to_examples(deduped_samples, mode=cfg.mode)\n\nn_ents_dedup = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in dedup_examples})\nprint(f\"{len(deduped_samples)} samples -> {n_ents_dedup} entities -> {len(dedup_examples)} examples (deduplicated)\")\n\ndist_dedup = Counter(e[\"label\"] for e in dedup_examples)\nprint(f\"Binary labels: yes={dist_dedup[1]}  no={dist_dedup[0]}\")\n\ntrain_dedup, val_dedup, test_dedup = split_data(\n    dedup_examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n)\nprint(f\"Train: {len(train_dedup)} | Val: {len(val_dedup)} | Test: {len(test_dedup)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7",
+   "metadata": {},
+   "source": [
+    "## 2. Tokenizer & Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n",
+    "\n",
+    "model = AutoModelForSequenceClassification.from_pretrained(\n",
+    "    cfg.model_name,\n",
+    "    num_labels=mode_cfg.labels.num_labels,\n",
+    "    id2label=mode_cfg.labels.id2label,\n",
+    "    label2id=mode_cfg.labels.label2id,\n",
+    ")\n",
+    "print(f\"Num labels: {mode_cfg.labels.num_labels} (binary)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "30b7d3a0",
+   "source": "### Truncation Analysis",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "5096fa98",
+   "source": "import matplotlib.pyplot as plt\n\nlengths = []\nfor ex in examples:\n    enc = tokenizer(ex[\"seg_a\"], ex[\"seg_b\"], truncation=False)\n    lengths.append(len(enc[\"input_ids\"]))\n\nlengths = np.array(lengths)\nn_truncated = (lengths > cfg.max_len).sum()\nprint(f\"Max length: {cfg.max_len}\")\nprint(f\"Total examples: {len(lengths)}\")\nprint(f\"Truncated: {n_truncated} ({100 * n_truncated / len(lengths):.1f}%)\")\nprint(f\"Token length — min: {lengths.min()}, median: {int(np.median(lengths))}, \"\n      f\"mean: {lengths.mean():.0f}, max: {lengths.max()}\")\n\nfig, ax = plt.subplots(figsize=(10, 3))\nax.hist(lengths, bins=50, edgecolor=\"black\", linewidth=0.5)\nax.axvline(cfg.max_len, color=\"red\", linestyle=\"--\", label=f\"max_len={cfg.max_len}\")\nax.set_xlabel(\"Token count\")\nax.set_ylabel(\"Examples\")\nax.set_title(\"Input Length Distribution (qa_b)\")\nax.legend()\nplt.tight_layout()\nplt.show()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c9",
+   "metadata": {},
+   "source": [
+    "## 3. Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_ds = EntitySentimentDataset(train_exs, tokenizer, cfg.max_len)\n",
+    "val_ds = EntitySentimentDataset(val_exs, tokenizer, cfg.max_len)\n",
+    "\n",
+    "training_args = cfg.to_training_arguments()\n",
+    "training_args.output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "training_args.logging_dir = os.path.join(\"..\", cfg.output_dir, \"logs\")\n",
+    "\n",
+    "class_weights = compute_class_weights(train_exs, n_classes=mode_cfg.labels.num_labels)\n",
+    "print(f\"Class weights: {class_weights.tolist()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = WeightedLossTrainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_ds,\n",
+    "    eval_dataset=val_ds,\n",
+    "    compute_metrics=make_compute_metrics(cfg.mode),\n",
+    "    callbacks=[EarlyStoppingCallback(\n",
+    "        early_stopping_patience=cfg.early_stopping_patience\n",
+    "    )],\n",
+    "    class_weights=class_weights,\n",
+    ")\n",
+    "\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "trainer.save_model(output_dir)\n",
+    "tokenizer.save_pretrained(output_dir)\n",
+    "\n",
+    "with open(os.path.join(output_dir, \"mode.txt\"), \"w\") as f:\n",
+    "    f.write(cfg.mode)\n",
+    "\n",
+    "print(f\"Model saved to '{output_dir}'\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fb8fe847",
+   "source": "### Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "393bc3ea",
+   "source": "import pandas as pd\n\neval_logs = [l for l in trainer.state.log_history if \"eval_loss\" in l]\n\nrows = []\nfor l in eval_logs:\n    rows.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1_3class\": round(l[\"eval_macro_f1\"], 4),\n        \"bin_accuracy\": round(l.get(\"eval_bin_accuracy\", 0), 4),\n        \"bin_f1_yes\": round(l.get(\"eval_bin_f1_yes\", 0), 4),\n    })\n\ndf_metrics = pd.DataFrame(rows)\ndf_metrics",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c13",
+   "metadata": {},
+   "source": [
+    "## 4. Evaluation\n",
+    "\n",
+    "QA-B uses binary classification per hypothesis. To get 3-class predictions,\n",
+    "we group consecutive triplets (negative, neutral, positive) and take the\n",
+    "argmax of P(yes) across the three."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "macro_f1, test_preds, test_labels = evaluate_qa_b_test(\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    test_exs=test_exs,\n",
+    "    max_len=cfg.max_len,\n",
+    "    batch_size=cfg.per_device_eval_batch_size,\n",
+    "    device=device,\n",
+    ")\n",
+    "\n",
+    "print(f\"Test macro-F1 (3-class reconstructed): {macro_f1:.4f}\")\n",
+    "print()\n",
+    "print(classification_report(\n",
+    "    test_labels, test_preds,\n",
+    "    target_names=list(SENTIMENT_LABELS.classes),\n",
+    "))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f7778e63",
+   "source": "### Entity-level Evaluation (highest-confidence aggregation)\n\nWhen an entity has multiple positions, each triplet produces a separate 3-class prediction.\nHere we aggregate by keeping the triplet with the highest P(yes) per entity,\nmatching the production inference behavior.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "1b3b49a9",
+   "source": "import torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nfrom src.models.distillbert import reconstruct_triplets\n\nds_test = EntitySentimentDataset(test_exs, tokenizer, cfg.max_len)\nloader = DataLoader(ds_test, batch_size=cfg.per_device_eval_batch_size, shuffle=False)\n\nall_yes_probs, all_bin_labels = [], []\nmodel.eval()\nwith torch.no_grad():\n    for batch in loader:\n        logits = model(\n            input_ids=batch[\"input_ids\"].to(device),\n            attention_mask=batch[\"attention_mask\"].to(device),\n        ).logits\n        all_yes_probs.extend(F.softmax(logits, dim=-1)[:, 1].cpu().tolist())\n        all_bin_labels.extend(batch[\"labels\"].tolist())\n\nsentiments = list(SENTIMENT_LABELS.classes)\nentity_best: dict[tuple, tuple[int, int, float]] = {}\n\nfor i in range(0, len(test_exs) - 2, 3):\n    triplet_probs = all_yes_probs[i:i + 3]\n    triplet_labels = all_bin_labels[i:i + 3]\n    pred_3class = int(np.argmax(triplet_probs))\n    label_3class = int(np.argmax(triplet_labels))\n    conf = max(triplet_probs)\n\n    ex = test_exs[i]\n    key = (ex[\"sample_id\"], ex[\"entity_id\"])\n    if key not in entity_best or conf > entity_best[key][2]:\n        entity_best[key] = (pred_3class, label_3class, conf)\n\nent_preds = [v[0] for v in entity_best.values()]\nent_labels = [v[1] for v in entity_best.values()]\n\nmacro_f1_entity = f1_score(ent_labels, ent_preds, average=\"macro\")\nprint(f\"Entity-level test macro-F1: {macro_f1_entity:.4f}\")\nprint(f\"({len(entity_best)} unique entities from {len(test_exs) // 3} position-level triplets)\")\nprint()\nprint(classification_report(\n    ent_labels, ent_preds,\n    target_names=sentiments,\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e7ba646b",
+   "source": "## 5. Training & Evaluation — Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "225f6ffb",
+   "source": "model_dedup = AutoModelForSequenceClassification.from_pretrained(\n    cfg.model_name,\n    num_labels=mode_cfg.labels.num_labels,\n    id2label=mode_cfg.labels.id2label,\n    label2id=mode_cfg.labels.label2id,\n)\n\ntrain_ds_dedup = EntitySentimentDataset(train_dedup, tokenizer, cfg.max_len)\nval_ds_dedup = EntitySentimentDataset(val_dedup, tokenizer, cfg.max_len)\n\ntraining_args_dedup = cfg.to_training_arguments()\ntraining_args_dedup.output_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntraining_args_dedup.logging_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\", \"logs\")\n\nclass_weights_dedup = compute_class_weights(train_dedup, n_classes=mode_cfg.labels.num_labels)\nprint(f\"Class weights (dedup): {class_weights_dedup.tolist()}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "3ff92fc7",
+   "source": "trainer_dedup = WeightedLossTrainer(\n    model=model_dedup,\n    args=training_args_dedup,\n    train_dataset=train_ds_dedup,\n    eval_dataset=val_ds_dedup,\n    compute_metrics=make_compute_metrics(cfg.mode),\n    callbacks=[EarlyStoppingCallback(\n        early_stopping_patience=cfg.early_stopping_patience\n    )],\n    class_weights=class_weights_dedup,\n)\n\ntrainer_dedup.train()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "7c4c73e3",
+   "source": "output_dir_dedup = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntrainer_dedup.save_model(output_dir_dedup)\ntokenizer.save_pretrained(output_dir_dedup)\n\nwith open(os.path.join(output_dir_dedup, \"mode.txt\"), \"w\") as f:\n    f.write(cfg.mode)\n\nprint(f\"Dedup model saved to '{output_dir_dedup}'\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "075b4cd2",
+   "source": "### Dedup Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "b92449e1",
+   "source": "eval_logs_d = [l for l in trainer_dedup.state.log_history if \"eval_loss\" in l]\n\nrows_d = []\nfor l in eval_logs_d:\n    rows_d.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1_3class\": round(l[\"eval_macro_f1\"], 4),\n        \"bin_accuracy\": round(l.get(\"eval_bin_accuracy\", 0), 4),\n        \"bin_f1_yes\": round(l.get(\"eval_bin_f1_yes\", 0), 4),\n    })\n\ndf_metrics_dedup = pd.DataFrame(rows_d)\ndf_metrics_dedup",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "36eee903",
+   "source": "### Dedup Test Evaluation",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "4b86af50",
+   "source": "macro_f1_dedup, test_preds_dedup, test_labels_dedup = evaluate_qa_b_test(\n    model=model_dedup,\n    tokenizer=tokenizer,\n    test_exs=test_dedup,\n    max_len=cfg.max_len,\n    batch_size=cfg.per_device_eval_batch_size,\n    device=device,\n)\n\nprint(f\"Test macro-F1 (3-class reconstructed, dedup): {macro_f1_dedup:.4f}\")\nprint()\nprint(classification_report(\n    test_labels_dedup, test_preds_dedup,\n    target_names=list(SENTIMENT_LABELS.classes),\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f4b3d712",
+   "source": "### Comparison: All Positions vs Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "5dd74e22",
+   "source": "comparison = pd.DataFrame({\n    \"variant\": [\"all positions (per-position)\", \"all positions (entity-level)\", \"deduplicated\"],\n    \"train_examples\": [len(train_exs), len(train_exs), len(train_dedup)],\n    \"test_examples\": [len(test_exs), f\"{len(test_exs) // 3} -> {len(entity_best)} entities\", len(test_dedup)],\n    \"test_macro_f1\": [round(macro_f1, 4), round(macro_f1_entity, 4), round(macro_f1_dedup, 4)],\n})\ncomparison",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/train_qa_m.ipynb

@@ -0,0 +1,373 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b1",
+   "metadata": {},
+   "source": [
+    "# Train: QA-M Mode\n",
+    "\n",
+    "Fine-tune DistilBERT with sentence-pair QA-M formulation (Sun et al. 2019).\n",
+    "\n",
+    "- Seg A: entity-centred context window\n",
+    "- Seg B: \"What do you think of the sentiment of the {entity_type} {entity} ?\"\n",
+    "- Label: 3-way (negative, neutral, positive)\n",
+    "\n",
+    "**Structure**\n",
+    "1. Load Config & Data\n",
+    "2. Tokenizer & Model\n",
+    "3. Training\n",
+    "4. Evaluation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "import random\n",
+    "from collections import Counter\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "torch.backends.cuda.enable_cudnn_sdp(False)\n",
+    "from sklearn.metrics import classification_report, f1_score\n",
+    "from transformers import (\n",
+    "    AutoModelForSequenceClassification,\n",
+    "    AutoTokenizer,\n",
+    "    EarlyStoppingCallback,\n",
+    ")\n",
+    "\n",
+    "sys.path.insert(0, os.path.abspath(\"..\"))\n",
+    "\n",
+    "from src.schemas.config import TrainingConfig\n",
+    "from src.schemas.labels import MODES, SENTIMENT_LABELS\n",
+    "from src.models.dataset import (\n",
+    "    EntitySentimentDataset,\n",
+    "    deduplicate_positions,\n",
+    "    flatten_to_examples,\n",
+    "    load_data,\n",
+    "    split_data,\n",
+    ")\n",
+    "from src.models.distillbert import (\n",
+    "    WeightedLossTrainer,\n",
+    "    compute_class_weights,\n",
+    "    make_compute_metrics,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b3",
+   "metadata": {},
+   "source": [
+    "## 1. Load Config & Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cfg = TrainingConfig.from_json(\"../data/config_qa_m.json\")\n",
+    "mode_cfg = MODES[cfg.mode]\n",
+    "\n",
+    "random.seed(cfg.seed)\n",
+    "np.random.seed(cfg.seed)\n",
+    "torch.manual_seed(cfg.seed)\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}  |  Mode: {cfg.mode}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5",
+   "metadata": {},
+   "outputs": [],
+   "source": "samples = load_data(os.path.join(\"..\", cfg.data_path))\nexamples = flatten_to_examples(samples, mode=cfg.mode)\n\nn_ents = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in examples})\nprint(f\"{len(samples)} samples -> {n_ents} entities -> {len(examples)} examples\")\n\ndist = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in examples)\nprint(f\"Label distribution: {dict(dist)}\")"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_exs, val_exs, test_exs = split_data(\n",
+    "    examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n",
+    ")\n",
+    "print(f\"Train: {len(train_exs)} | Val: {len(val_exs)} | Test: {len(test_exs)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4f074039",
+   "source": "### Deduplicated dataset (one position per entity)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "417c40e2",
+   "source": "deduped_samples = deduplicate_positions(samples)\ndedup_examples = flatten_to_examples(deduped_samples, mode=cfg.mode)\n\nn_ents_dedup = len({(e[\"sample_id\"], e[\"entity_id\"]) for e in dedup_examples})\nprint(f\"{len(deduped_samples)} samples -> {n_ents_dedup} entities -> {len(dedup_examples)} examples (deduplicated)\")\n\ndist_dedup = Counter(SENTIMENT_LABELS.id2label[e[\"label\"]] for e in dedup_examples)\nprint(f\"Label distribution: {dict(dist_dedup)}\")\n\ntrain_dedup, val_dedup, test_dedup = split_data(\n    dedup_examples, cfg.val_split, cfg.test_split, seed=cfg.seed\n)\nprint(f\"Train: {len(train_dedup)} | Val: {len(val_dedup)} | Test: {len(test_dedup)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7",
+   "metadata": {},
+   "source": [
+    "## 2. Tokenizer & Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n",
+    "\n",
+    "model = AutoModelForSequenceClassification.from_pretrained(\n",
+    "    cfg.model_name,\n",
+    "    num_labels=mode_cfg.labels.num_labels,\n",
+    "    id2label=mode_cfg.labels.id2label,\n",
+    "    label2id=mode_cfg.labels.label2id,\n",
+    ")\n",
+    "print(f\"Num labels: {mode_cfg.labels.num_labels}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "08b28688",
+   "source": "### Truncation Analysis",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "a69f3e30",
+   "source": "import matplotlib.pyplot as plt\n\nlengths = []\nfor ex in examples:\n    enc = tokenizer(ex[\"seg_a\"], ex[\"seg_b\"], truncation=False)\n    lengths.append(len(enc[\"input_ids\"]))\n\nlengths = np.array(lengths)\nn_truncated = (lengths > cfg.max_len).sum()\nprint(f\"Max length: {cfg.max_len}\")\nprint(f\"Total examples: {len(lengths)}\")\nprint(f\"Truncated: {n_truncated} ({100 * n_truncated / len(lengths):.1f}%)\")\nprint(f\"Token length — min: {lengths.min()}, median: {int(np.median(lengths))}, \"\n      f\"mean: {lengths.mean():.0f}, max: {lengths.max()}\")\n\nfig, ax = plt.subplots(figsize=(10, 3))\nax.hist(lengths, bins=50, edgecolor=\"black\", linewidth=0.5)\nax.axvline(cfg.max_len, color=\"red\", linestyle=\"--\", label=f\"max_len={cfg.max_len}\")\nax.set_xlabel(\"Token count\")\nax.set_ylabel(\"Examples\")\nax.set_title(\"Input Length Distribution (qa_m)\")\nax.legend()\nplt.tight_layout()\nplt.show()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b9",
+   "metadata": {},
+   "source": [
+    "## 3. Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_ds = EntitySentimentDataset(train_exs, tokenizer, cfg.max_len)\n",
+    "val_ds = EntitySentimentDataset(val_exs, tokenizer, cfg.max_len)\n",
+    "\n",
+    "training_args = cfg.to_training_arguments()\n",
+    "training_args.output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "training_args.logging_dir = os.path.join(\"..\", cfg.output_dir, \"logs\")\n",
+    "\n",
+    "class_weights = compute_class_weights(train_exs, n_classes=mode_cfg.labels.num_labels)\n",
+    "print(f\"Class weights: {class_weights.tolist()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = WeightedLossTrainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_ds,\n",
+    "    eval_dataset=val_ds,\n",
+    "    compute_metrics=make_compute_metrics(cfg.mode),\n",
+    "    callbacks=[EarlyStoppingCallback(\n",
+    "        early_stopping_patience=cfg.early_stopping_patience\n",
+    "    )],\n",
+    "    class_weights=class_weights,\n",
+    ")\n",
+    "\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output_dir = os.path.join(\"..\", cfg.output_dir)\n",
+    "trainer.save_model(output_dir)\n",
+    "tokenizer.save_pretrained(output_dir)\n",
+    "\n",
+    "with open(os.path.join(output_dir, \"mode.txt\"), \"w\") as f:\n",
+    "    f.write(cfg.mode)\n",
+    "\n",
+    "print(f\"Model saved to '{output_dir}'\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c2bd2384",
+   "source": "### Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "dc800309",
+   "source": "import pandas as pd\n\neval_logs = [l for l in trainer.state.log_history if \"eval_loss\" in l]\n\nrows = []\nfor l in eval_logs:\n    rows.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1\": round(l[\"eval_macro_f1\"], 4),\n        \"f1_negative\": round(l.get(\"eval_f1_negative\", 0), 4),\n        \"f1_neutral\": round(l.get(\"eval_f1_neutral\", 0), 4),\n        \"f1_positive\": round(l.get(\"eval_f1_positive\", 0), 4),\n    })\n\ndf_metrics = pd.DataFrame(rows)\ndf_metrics",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b13",
+   "metadata": {},
+   "source": [
+    "## 4. Evaluation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_ds = EntitySentimentDataset(test_exs, tokenizer, cfg.max_len)\n",
+    "pred_output = trainer.predict(test_ds)\n",
+    "\n",
+    "test_preds = np.argmax(pred_output.predictions, axis=-1).tolist()\n",
+    "test_labels = pred_output.label_ids.tolist()\n",
+    "macro_f1 = f1_score(test_labels, test_preds, average=\"macro\")\n",
+    "\n",
+    "print(f\"Test macro-F1: {macro_f1:.4f}\")\n",
+    "print()\n",
+    "print(classification_report(\n",
+    "    test_labels, test_preds,\n",
+    "    target_names=list(SENTIMENT_LABELS.classes),\n",
+    "))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd99b7e0",
+   "source": "### Entity-level Evaluation (highest-confidence aggregation)\n\nWhen an entity has multiple positions, each produces a separate prediction.\nHere we aggregate by keeping the prediction with the highest confidence per entity,\nmatching the production inference behavior.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "aa42bf08",
+   "source": "import torch.nn.functional as F\n\nprobs = F.softmax(torch.tensor(pred_output.predictions), dim=-1).numpy()\nconfidences = probs.max(axis=-1)\npreds_all = np.argmax(probs, axis=-1)\n\nsentiments = list(SENTIMENT_LABELS.classes)\nentity_best: dict[tuple, tuple[int, int, float]] = {}\nfor ex, pred_id, conf, label in zip(test_exs, preds_all, confidences, test_labels):\n    key = (ex[\"sample_id\"], ex[\"entity_id\"])\n    if key not in entity_best or conf > entity_best[key][2]:\n        entity_best[key] = (int(pred_id), int(label), float(conf))\n\nent_preds = [v[0] for v in entity_best.values()]\nent_labels = [v[1] for v in entity_best.values()]\n\nmacro_f1_entity = f1_score(ent_labels, ent_preds, average=\"macro\")\nprint(f\"Entity-level test macro-F1: {macro_f1_entity:.4f}\")\nprint(f\"({len(entity_best)} unique entities from {len(test_exs)} position-level examples)\")\nprint()\nprint(classification_report(\n    ent_labels, ent_preds,\n    target_names=sentiments,\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "925dff05",
+   "source": "## 5. Training & Evaluation — Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "88e587d7",
+   "source": "model_dedup = AutoModelForSequenceClassification.from_pretrained(\n    cfg.model_name,\n    num_labels=mode_cfg.labels.num_labels,\n    id2label=mode_cfg.labels.id2label,\n    label2id=mode_cfg.labels.label2id,\n)\n\ntrain_ds_dedup = EntitySentimentDataset(train_dedup, tokenizer, cfg.max_len)\nval_ds_dedup = EntitySentimentDataset(val_dedup, tokenizer, cfg.max_len)\n\ntraining_args_dedup = cfg.to_training_arguments()\ntraining_args_dedup.output_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntraining_args_dedup.logging_dir = os.path.join(\"..\", cfg.output_dir + \"_dedup\", \"logs\")\n\nclass_weights_dedup = compute_class_weights(train_dedup, n_classes=mode_cfg.labels.num_labels)\nprint(f\"Class weights (dedup): {class_weights_dedup.tolist()}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "a4f5bc5e",
+   "source": "trainer_dedup = WeightedLossTrainer(\n    model=model_dedup,\n    args=training_args_dedup,\n    train_dataset=train_ds_dedup,\n    eval_dataset=val_ds_dedup,\n    compute_metrics=make_compute_metrics(cfg.mode),\n    callbacks=[EarlyStoppingCallback(\n        early_stopping_patience=cfg.early_stopping_patience\n    )],\n    class_weights=class_weights_dedup,\n)\n\ntrainer_dedup.train()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "20ebd649",
+   "source": "output_dir_dedup = os.path.join(\"..\", cfg.output_dir + \"_dedup\")\ntrainer_dedup.save_model(output_dir_dedup)\ntokenizer.save_pretrained(output_dir_dedup)\n\nwith open(os.path.join(output_dir_dedup, \"mode.txt\"), \"w\") as f:\n    f.write(cfg.mode)\n\nprint(f\"Dedup model saved to '{output_dir_dedup}'\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4a0e255e",
+   "source": "### Dedup Validation Metrics per Epoch",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "69f51f21",
+   "source": "eval_logs_d = [l for l in trainer_dedup.state.log_history if \"eval_loss\" in l]\n\nrows_d = []\nfor l in eval_logs_d:\n    rows_d.append({\n        \"epoch\": int(l[\"epoch\"]),\n        \"eval_loss\": round(l[\"eval_loss\"], 4),\n        \"macro_f1\": round(l[\"eval_macro_f1\"], 4),\n        \"f1_negative\": round(l.get(\"eval_f1_negative\", 0), 4),\n        \"f1_neutral\": round(l.get(\"eval_f1_neutral\", 0), 4),\n        \"f1_positive\": round(l.get(\"eval_f1_positive\", 0), 4),\n    })\n\ndf_metrics_dedup = pd.DataFrame(rows_d)\ndf_metrics_dedup",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "54409311",
+   "source": "### Dedup Test Evaluation",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "b9f18bc4",
+   "source": "test_ds_dedup = EntitySentimentDataset(test_dedup, tokenizer, cfg.max_len)\npred_output_dedup = trainer_dedup.predict(test_ds_dedup)\n\ntest_preds_dedup = np.argmax(pred_output_dedup.predictions, axis=-1).tolist()\ntest_labels_dedup = pred_output_dedup.label_ids.tolist()\nmacro_f1_dedup = f1_score(test_labels_dedup, test_preds_dedup, average=\"macro\")\n\nprint(f\"Test macro-F1 (dedup): {macro_f1_dedup:.4f}\")\nprint()\nprint(classification_report(\n    test_labels_dedup, test_preds_dedup,\n    target_names=list(SENTIMENT_LABELS.classes),\n))",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "900bdb55",
+   "source": "### Comparison: All Positions vs Deduplicated",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "111ce91e",
+   "source": "comparison = pd.DataFrame({\n    \"variant\": [\"all positions (per-position)\", \"all positions (entity-level)\", \"deduplicated\"],\n    \"train_examples\": [len(train_exs), len(train_exs), len(train_dedup)],\n    \"test_examples\": [len(test_exs), f\"{len(test_exs)} -> {len(entity_best)} entities\", len(test_dedup)],\n    \"test_macro_f1\": [round(macro_f1, 4), round(macro_f1_entity, 4), round(macro_f1_dedup, 4)],\n})\ncomparison",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}