feat(evaluation): add beam search, metrics pipeline, and stabilized training workflow

configs/train/stabilized.yaml

@@ -0,0 +1,84 @@
+# =============================================================================
+# configs/train/stabilized.yaml — first experimental run with the opt-in
+# training-stability primitives turned on.
+# -----------------------------------------------------------------------------
+# Identical to configs/base.yaml except for the four flags called out in the
+# `train:` section below. Every other field mirrors the IEEE notebook verbatim
+# so this run is comparable to the baseline at the same seed and architecture.
+#
+# Why a complete config (not a thin override)?
+#   scripts/train.py only accepts --config; there is no --override merge mode
+#   in the CLI (the README mentions one but it's aspirational, not implemented).
+#   Duplicating the values here is the smallest correct change that keeps
+#   base.yaml itself untouched — which was the explicit requirement for this
+#   experiment phase.
+#
+# Usage:
+#   python -m scripts.train --config configs/train/stabilized.yaml \
+#       --output-dir outputs/runs/stabilized
+#
+# Compare against the baseline by training the same code twice — once with
+# configs/base.yaml, once with this file — and diffing the resulting
+# results/<run_id>/metrics.json files.
+# =============================================================================
+
+data:
+  base_path: data/coco2017
+  annotations_filename: captions_train2017.json
+  images_subdir: train2017
+  sample_size: 120000           # Same sample as base.yaml — comparability matters
+  train_val_split: 0.8
+
+model:
+  embedding_dim: 512
+  units: 512
+  max_length: 40
+  vocabulary_size: 15000
+  encoder_num_heads: 1
+  decoder_num_heads: 8
+  decoder_dropout_inner: 0.3
+  decoder_dropout_outer: 0.5
+  decoder_attention_dropout: 0.1
+
+train:
+  epochs: 10
+  batch_size: 64
+  buffer_size: 1000
+  early_stopping_patience: 3
+  seed: 42
+  learning_rate: 0.001
+  weights_filename: model.h5
+
+  # ---- the four flags this experiment is actually testing -------------------
+  # Label smoothing 0.1 softens the cross-entropy target so the decoder
+  # cannot collapse onto a handful of high-frequency tokens. Standard
+  # transformer captioning recipe (BLIP, ViT-GPT2, GIT all use it).
+  label_smoothing: 0.1
+
+  # Warmup + cosine decay replaces the bare constant Adam LR. Transformers
+  # trained from scratch with no warmup tend to settle into a "safe captions"
+  # basin where every output looks like "a man standing ...". Cosine decay
+  # then anneals smoothly toward min_learning_rate.
+  lr_schedule: cosine
+  warmup_steps: 500             # ~1/3 of an epoch at batch 64, sample 120k
+  cosine_decay_steps: null      # null -> trainer derives from steps_per_epoch * epochs
+  min_learning_rate: 0.0
+
+  # Restore conventional behaviour: dropout OFF during validation, accuracy
+  # tracker weighted by token count. This gives a clean val_loss signal so
+  # EarlyStopping fires on a real plateau rather than on dropout noise.
+  honour_training_flag_in_test_step: true
+
+serve:
+  max_upload_bytes: 10485760
+  decode_strategy: greedy       # Decode strategy is selected at evaluate time
+  beam_width: 4                 # Stored defaults for `scripts.evaluate --decode-strategy beam`
+  length_penalty: 0.7
+  repetition_penalty: 1.0
+  no_repeat_ngram_size: 3
+  cors_allowed_origins:
+    - http://localhost:3000
+    - http://localhost:5173
+    - http://localhost:5174
+    - http://127.0.0.1:5173
+    - http://127.0.0.1:5174

docs/STABILIZED_TRAINING_RUNBOOK.md

@@ -0,0 +1,236 @@
+# Stabilized-config training runbook
+
+Hand-off doc for running the first real `configs/train/stabilized.yaml`
+experiment on Kaggle (or any free GPU notebook host). Stays out of the
+package proper because it's a one-shot operational guide, not project
+documentation.
+
+## What you'll produce
+
+By the end of this runbook you will have, **per decode strategy**:
+
+```
+results/<run_id>/
+    metrics.json        # BLEU-1..4, ROUGE-L, METEOR, CIDEr
+    predictions.jsonl   # one row per validation image
+    diagnostics.jsonl   # per-sample length / repetition / sentence BLEU
+    run_meta.json       # decode flags, timestamp, model id
+    report.md           # human-readable summary
+```
+
+Download both directories (`results/<greedy_run_id>/` and
+`results/<beam_run_id>/`) and the trained `models/v1.0.0/` artefacts as
+Kaggle outputs. That's the entire input for the Phase-2-stabilization
+comparison.
+
+## Caveats before you start
+
+1. **`requirements.txt` pins `tensorflow-cpu==2.15.0`.** On Kaggle GPU you
+   must install `tensorflow==2.15.0` instead — see step 3 below. The
+   `Keras 3` warning in the original pin comment still applies: do not
+   upgrade to 2.16+, save/load semantics for `TextVectorization` change.
+2. **Kaggle session limit is 9 hours per run.** A full 10-epoch pass on
+   120k captions at batch 64 fits comfortably on a T4 (well under 9h),
+   but plan to checkpoint after every epoch in case the session restarts.
+   `default_callbacks(...)` already writes `best.h5` on val_loss
+   improvement, so this is handled.
+3. **Free-tier Kaggle gives ~30 GB of attached-dataset space.** COCO 2017
+   train2017 is ~19 GB. Use the public `awsaf49/coco-2017-dataset`
+   mount; do not re-upload it as your own dataset.
+4. **The `data.base_path` in `stabilized.yaml` points at `data/coco2017`**
+   (the project-relative path). On Kaggle the COCO mount is at
+   `/kaggle/input/coco-2017-dataset/coco2017`. You override at runtime
+   without editing the YAML using the env-var override pattern (step 4).
+
+## Kaggle notebook cells
+
+Paste each block into a separate cell so you can re-run individual
+steps without restarting.
+
+### Cell 1 — Attach the COCO dataset
+
+In the Kaggle notebook UI: **+ Add Data → Search "coco-2017" →**
+`awsaf49/coco-2017-dataset`. After attaching, verify:
+
+```python
+!ls /kaggle/input/coco-2017-dataset/coco2017
+# expect: annotations/  train2017/  val2017/
+```
+
+### Cell 2 — Pull the repo
+
+```python
+!git clone https://github.com/<your-user>/image-captioning-system.git
+%cd image-captioning-system
+```
+
+(If your repo isn't public, upload it as a Kaggle dataset and reference
+it via `/kaggle/input/<dataset-slug>/`.)
+
+### Cell 3 — Install with the GPU TF wheel
+
+```python
+# Replace the CPU TF pin with the GPU-capable wheel. The rest of the
+# pinned versions in requirements.txt are GPU/CPU-agnostic.
+!pip install -q tensorflow==2.15.0
+!pip install -q -r requirements-dev.txt -r requirements-eval.txt
+!pip install -q -e .
+```
+
+Verify GPU is visible to TF:
+
+```python
+import tensorflow as tf
+print("GPUs:", tf.config.list_physical_devices("GPU"))
+```
+
+Expected output: at least one GPU listed (typically `T4` or `P100`).
+
+### Cell 4 — Train with the stabilized config
+
+The `data.base_path` field in `stabilized.yaml` is overridden via
+env-var so the YAML stays unedited:
+
+```python
+import os
+os.environ["CAPTIONING__DATA__BASE_PATH"] = "/kaggle/input/coco-2017-dataset/coco2017"
+
+!python -m scripts.train \
+    --config configs/train/stabilized.yaml \
+    --output-dir outputs/runs/stabilized
+```
+
+Expected wall-clock on a T4: ~30-50 min per epoch, ~5-8 hours for 10
+epochs. EarlyStopping (patience 3) typically fires before epoch 10.
+
+What lands in `outputs/runs/stabilized/`:
+- `best.h5`         — best val_loss checkpoint (ModelCheckpoint)
+- `model.h5`        — final-epoch weights
+- `vocab.pkl/json`  — fitted tokenizer
+- `history.json`    — train/val loss per epoch
+- `training_log.csv`— CSVLogger output
+
+### Cell 5 — Promote the trained checkpoint
+
+The evaluation and serving paths expect `models/v1.0.0/model.h5` +
+tokenizer. Promote the best checkpoint there:
+
+```python
+!mkdir -p models/v1.0.0
+!cp outputs/runs/stabilized/best.h5    models/v1.0.0/model.h5
+!cp outputs/runs/stabilized/vocab.pkl  models/v1.0.0/vocab.pkl
+!cp outputs/runs/stabilized/vocab.json models/v1.0.0/vocab.json
+```
+
+### Cell 6 — Greedy evaluation
+
+```python
+os.environ["CAPTIONING__DATA__BASE_PATH"] = "/kaggle/input/coco-2017-dataset/coco2017"
+
+!python -m scripts.evaluate \
+    --config configs/train/stabilized.yaml \
+    --weights models/v1.0.0/model.h5 \
+    --tokenizer-dir models/v1.0.0 \
+    --results-root results \
+    --run-id stabilized-greedy \
+    --model-id inceptionv3-transformer-stabilized \
+    --decode-strategy greedy \
+    --max-samples 500
+```
+
+If METEOR/Java is unavailable on the Kaggle image (it usually is, but
+the wheel can fail at import):
+
+```python
+!apt-get install -y openjdk-11-jre-headless  # if METEOR errors out
+# or:
+!python -m scripts.evaluate ... --skip-meteor
+```
+
+### Cell 7 — Beam evaluation
+
+```python
+!python -m scripts.evaluate \
+    --config configs/train/stabilized.yaml \
+    --weights models/v1.0.0/model.h5 \
+    --tokenizer-dir models/v1.0.0 \
+    --results-root results \
+    --run-id stabilized-beam-w4-lp07-nrn3 \
+    --model-id inceptionv3-transformer-stabilized \
+    --decode-strategy beam \
+    --beam-width 4 \
+    --length-penalty 0.7 \
+    --no-repeat-ngram-size 3 \
+    --max-samples 500
+```
+
+### Cell 8 — Per-sample inspection (qualitative review)
+
+```python
+!python -m scripts.inspect_predictions \
+    --config configs/train/stabilized.yaml \
+    --weights models/v1.0.0/model.h5 \
+    --tokenizer-dir models/v1.0.0 \
+    --decode-strategy beam \
+    --beam-width 4 \
+    --n-samples 30 \
+    --output results/stabilized-beam-w4-lp07-nrn3/qualitative.jsonl
+```
+
+### Cell 9 — Persist outputs
+
+Kaggle persists anything under `/kaggle/working/` between sessions and
+makes it downloadable as a notebook output. Move the results there:
+
+```python
+!mkdir -p /kaggle/working/handoff
+!cp -r results /kaggle/working/handoff/
+!cp -r models /kaggle/working/handoff/
+!ls -la /kaggle/working/handoff/results
+```
+
+### Cell 10 — Pre-flight summary print
+
+```python
+import json
+for run in ("stabilized-greedy", "stabilized-beam-w4-lp07-nrn3"):
+    m = json.load(open(f"results/{run}/metrics.json"))
+    print(run, "->", {k: m[k] for k in ("bleu1", "bleu4", "rouge_l", "meteor", "cider")})
+```
+
+This should print two lines that you copy back into the next chat
+session — that's all I need to do Steps 4-6 of the original request.
+
+## What to bring back to this conversation
+
+When you return:
+1. The two `metrics.json` files (or just the dicts from Cell 10 — same content).
+2. The two `report.md` files for the human-readable view.
+3. The `diagnostics.jsonl` from at least the beam run, so qualitative
+   inspection isn't blind.
+4. (Optional but useful) the `history.json` from the training run — lets
+   us check whether warmup+cosine actually flattened the validation
+   curve as predicted.
+
+With those four artefacts I can do the full quantitative comparison,
+qualitative inspection, and bottleneck analysis you originally asked
+for, with real numbers and no fabrication.
+
+## If something goes sideways
+
+* **OOM on the T4 during training** — drop `train.batch_size` to 32
+  (it's already 64; halving keeps the same effective optimization but
+  doubles steps_per_epoch and doubles `warmup_steps` too — the trainer
+  auto-derives `cosine_decay_steps`, so just batch_size and warmup_steps
+  need adjusting).
+* **METEOR fails to import** — pass `--skip-meteor`; the other four
+  metrics still write to `metrics.json`.
+* **Session times out mid-training** — Kaggle saves the working
+  directory; re-attach the notebook, `cp` the partial weights from
+  `outputs/runs/stabilized/best.h5` to `models/v1.0.0/`, and continue
+  with evaluation only. EarlyStopping likely fired anyway.
+* **`scripts/train.py` errors with `FileNotFoundError`** — verify
+  `os.environ["CAPTIONING__DATA__BASE_PATH"]` is set in the *same cell*
+  as the `!python` call, not a prior cell (Jupyter cell-level env vars
+  only propagate to subprocesses started from the same cell on some
+  Kaggle images).

scripts/evaluate.py

@@ -1,25 +1,47 @@
 """Evaluate a trained model on the COCO validation split.
 
 Usage:
+    # Full benchmark-ready evaluation (recommended) — writes
+    # results/<run_id>/{metrics.json, predictions.jsonl, diagnostics.jsonl, ...}
     python -m scripts.evaluate \\
         --config configs/base.yaml \\
         --weights models/v1.0.0/model.h5 \\
         --tokenizer-dir models/v1.0.0 \\
-        --report docs/results/v1.0.0.md \\
+        --results-root results \\
         --max-samples 500
+
+    # Optional: produce a single Markdown report at a chosen path
+    python -m scripts.evaluate ... --report docs/results/v1.0.0.md
+
+What this script produces (per run):
+    metrics.json        — corpus BLEU-1..4, ROUGE-L, METEOR, CIDEr
+    predictions.jsonl   — image / prediction / references for downstream tools
+    diagnostics.jsonl   — per-sample length / repetition / sentence BLEU flags
+    run_meta.json       — model id, decode strategy, beam width, timestamp
+    report.md           — human-readable summary
+
+Phase 3 benchmark code joins multiple ``results/<run_id>/`` directories to
+plot BLEU-4 / CIDEr / latency across models.
 """
 
 from __future__ import annotations
 
-import json
+from datetime import datetime, timezone
 from pathlib import Path
+from typing import cast
 
 import click
 
 from captioning.config import load_config
 from captioning.data import load_coco_annotations, make_image_level_splits
-from captioning.evaluation import corpus_bleu_score
+from captioning.evaluation import (
+    RunMeta,
+    compute_all_metrics,
+    diagnose_many,
+    write_run_artifacts,
+)
 from captioning.inference import CaptionPredictor
+from captioning.inference.predictor import DecodeStrategy
 from captioning.preprocessing import preprocess_caption
 from captioning.utils import configure_logging, get_logger, set_global_seed
 
@@ -32,12 +54,39 @@ log = get_logger(__name__)
 )
 @click.option("--weights", required=True, type=click.Path(exists=True, path_type=Path))
 @click.option("--tokenizer-dir", required=True, type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "--results-root",
+    type=click.Path(path_type=Path),
+    default=Path("results"),
+    help="Parent directory for the per-run sub-folder (results/<run_id>/).",
+)
+@click.option(
+    "--run-id",
+    type=str,
+    default=None,
+    help="Sub-folder name under --results-root. Defaults to a UTC timestamp.",
+)
+@click.option(
+    "--model-id",
+    type=str,
+    default="inceptionv3-transformer-v1",
+    help="Identifier used by Phase 3 cross-model joining of metrics.",
+)
+@click.option(
+    "--decode-strategy",
+    type=click.Choice(["greedy", "beam"]),
+    default=None,
+    help="Override config.serve.decode_strategy for this run.",
+)
+@click.option("--beam-width", type=int, default=None, help="Beam width (only used with beam).")
+@click.option("--length-penalty", type=float, default=None)
+@click.option("--repetition-penalty", type=float, default=None)
 @click.option(
     "--report",
     "report_path",
     default=None,
     type=click.Path(path_type=Path),
-    help="Optional path to write a Markdown report.",
+    help="Optional path to an additional human-readable Markdown report.",
 )
 @click.option(
     "--max-samples",
@@ -45,14 +94,35 @@ log = get_logger(__name__)
     type=int,
     help="Cap on validation examples (full val takes hours on CPU).",
 )
-def main(
+@click.option(
+    "--skip-meteor",
+    is_flag=True,
+    default=False,
+    help="Skip METEOR (avoids needing Java).",
+)
+@click.option(
+    "--skip-cider",
+    is_flag=True,
+    default=False,
+    help="Skip CIDEr.",
+)
+def main(  # — CLI option count is unavoidable
     config_path: Path,
     weights: Path,
     tokenizer_dir: Path,
+    results_root: Path,
+    run_id: str | None,
+    model_id: str,
+    decode_strategy: str | None,
+    beam_width: int | None,
+    length_penalty: float | None,
+    repetition_penalty: float | None,
     report_path: Path | None,
     max_samples: int,
+    skip_meteor: bool,
+    skip_cider: bool,
 ) -> None:
-    """Compute corpus BLEU-4 on the val split and (optionally) write a report."""
+    """Evaluate the model on the val split and write benchmark artefacts."""
     configure_logging()
     config = load_config(config_path)
     set_global_seed(config.train.seed)
@@ -75,8 +145,23 @@ def main(
         refs_by_image.setdefault(img, []).append(cap)
     image_paths = list(refs_by_image.keys())[:max_samples]
 
+    effective_strategy = decode_strategy or config.serve.decode_strategy
+    effective_beam_width = beam_width if beam_width is not None else config.serve.beam_width
+    effective_length_penalty = (
+        length_penalty if length_penalty is not None else config.serve.length_penalty
+    )
+    effective_repetition_penalty = (
+        repetition_penalty if repetition_penalty is not None else config.serve.repetition_penalty
+    )
+
     predictor = CaptionPredictor.from_artifacts(
-        weights_path=weights, tokenizer_dir=tokenizer_dir, config=config
+        weights_path=weights,
+        tokenizer_dir=tokenizer_dir,
+        config=config,
+        decode_strategy=cast("DecodeStrategy", effective_strategy),
+        beam_width=effective_beam_width,
+        length_penalty=effective_length_penalty,
+        repetition_penalty=effective_repetition_penalty,
     )
     predictor.warmup()
 
@@ -86,25 +171,70 @@ def main(
         predictions.append(predictor.predict_path(path))
         references.append(refs_by_image[path])
 
-    bleu = corpus_bleu_score(predictions, references)
-    log.info("evaluation_done", bleu4=bleu, n=len(predictions))
-    click.echo(f"BLEU-4: {bleu:.2f}  (n={len(predictions)})")
+    metrics = compute_all_metrics(
+        predictions,
+        references,
+        include_meteor=not skip_meteor,
+        include_cider=not skip_cider,
+    )
+    diagnostics = diagnose_many(image_paths, predictions, references)
+
+    run_id = run_id or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+    run_dir = Path(results_root) / run_id
+    meta = RunMeta(
+        model_id=model_id,
+        decode_strategy=effective_strategy,
+        weights_path=str(weights),
+        tokenizer_dir=str(tokenizer_dir),
+        n_samples=len(predictions),
+        max_length=config.model.max_length,
+        beam_width=effective_beam_width if effective_strategy == "beam" else None,
+        length_penalty=effective_length_penalty if effective_strategy == "beam" else None,
+        repetition_penalty=effective_repetition_penalty,
+    )
+    write_run_artifacts(
+        run_dir,
+        metrics=metrics,
+        meta=meta,
+        images=image_paths,
+        predictions=predictions,
+        references=references,
+        diagnostics=diagnostics,
+    )
+
+    log.info(
+        "evaluation_done",
+        run_dir=str(run_dir),
+        n=metrics.n_examples,
+        bleu4=metrics.bleu4,
+        rouge_l=metrics.rouge_l,
+        meteor=metrics.meteor,
+        cider=metrics.cider,
+    )
+    click.echo(f"Run directory: {run_dir}")
+    _echo_metric("BLEU-1", metrics.bleu1)
+    _echo_metric("BLEU-2", metrics.bleu2)
+    _echo_metric("BLEU-3", metrics.bleu3)
+    _echo_metric("BLEU-4", metrics.bleu4)
+    _echo_metric("ROUGE-L", metrics.rouge_l)
+    _echo_metric("METEOR", metrics.meteor)
+    _echo_metric("CIDEr", metrics.cider)
+    if metrics.errors:
+        click.echo(f"Skipped/failed: {sorted(metrics.errors)}")
 
     if report_path is not None:
         report_path.parent.mkdir(parents=True, exist_ok=True)
         report_path.write_text(
-            f"# Evaluation v1\n\n"
-            f"- BLEU-4: **{bleu:.2f}**\n"
-            f"- Examples: {len(predictions)}\n"
-            f"- Weights: `{weights}`\n",
-            encoding="utf-8",
-        )
-        json.dump(
-            {"bleu4": bleu, "n": len(predictions)},
-            (report_path.with_suffix(".json")).open("w", encoding="utf-8"),
-            indent=2,
+            (run_dir / "report.md").read_text(encoding="utf-8"), encoding="utf-8"
         )
 
 
+def _echo_metric(name: str, value: float | None) -> None:
+    if value is None:
+        click.echo(f"{name}: n/a")
+    else:
+        click.echo(f"{name}: {value:.2f}")
+
+
 if __name__ == "__main__":
     main()

scripts/inspect_predictions.py

@@ -0,0 +1,128 @@
+"""Per-sample inspection — print N validation predictions vs. ground truth.
+
+Usage:
+    python -m scripts.inspect_predictions \\
+        --config configs/base.yaml \\
+        --weights models/v1.0.0/model.h5 \\
+        --tokenizer-dir models/v1.0.0 \\
+        --n-samples 25
+
+This script answers the diagnostic question: *when the model is wrong, **how**
+is it wrong?* Each row shows the image filename, the predicted caption, one
+reference caption, sentence-level BLEU-4 / ROUGE-L, the prediction length,
+the longest repeated-token run, and a set of failure flags
+(``empty`` / ``very_short`` / ``repetitive`` / ``under_length``).
+
+Output also lands as ``diagnostics.jsonl`` so the same data can be loaded
+into pandas / DuckDB for ad-hoc grouping.
+"""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import cast
+
+import click
+
+from captioning.config import load_config
+from captioning.data import load_coco_annotations, make_image_level_splits
+from captioning.evaluation import (
+    diagnose_many,
+    format_diagnostic_row,
+    write_diagnostics_jsonl,
+)
+from captioning.inference import CaptionPredictor
+from captioning.inference.predictor import DecodeStrategy
+from captioning.preprocessing import preprocess_caption
+from captioning.utils import configure_logging, get_logger, set_global_seed
+
+log = get_logger(__name__)
+
+
+@click.command()
+@click.option(
+    "--config", "config_path", required=True, type=click.Path(exists=True, path_type=Path)
+)
+@click.option("--weights", required=True, type=click.Path(exists=True, path_type=Path))
+@click.option("--tokenizer-dir", required=True, type=click.Path(exists=True, path_type=Path))
+@click.option("--n-samples", type=int, default=25, help="Number of random val samples to inspect.")
+@click.option(
+    "--decode-strategy",
+    type=click.Choice(["greedy", "beam"]),
+    default=None,
+    help="Override decode strategy for this inspection.",
+)
+@click.option("--beam-width", type=int, default=None)
+@click.option(
+    "--output",
+    "output_path",
+    type=click.Path(path_type=Path),
+    default=None,
+    help="Optional path to write diagnostics.jsonl. Defaults to printing only.",
+)
+@click.option(
+    "--seed",
+    type=int,
+    default=None,
+    help="RNG seed for sample selection (defaults to config.train.seed).",
+)
+def main(
+    config_path: Path,
+    weights: Path,
+    tokenizer_dir: Path,
+    n_samples: int,
+    decode_strategy: str | None,
+    beam_width: int | None,
+    output_path: Path | None,
+    seed: int | None,
+) -> None:
+    """Sample N val images, run inference, and print prediction diagnostics."""
+    configure_logging()
+    config = load_config(config_path)
+    set_global_seed(config.train.seed)
+
+    df = load_coco_annotations(
+        base_path=config.data.base_path,
+        annotations_filename=config.data.annotations_filename,
+        images_subdir=config.data.images_subdir,
+        sample_size=config.data.sample_size,
+        seed=config.train.seed,
+        caption_preprocessor=preprocess_caption,
+    )
+    _, _, val_imgs, val_caps = make_image_level_splits(
+        df, train_fraction=config.data.train_val_split, seed=config.train.seed
+    )
+    refs_by_image: dict[str, list[str]] = {}
+    for img, cap in zip(val_imgs, val_caps, strict=True):
+        refs_by_image.setdefault(img, []).append(cap)
+
+    rng = random.Random(seed if seed is not None else config.train.seed)
+    picks = rng.sample(sorted(refs_by_image.keys()), k=min(n_samples, len(refs_by_image)))
+
+    effective_strategy = decode_strategy or config.serve.decode_strategy
+    effective_beam_width = beam_width if beam_width is not None else config.serve.beam_width
+    predictor = CaptionPredictor.from_artifacts(
+        weights_path=weights,
+        tokenizer_dir=tokenizer_dir,
+        config=config,
+        decode_strategy=cast("DecodeStrategy", effective_strategy),
+        beam_width=effective_beam_width,
+    )
+    predictor.warmup()
+
+    predictions = [predictor.predict_path(p) for p in picks]
+    references = [refs_by_image[p] for p in picks]
+    diagnostics = diagnose_many(picks, predictions, references)
+
+    for d in diagnostics:
+        click.echo(format_diagnostic_row(d))
+        click.echo("-" * 80)
+
+    if output_path is not None:
+        write_diagnostics_jsonl(diagnostics, output_path)
+        click.echo(f"Wrote diagnostics: {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/predict.py

@@ -1,21 +1,28 @@
 """CLI single-image inference.
 
 Usage:
+    # Greedy (default — same as the IEEE notebook)
     python -m scripts.predict \\
         --config configs/base.yaml \\
         --weights models/v1.0.0/model.h5 \\
         --tokenizer-dir models/v1.0.0 \\
         --image path/to/photo.jpg
+
+    # Beam search with explicit parameters
+    python -m scripts.predict ... --decode-strategy beam --beam-width 4 \\
+        --length-penalty 0.7 --repetition-penalty 1.1 --no-repeat-ngram-size 3
 """
 
 from __future__ import annotations
 
 from pathlib import Path
+from typing import cast
 
 import click
 
 from captioning.config import load_config
 from captioning.inference import CaptionPredictor
+from captioning.inference.predictor import DecodeStrategy
 from captioning.utils import configure_logging, get_logger
 
 log = get_logger(__name__)
@@ -28,7 +35,27 @@ log = get_logger(__name__)
 @click.option("--weights", required=True, type=click.Path(exists=True, path_type=Path))
 @click.option("--tokenizer-dir", required=True, type=click.Path(exists=True, path_type=Path))
 @click.option("--image", required=True, type=click.Path(exists=True, path_type=Path))
-def main(config_path: Path, weights: Path, tokenizer_dir: Path, image: Path) -> None:
+@click.option(
+    "--decode-strategy",
+    type=click.Choice(["greedy", "beam"]),
+    default=None,
+    help="Override config.serve.decode_strategy for this run.",
+)
+@click.option("--beam-width", type=int, default=None)
+@click.option("--length-penalty", type=float, default=None)
+@click.option("--repetition-penalty", type=float, default=None)
+@click.option("--no-repeat-ngram-size", type=int, default=None)
+def main(
+    config_path: Path,
+    weights: Path,
+    tokenizer_dir: Path,
+    image: Path,
+    decode_strategy: str | None,
+    beam_width: int | None,
+    length_penalty: float | None,
+    repetition_penalty: float | None,
+    no_repeat_ngram_size: int | None,
+) -> None:
     """Generate a caption for one image."""
     configure_logging()
     config = load_config(config_path)
@@ -37,6 +64,11 @@ def main(config_path: Path, weights: Path, tokenizer_dir: Path, image: Path) ->
         weights_path=weights,
         tokenizer_dir=tokenizer_dir,
         config=config,
+        decode_strategy=cast("DecodeStrategy | None", decode_strategy),
+        beam_width=beam_width,
+        length_penalty=length_penalty,
+        repetition_penalty=repetition_penalty,
+        no_repeat_ngram_size=no_repeat_ngram_size,
     )
     predictor.warmup()
     caption = predictor.predict_path(image)

src/captioning/config/schema.py

@@ -92,7 +92,15 @@ class ModelConfig(_StrictModel):
 
 
 class TrainConfig(_StrictModel):
-    """Optimisation hyperparameters."""
+    """Optimisation hyperparameters.
+
+    The Phase 1 baseline mirrors the IEEE notebook: constant LR, no label
+    smoothing, dropout-active validation (a notebook quirk preserved for
+    parity). The fields below the comment line are *opt-in* training-
+    stability knobs added during the caption-quality stabilisation phase.
+    Defaults keep every existing run byte-for-byte identical to the
+    notebook; flipping the flag in YAML opts a run into the modern recipe.
+    """
 
     epochs: int = 10
     batch_size: int = 64
@@ -102,16 +110,81 @@ class TrainConfig(_StrictModel):
     learning_rate: float = 1e-3  # Notebook uses Keras Adam default == 1e-3
     weights_filename: str = "model.h5"
 
+    # ---- opt-in stability flags (default values preserve notebook parity) ----
+    label_smoothing: float = 0.0
+    lr_schedule: str = "constant"  # "constant" | "cosine"
+    warmup_steps: int = 0
+    cosine_decay_steps: int | None = None  # If None, derived from epochs * steps_per_epoch
+    min_learning_rate: float = 0.0
+    honour_training_flag_in_test_step: bool = False  # parity-quirk override
+
+    @field_validator("label_smoothing")
+    @classmethod
+    def _validate_label_smoothing(cls, v: float) -> float:
+        if not 0.0 <= v < 1.0:
+            raise ValueError(f"label_smoothing must be in [0, 1), got {v}")
+        return v
+
+    @field_validator("lr_schedule")
+    @classmethod
+    def _validate_lr_schedule(cls, v: str) -> str:
+        if v not in {"constant", "cosine"}:
+            raise ValueError(f"lr_schedule must be 'constant' or 'cosine', got {v!r}")
+        return v
+
+    @field_validator("warmup_steps")
+    @classmethod
+    def _validate_warmup_steps(cls, v: int) -> int:
+        if v < 0:
+            raise ValueError(f"warmup_steps must be >= 0, got {v}")
+        return v
+
 
 class ServeConfig(_StrictModel):
     """Settings for the FastAPI backend (Phase 2). Defined here so the schema
-    is complete and tests don't have to mock a sub-config's existence."""
+    is complete and tests don't have to mock a sub-config's existence.
+
+    Decoding-related defaults are deliberately conservative: ``greedy`` stays
+    the default for byte-for-byte parity with the IEEE notebook. Switching to
+    beam at deploy time is a one-line YAML override:
+
+        serve:
+          decode_strategy: beam
+          beam_width: 4
+          length_penalty: 0.7
+          repetition_penalty: 1.1
+          no_repeat_ngram_size: 3
+    """
 
     max_upload_bytes: int = 10 * 1024 * 1024  # 10 MB
-    decode_strategy: str = "greedy"  # Phase 1b adds "beam"
+    decode_strategy: str = "greedy"
     beam_width: int = 3
+    length_penalty: float = 1.0
+    repetition_penalty: float = 1.0
+    no_repeat_ngram_size: int = 0
     cors_allowed_origins: list[str] = Field(default_factory=lambda: ["http://localhost:3000"])
 
+    @field_validator("decode_strategy")
+    @classmethod
+    def _validate_decode_strategy(cls, v: str) -> str:
+        if v not in {"greedy", "beam"}:
+            raise ValueError(f"decode_strategy must be 'greedy' or 'beam', got {v!r}")
+        return v
+
+    @field_validator("beam_width")
+    @classmethod
+    def _validate_beam_width(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError(f"beam_width must be >= 1, got {v}")
+        return v
+
+    @field_validator("repetition_penalty")
+    @classmethod
+    def _validate_repetition_penalty(cls, v: float) -> float:
+        if v < 1.0:
+            raise ValueError(f"repetition_penalty must be >= 1.0 (1.0 disables it), got {v}")
+        return v
+
 
 class AppConfig(BaseSettings):
     """Top-level config aggregating every sub-config.

src/captioning/evaluation/__init__.py

@@ -1,9 +1,49 @@
-"""Evaluation — caption-quality metrics.
+"""Evaluation — caption-quality metrics + per-sample diagnostics.
 
-Phase 1 ships a corpus-BLEU implementation only; Phase 1b expands to CIDEr,
-METEOR, and ROUGE-L (which is why this is its own package, not a single file).
+Available metrics (all corpus-level, 0-100 scale where applicable):
+    * BLEU-1..4 — :mod:`bleu`
+    * ROUGE-L   — :mod:`rouge`
+    * METEOR    — :mod:`meteor`  (requires a JRE on PATH)
+    * CIDEr     — :mod:`cider`   (requires >= 2 examples)
+
+:func:`compute_all_metrics` in :mod:`runner` is the single entry point used
+by the CLI and by future Phase 3 benchmark comparisons; per-sample
+diagnostics live in :mod:`inspection`.
 """
 
-from captioning.evaluation.bleu import corpus_bleu_score
+from captioning.evaluation.benchmark import RunMeta, write_run_artifacts
+from captioning.evaluation.bleu import (
+    BleuBreakdown,
+    corpus_bleu_breakdown,
+    corpus_bleu_score,
+)
+from captioning.evaluation.cider import MIN_SAMPLES_FOR_CIDER, corpus_cider_score
+from captioning.evaluation.inspection import (
+    SampleDiagnostics,
+    diagnose_many,
+    diagnose_sample,
+    format_diagnostic_row,
+    write_diagnostics_jsonl,
+)
+from captioning.evaluation.meteor import corpus_meteor_score
+from captioning.evaluation.rouge import corpus_rouge_l_score
+from captioning.evaluation.runner import MetricsReport, compute_all_metrics
 
-__all__ = ["corpus_bleu_score"]
+__all__ = [
+    "MIN_SAMPLES_FOR_CIDER",
+    "BleuBreakdown",
+    "MetricsReport",
+    "RunMeta",
+    "SampleDiagnostics",
+    "compute_all_metrics",
+    "corpus_bleu_breakdown",
+    "corpus_bleu_score",
+    "corpus_cider_score",
+    "corpus_meteor_score",
+    "corpus_rouge_l_score",
+    "diagnose_many",
+    "diagnose_sample",
+    "format_diagnostic_row",
+    "write_diagnostics_jsonl",
+    "write_run_artifacts",
+]

src/captioning/evaluation/benchmark.py

@@ -0,0 +1,128 @@
+"""Benchmark-ready run artefacts.
+
+Every evaluation pass writes a consistent set of files under
+``<run_root>/<run_id>/`` so Phase 3 cross-model comparisons can join them
+without bespoke parsing per model:
+
+    metrics.json            — :class:`MetricsReport` dumped via dataclass-asdict
+    predictions.jsonl       — one row per (image, prediction, references)
+    diagnostics.jsonl       — one :class:`SampleDiagnostics` per row
+    run_meta.json           — model id, decode strategy, n_samples, timestamp
+    report.md               — Markdown summary humans actually read
+
+A "run" is one (model, decode_strategy, dataset_slice) tuple. ``run_id`` is
+a free-form string — the CLI defaults to a timestamp; comparison code groups
+by ``model_id`` to plot bars across models.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+from captioning.evaluation.inspection import SampleDiagnostics, write_diagnostics_jsonl
+from captioning.evaluation.runner import MetricsReport
+
+
+@dataclass(frozen=True)
+class RunMeta:
+    """Per-evaluation-run metadata persisted next to metrics."""
+
+    model_id: str
+    decode_strategy: str
+    weights_path: str
+    tokenizer_dir: str
+    n_samples: int
+    max_length: int
+    beam_width: int | None = None
+    length_penalty: float | None = None
+    repetition_penalty: float | None = None
+    timestamp_utc: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "model_id": self.model_id,
+            "decode_strategy": self.decode_strategy,
+            "weights_path": self.weights_path,
+            "tokenizer_dir": self.tokenizer_dir,
+            "n_samples": self.n_samples,
+            "max_length": self.max_length,
+            "beam_width": self.beam_width,
+            "length_penalty": self.length_penalty,
+            "repetition_penalty": self.repetition_penalty,
+            "timestamp_utc": self.timestamp_utc,
+        }
+
+
+def write_run_artifacts(
+    run_dir: str | Path,
+    *,
+    metrics: MetricsReport,
+    meta: RunMeta,
+    images: list[str],
+    predictions: list[str],
+    references: list[list[str]],
+    diagnostics: list[SampleDiagnostics],
+) -> Path:
+    """Write every benchmark artefact to ``run_dir`` and return the directory.
+
+    Idempotent over a clean ``run_dir``; overwrites existing files inside.
+    """
+    out = Path(run_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    (out / "metrics.json").write_text(json.dumps(metrics.to_dict(), indent=2), encoding="utf-8")
+    (out / "run_meta.json").write_text(json.dumps(meta.to_dict(), indent=2), encoding="utf-8")
+
+    with (out / "predictions.jsonl").open("w", encoding="utf-8") as f:
+        for img, pred, refs in zip(images, predictions, references, strict=True):
+            row = {"image": img, "prediction": pred, "references": list(refs)}
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    write_diagnostics_jsonl(diagnostics, out / "diagnostics.jsonl")
+    (out / "report.md").write_text(_render_report_markdown(meta, metrics), encoding="utf-8")
+    return out
+
+
+def _render_report_markdown(meta: RunMeta, m: MetricsReport) -> str:
+    """Render the human-facing Markdown summary of a single run."""
+
+    def fmt(v: float | None) -> str:
+        return "n/a" if v is None else f"{v:.2f}"
+
+    lines = [
+        f"# Evaluation run — {meta.model_id}",
+        "",
+        f"- Decode strategy: `{meta.decode_strategy}`",
+        f"- Weights: `{meta.weights_path}`",
+        f"- Tokenizer dir: `{meta.tokenizer_dir}`",
+        f"- Samples: **{meta.n_samples}**",
+        f"- Timestamp (UTC): {meta.timestamp_utc}",
+    ]
+    if meta.beam_width is not None:
+        lines.append(f"- Beam width: {meta.beam_width}")
+    if meta.length_penalty is not None:
+        lines.append(f"- Length penalty: {meta.length_penalty}")
+    if meta.repetition_penalty is not None:
+        lines.append(f"- Repetition penalty: {meta.repetition_penalty}")
+    lines += [
+        "",
+        "## Metrics",
+        "",
+        "| Metric | Value |",
+        "|---|---|",
+        f"| BLEU-1 | {fmt(m.bleu1)} |",
+        f"| BLEU-2 | {fmt(m.bleu2)} |",
+        f"| BLEU-3 | {fmt(m.bleu3)} |",
+        f"| BLEU-4 | {fmt(m.bleu4)} |",
+        f"| ROUGE-L | {fmt(m.rouge_l)} |",
+        f"| METEOR | {fmt(m.meteor)} |",
+        f"| CIDEr | {fmt(m.cider)} |",
+    ]
+    if m.errors:
+        lines += ["", "## Skipped or failed metrics", ""]
+        for name, err in m.errors.items():
+            lines.append(f"- `{name}`: {err}")
+    return "\n".join(lines) + "\n"

src/captioning/evaluation/bleu.py

@@ -1,21 +1,40 @@
-"""Corpus BLEU score (Phase 1 minimal implementation).
-
-The IEEE paper reports BLEU ~24 on COCO val. The notebook does not include
-the evaluation code that produced this number — we add it here so the new
-modular pipeline can verify it matches the paper.
-
-Phase 1 ships *one* metric (corpus BLEU-4 via ``sacrebleu``) on purpose:
-    * sacrebleu is the de-facto BLEU implementation. NLTK's BLEU has
-      idiosyncratic smoothing and produces slightly different numbers; we
-      use sacrebleu so the published number is reproducible by anyone with
-      pip.
-    * Phase 1b expands to BLEU-1..4, CIDEr, METEOR, ROUGE-L, all in this
-      package, all behind the same ``runner.py`` interface.
+"""Corpus BLEU score.
+
+The IEEE paper reports BLEU-4 ~24 on COCO val. ``sacrebleu`` is the de-facto
+BLEU implementation; NLTK's BLEU has idiosyncratic smoothing and would not
+reproduce the published number across machines.
+
+``corpus_bleu_score`` returns BLEU-4 (the default n=4 score) so existing
+callers keep working. ``corpus_bleu_breakdown`` additionally exposes BLEU-1,
+BLEU-2, BLEU-3, BLEU-4 in one pass — useful for the inspection utility and
+for the JSON report consumed by Phase 3 cross-model comparison.
 """
 
 from __future__ import annotations
 
 from collections.abc import Sequence
+from dataclasses import dataclass
+
+from captioning.evaluation.tokenization import (
+    strip_sentinels_many,
+    strip_sentinels_references,
+)
+
+
+@dataclass(frozen=True)
+class BleuBreakdown:
+    """Per-n BLEU precisions plus the overall BLEU-4 score (0-100 scale)."""
+
+    bleu1: float
+    bleu2: float
+    bleu3: float
+    bleu4: float
+
+
+def _refs_by_slot(references: Sequence[Sequence[str]]) -> list[list[str]]:
+    """Convert ragged per-example references to sacrebleu's per-slot layout."""
+    max_refs = max(len(r) for r in references) if references else 0
+    return [[refs[i] if i < len(refs) else "" for refs in references] for i in range(max_refs)]
 
 
 def corpus_bleu_score(
@@ -38,6 +57,27 @@ def corpus_bleu_score(
         ImportError: If sacrebleu is not installed. Install via the eval
             extras: ``pip install -e ".[eval]"`` or the requirements file.
     """
+    return corpus_bleu_breakdown(predictions, references).bleu4
+
+
+def corpus_bleu_breakdown(
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+) -> BleuBreakdown:
+    """Compute BLEU-1, BLEU-2, BLEU-3, BLEU-4 in a single pass.
+
+    Args:
+        predictions: One generated caption per example.
+        references: One *list* of reference captions per example.
+
+    Returns:
+        :class:`BleuBreakdown` with all four cumulative BLEU-n scores on the
+        0-100 scale (sacrebleu's convention).
+
+    Raises:
+        ImportError: If sacrebleu is not installed.
+        ValueError: On mismatched lengths.
+    """
     try:
         import sacrebleu
     except ImportError as e:
@@ -52,12 +92,17 @@ def corpus_bleu_score(
             f"({len(references)}) must have the same length"
         )
 
-    # sacrebleu's `corpus_bleu` expects parallel lists, one *per reference
-    # slot*: refs_by_slot[slot_index][example_index].
-    max_refs = max(len(r) for r in references) if references else 0
-    refs_by_slot = [
-        [refs[i] if i < len(refs) else "" for refs in references] for i in range(max_refs)
-    ]
+    preds = strip_sentinels_many(predictions)
+    refs = strip_sentinels_references(references)
+    refs_by_slot = _refs_by_slot(refs)
 
-    bleu = sacrebleu.corpus_bleu(list(predictions), refs_by_slot)
-    return float(bleu.score)
+    # ``corpus_bleu`` only returns BLEU-4. To get cumulative BLEU-1..3 we
+    # instantiate ``BLEU`` directly with ``max_ngram_order=n``, which weights
+    # the geometric mean over precisions[:n] (same convention as NLTK's
+    # cumulative BLEU and the COCO eval scripts).
+    bleu_cls = sacrebleu.metrics.BLEU
+    scores: list[float] = []
+    for n in (1, 2, 3, 4):
+        scorer = bleu_cls(max_ngram_order=n, effective_order=True)
+        scores.append(float(scorer.corpus_score(preds, refs_by_slot).score))
+    return BleuBreakdown(bleu1=scores[0], bleu2=scores[1], bleu3=scores[2], bleu4=scores[3])

src/captioning/evaluation/cider.py

@@ -0,0 +1,81 @@
+"""CIDEr (Consensus-based Image Description Evaluation) corpus metric.
+
+CIDEr is the metric the COCO captioning leaderboard ranks by. It computes a
+TF-IDF weighting over n-grams of the references and measures cosine similarity
+to the prediction. Higher is better; correctly trained models score in the
+range 0.6 - 1.4 on COCO val.
+
+Implementation notes:
+    * We delegate to ``pycocoevalcap`` — the reference implementation used by
+      the original CIDEr paper and by every COCO submission.
+    * CIDEr's TF-IDF is corpus-level: scoring a *single* example returns 0
+      because every n-gram is "common" to that one-document corpus. The
+      ``runner`` aggregator and the CLI guard against calling CIDEr with
+      fewer than ``MIN_SAMPLES_FOR_CIDER`` examples and return ``None``
+      in that case.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from captioning.evaluation.tokenization import (
+    strip_sentinels_many,
+    strip_sentinels_references,
+)
+
+# CIDEr's TF-IDF is degenerate below this — every n-gram is "common"
+# to the entire corpus, so the score collapses to 0. We surface ``None``
+# instead of a misleading value below this threshold.
+MIN_SAMPLES_FOR_CIDER = 2
+
+
+def corpus_cider_score(
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+) -> float:
+    """Compute corpus CIDEr.
+
+    Args:
+        predictions: One generated caption per example.
+        references: One *list* of reference captions per example.
+
+    Returns:
+        CIDEr in the 0-10 range (pycocoevalcap convention; the typical COCO
+        leaderboard value is in [0, 2]).
+
+    Raises:
+        ImportError: If ``pycocoevalcap`` is not installed.
+        ValueError: On mismatched lengths or if called with fewer than
+            ``MIN_SAMPLES_FOR_CIDER`` examples (in which case CIDEr's TF-IDF
+            is degenerate and the score is meaningless).
+    """
+    if len(predictions) != len(references):
+        raise ValueError(
+            f"predictions ({len(predictions)}) and references "
+            f"({len(references)}) must have the same length"
+        )
+    if len(predictions) < MIN_SAMPLES_FOR_CIDER:
+        raise ValueError(
+            f"CIDEr requires at least {MIN_SAMPLES_FOR_CIDER} examples; "
+            f"got {len(predictions)}. TF-IDF is degenerate on smaller corpora."
+        )
+
+    try:
+        from pycocoevalcap.cider.cider import Cider
+    except ImportError as e:
+        raise ImportError(
+            "pycocoevalcap is required for CIDEr evaluation. "
+            "Install via `pip install -r requirements-eval.txt`."
+        ) from e
+
+    preds = strip_sentinels_many(predictions)
+    refs = strip_sentinels_references(references)
+
+    # pycocoevalcap expects {image_id: [captions]} dicts.
+    gts = {str(i): [r for r in ref_list if r] for i, ref_list in enumerate(refs)}
+    res = {str(i): [p] for i, p in enumerate(preds)}
+
+    scorer = Cider()
+    score, _ = scorer.compute_score(gts, res)
+    return float(score)

src/captioning/evaluation/inspection.py

@@ -0,0 +1,163 @@
+"""Per-sample inspection utilities for diagnosing weak captions.
+
+The aggregate corpus metric tells you *how bad* the model is; this module
+tells you *why*. For each (image, prediction, reference-set) triple it
+records per-sample BLEU-4, sentence-level ROUGE-L, the prediction length,
+the longest repeated token run, and whether the prediction is empty after
+stripping sentinels.
+
+Three failure modes the evaluation pass is trying to surface:
+    * **Generic captions** — high BLEU-1, low BLEU-4 (n-gram trickle out).
+    * **Repetition** — large ``repeat_run`` value.
+    * **Early stopping** — ``length_tokens`` far below reference median.
+
+Output JSONL is intentionally flat (one line per sample) so it can be loaded
+with ``pandas.read_json(..., lines=True)`` or grep'd from the shell. The
+runner that uses this module writes one such file per evaluation pass
+alongside ``metrics.json`` for the same run.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Sequence
+from dataclasses import asdict, dataclass
+from itertools import pairwise
+from pathlib import Path
+
+from captioning.evaluation.tokenization import strip_sentinels
+
+
+@dataclass(frozen=True)
+class SampleDiagnostics:
+    """Inspectable record for one (image, prediction, reference-set) triple."""
+
+    image: str
+    prediction: str
+    references: list[str]
+    length_tokens: int
+    longest_repeat_run: int
+    sentence_bleu4: float | None
+    sentence_rouge_l: float | None
+    flags: list[str]
+
+
+def _longest_repeat_run(tokens: Sequence[str]) -> int:
+    """Return the longest run of immediately-repeated tokens.
+
+    Example: ``["a", "a", "a", "dog"]`` -> ``3``. Used to flag the classic
+    transformer-decoder collapse where the same token is emitted on every step.
+    """
+    if not tokens:
+        return 0
+    best = current = 1
+    for prev, cur in pairwise(tokens):
+        current = current + 1 if cur == prev else 1
+        best = max(best, current)
+    return best
+
+
+def _sentence_bleu4(prediction: str, references: Sequence[str]) -> float | None:
+    """Sentence-level BLEU-4 via sacrebleu's effective-order smoothing."""
+    try:
+        import sacrebleu
+    except ImportError:
+        return None
+    if not references or not prediction:
+        return None
+    scorer = sacrebleu.metrics.BLEU(effective_order=True, max_ngram_order=4)
+    return float(scorer.sentence_score(prediction, list(references)).score)
+
+
+def _sentence_rouge_l(prediction: str, references: Sequence[str]) -> float | None:
+    """Best-of-references sentence-level ROUGE-L F-measure (0-100 scale)."""
+    try:
+        from rouge_score import rouge_scorer
+    except ImportError:
+        return None
+    valid_refs = [r for r in references if r]
+    if not valid_refs or not prediction:
+        return None
+    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
+    best = max(scorer.score(r, prediction)["rougeL"].fmeasure for r in valid_refs)
+    return float(100.0 * best)
+
+
+def diagnose_sample(
+    image: str,
+    prediction: str,
+    references: Sequence[str],
+) -> SampleDiagnostics:
+    """Return :class:`SampleDiagnostics` for one prediction-vs-references row."""
+    pred_clean = strip_sentinels(prediction)
+    ref_clean = [strip_sentinels(r) for r in references if r]
+    tokens = pred_clean.split()
+
+    flags: list[str] = []
+    if not pred_clean:
+        flags.append("empty")
+    if len(tokens) <= 2:
+        flags.append("very_short")
+    repeat = _longest_repeat_run(tokens)
+    if repeat >= 3:
+        flags.append("repetitive")
+    if ref_clean and tokens and len(tokens) < min(len(r.split()) for r in ref_clean) // 2:
+        flags.append("under_length")
+
+    return SampleDiagnostics(
+        image=image,
+        prediction=pred_clean,
+        references=ref_clean,
+        length_tokens=len(tokens),
+        longest_repeat_run=repeat,
+        sentence_bleu4=_sentence_bleu4(pred_clean, ref_clean),
+        sentence_rouge_l=_sentence_rouge_l(pred_clean, ref_clean),
+        flags=flags,
+    )
+
+
+def diagnose_many(
+    images: Sequence[str],
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+) -> list[SampleDiagnostics]:
+    """Vectorised :func:`diagnose_sample` over parallel sequences."""
+    if not (len(images) == len(predictions) == len(references)):
+        raise ValueError(
+            "images, predictions, references must be the same length: "
+            f"got {len(images)} / {len(predictions)} / {len(references)}"
+        )
+    return [
+        diagnose_sample(img, pred, refs)
+        for img, pred, refs in zip(images, predictions, references, strict=True)
+    ]
+
+
+def write_diagnostics_jsonl(
+    diagnostics: Iterable[SampleDiagnostics],
+    path: str | Path,
+) -> None:
+    """Write one JSON object per line — pandas/jq friendly.
+
+    Args:
+        diagnostics: An iterable of :class:`SampleDiagnostics` (typically the
+            output of :func:`diagnose_many`).
+        path: Destination file. Parent directory is created if needed.
+    """
+    out = Path(path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with out.open("w", encoding="utf-8") as f:
+        for d in diagnostics:
+            f.write(json.dumps(asdict(d), ensure_ascii=False) + "\n")
+
+
+def format_diagnostic_row(d: SampleDiagnostics) -> str:
+    """Return a one-line human-readable summary — used by the CLI tail print."""
+    bleu = f"BLEU4={d.sentence_bleu4:5.1f}" if d.sentence_bleu4 is not None else "BLEU4=    n/a"
+    rouge = f"R-L={d.sentence_rouge_l:5.1f}" if d.sentence_rouge_l is not None else "R-L=    n/a"
+    flagstr = ",".join(d.flags) if d.flags else "-"
+    return (
+        f"{Path(d.image).name:35s}  "
+        f"{bleu}  {rouge}  len={d.length_tokens:>2}  repeat={d.longest_repeat_run:>2}  "
+        f"flags={flagstr}\n  pred: {d.prediction}\n  ref : {d.references[0] if d.references else ''}"
+    )

src/captioning/evaluation/meteor.py

@@ -0,0 +1,77 @@
+"""METEOR (Metric for Evaluation of Translation with Explicit Ordering).
+
+METEOR is part of the standard COCO captioning report alongside BLEU, ROUGE-L,
+and CIDEr. It complements BLEU by rewarding semantic matches (synonyms,
+stems) rather than only surface n-gram overlap.
+
+Implementation notes:
+    * We use the ``pycocoevalcap`` METEOR adapter, which shells out to the
+      original Java implementation. METEOR therefore needs a JRE on PATH at
+      runtime; the import succeeds either way, the Java process is spawned
+      lazily on first scoring call.
+    * METEOR's process is long-lived and accepts batches over stdin/stdout —
+      a single ``compute_score`` call handles the whole corpus in one round
+      trip, so this scales to thousands of examples without thrashing the JVM.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from captioning.evaluation.tokenization import (
+    strip_sentinels_many,
+    strip_sentinels_references,
+)
+
+
+def corpus_meteor_score(
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+) -> float:
+    """Compute corpus METEOR via ``pycocoevalcap``.
+
+    Args:
+        predictions: One generated caption per example.
+        references: One *list* of reference captions per example.
+
+    Returns:
+        Corpus METEOR in the 0-100 range to match the rest of this package.
+        pycocoevalcap returns 0-1; we multiply by 100 for report parity.
+
+    Raises:
+        ImportError: If ``pycocoevalcap`` is not installed.
+        ValueError: On mismatched lengths.
+        RuntimeError: If the Java METEOR process cannot be launched.
+    """
+    if len(predictions) != len(references):
+        raise ValueError(
+            f"predictions ({len(predictions)}) and references "
+            f"({len(references)}) must have the same length"
+        )
+    if not predictions:
+        return 0.0
+
+    try:
+        from pycocoevalcap.meteor.meteor import Meteor
+    except ImportError as e:
+        raise ImportError(
+            "pycocoevalcap is required for METEOR evaluation. "
+            "Install via `pip install -r requirements-eval.txt`."
+        ) from e
+
+    preds = strip_sentinels_many(predictions)
+    refs = strip_sentinels_references(references)
+
+    gts = {str(i): [r for r in ref_list if r] for i, ref_list in enumerate(refs)}
+    res = {str(i): [p] for i, p in enumerate(preds)}
+
+    scorer = Meteor()
+    try:
+        score, _ = scorer.compute_score(gts, res)
+    except Exception as e:  # — meteor.py raises bare Exceptions
+        raise RuntimeError(
+            "METEOR scoring failed. METEOR requires a Java runtime on PATH. "
+            f"Underlying error: {e}"
+        ) from e
+
+    return float(100.0 * score)

src/captioning/evaluation/rouge.py

@@ -0,0 +1,75 @@
+"""Corpus ROUGE-L for caption evaluation.
+
+ROUGE-L measures the longest common subsequence between a prediction and its
+references and is part of the standard COCO captioning report (BLEU, METEOR,
+ROUGE-L, CIDEr).
+
+Implementation notes:
+    * We use Google's ``rouge_score`` package (the canonical implementation
+      since the original perl scripts were retired). It returns precision /
+      recall / fmeasure per (prediction, reference) pair.
+    * COCO captions ship up to 5 references per image. We take the maximum
+      F-measure across references — same convention as pycocoevalcap.
+    * The corpus score is the mean of per-sample F-measures, matching how
+      sacrebleu and pycocoevalcap aggregate metrics over a dataset.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from captioning.evaluation.tokenization import (
+    strip_sentinels_many,
+    strip_sentinels_references,
+)
+
+
+def corpus_rouge_l_score(
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+) -> float:
+    """Compute corpus ROUGE-L F-measure.
+
+    Args:
+        predictions: One generated caption per example.
+        references: One *list* of reference captions per example.
+
+    Returns:
+        Mean ROUGE-L F-measure across examples, in the 0-100 range to match
+        sacrebleu's convention (so the report shows BLEU/ROUGE/METEOR/CIDEr
+        on comparable scales).
+
+    Raises:
+        ImportError: If ``rouge_score`` is not installed
+            (``pip install -r requirements-eval.txt``).
+        ValueError: On mismatched lengths or an empty references slot.
+    """
+    if len(predictions) != len(references):
+        raise ValueError(
+            f"predictions ({len(predictions)}) and references "
+            f"({len(references)}) must have the same length"
+        )
+    if not predictions:
+        return 0.0
+
+    try:
+        from rouge_score import rouge_scorer
+    except ImportError as e:
+        raise ImportError(
+            "rouge_score is required for ROUGE-L evaluation. "
+            "Install via `pip install -r requirements-eval.txt`."
+        ) from e
+
+    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
+    preds = strip_sentinels_many(predictions)
+    refs = strip_sentinels_references(references)
+
+    total = 0.0
+    for hypothesis, ref_list in zip(preds, refs, strict=True):
+        valid_refs = [r for r in ref_list if r]
+        if not valid_refs or not hypothesis:
+            continue
+        best = max(scorer.score(r, hypothesis)["rougeL"].fmeasure for r in valid_refs)
+        total += best
+
+    return float(100.0 * total / len(preds))

src/captioning/evaluation/runner.py

@@ -0,0 +1,128 @@
+"""Single entry point that returns every implemented caption-quality metric.
+
+``compute_all_metrics`` is the shared aggregator used by the CLI
+(:mod:`scripts.evaluate`) and the per-sample inspection utility. It produces
+a single :class:`MetricsReport` so downstream code never has to know which
+metrics exist in the package — only how to read fields off the dataclass.
+
+Adding a new metric is the four-step pattern this package already follows
+elsewhere:
+    1. Implement ``corpus_<metric>_score`` in a sibling module.
+    2. Add an entry to :class:`MetricsReport`.
+    3. Call it from :func:`compute_all_metrics` (wrapped in a try/except so a
+       single broken metric never poisons the whole report).
+    4. Add a unit test on a toy fixture.
+
+The exception swallowing is deliberate — METEOR needs Java, CIDEr needs
+multiple samples, sacrebleu is always available. We do NOT want one
+unavailable metric to kill the entire evaluation pass; instead we record
+``None`` for that metric and surface a per-metric ``errors`` field so callers
+(and the CLI) can flag the issue without losing the metrics that did work.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import asdict, dataclass, field
+
+from captioning.evaluation.bleu import corpus_bleu_breakdown
+from captioning.evaluation.cider import MIN_SAMPLES_FOR_CIDER, corpus_cider_score
+from captioning.evaluation.meteor import corpus_meteor_score
+from captioning.evaluation.rouge import corpus_rouge_l_score
+
+
+@dataclass(frozen=True)
+class MetricsReport:
+    """Aggregate metric snapshot for one evaluation pass.
+
+    Every metric is ``float | None`` — ``None`` means the metric was skipped
+    (uninstalled, environment missing Java, too few samples for CIDEr, ...).
+    The reason for skipping is in :attr:`errors` keyed by metric name.
+    """
+
+    n_examples: int
+    bleu1: float | None = None
+    bleu2: float | None = None
+    bleu3: float | None = None
+    bleu4: float | None = None
+    rouge_l: float | None = None
+    meteor: float | None = None
+    cider: float | None = None
+    errors: dict[str, str] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, object]:
+        """Return a JSON-serialisable dict (``errors`` becomes a sub-object)."""
+        return asdict(self)
+
+
+def compute_all_metrics(
+    predictions: Sequence[str],
+    references: Sequence[Sequence[str]],
+    *,
+    include_meteor: bool = True,
+    include_cider: bool = True,
+) -> MetricsReport:
+    """Compute every available metric on a single ``(preds, refs)`` corpus.
+
+    Args:
+        predictions: One generated caption per example.
+        references: One *list* of reference captions per example.
+        include_meteor: Set False to skip METEOR (avoids the JVM spawn —
+            helpful in CI where Java isn't installed).
+        include_cider: Set False to skip CIDEr (avoids the warning when
+            running on tiny corpora; the runner also auto-skips below
+            ``MIN_SAMPLES_FOR_CIDER``).
+
+    Returns:
+        A :class:`MetricsReport` with every field populated by a corpus
+        metric or recorded as failed in ``errors``.
+    """
+    if len(predictions) != len(references):
+        raise ValueError(
+            f"predictions ({len(predictions)}) and references "
+            f"({len(references)}) must have the same length"
+        )
+
+    errors: dict[str, str] = {}
+    bleu1 = bleu2 = bleu3 = bleu4 = None
+    rouge_l = meteor = cider = None
+
+    try:
+        bleu = corpus_bleu_breakdown(predictions, references)
+        bleu1, bleu2, bleu3, bleu4 = bleu.bleu1, bleu.bleu2, bleu.bleu3, bleu.bleu4
+    except Exception as e:  # — surface, don't crash the run
+        errors["bleu"] = repr(e)
+
+    try:
+        rouge_l = corpus_rouge_l_score(predictions, references)
+    except Exception as e:
+        errors["rouge_l"] = repr(e)
+
+    if include_meteor:
+        try:
+            meteor = corpus_meteor_score(predictions, references)
+        except Exception as e:
+            errors["meteor"] = repr(e)
+
+    if include_cider:
+        if len(predictions) < MIN_SAMPLES_FOR_CIDER:
+            errors["cider"] = (
+                f"skipped: needs >= {MIN_SAMPLES_FOR_CIDER} examples, " f"got {len(predictions)}"
+            )
+        else:
+            try:
+                cider = corpus_cider_score(predictions, references)
+            except Exception as e:
+                errors["cider"] = repr(e)
+
+    return MetricsReport(
+        n_examples=len(predictions),
+        bleu1=bleu1,
+        bleu2=bleu2,
+        bleu3=bleu3,
+        bleu4=bleu4,
+        rouge_l=rouge_l,
+        meteor=meteor,
+        cider=cider,
+        errors=errors,
+    )

src/captioning/evaluation/tokenization.py

@@ -0,0 +1,44 @@
+"""Reference/hypothesis tokenisation helpers shared across metric modules.
+
+CIDEr, METEOR (via pycocoevalcap) and ROUGE-L all expect *string* inputs but
+each implements its own tokenisation internally. To compare metrics on the
+same footing we strip our sentinel tokens once, up front, before any metric
+sees a caption.
+
+These helpers exist because every metric module would otherwise re-implement
+the same ``[start]`` / ``[end]`` stripping inline — and bugs would diverge.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+
+from captioning.preprocessing.caption import END_TOKEN, START_TOKEN
+
+
+def strip_sentinels(caption: str) -> str:
+    """Remove ``[start]`` / ``[end]`` sentinels and collapse whitespace.
+
+    Args:
+        caption: A caption string that may carry our training sentinels.
+
+    Returns:
+        The same caption with the sentinels removed and consecutive whitespace
+        collapsed to a single space. Empty input returns ``""``.
+    """
+    if not caption:
+        return ""
+    cleaned = caption.replace(START_TOKEN, " ").replace(END_TOKEN, " ")
+    return " ".join(cleaned.split())
+
+
+def strip_sentinels_many(captions: Iterable[str]) -> list[str]:
+    """Apply :func:`strip_sentinels` to every caption in ``captions``."""
+    return [strip_sentinels(c) for c in captions]
+
+
+def strip_sentinels_references(
+    references: Sequence[Sequence[str]],
+) -> list[list[str]]:
+    """Apply :func:`strip_sentinels` to every reference list."""
+    return [strip_sentinels_many(refs) for refs in references]

src/captioning/inference/__init__.py

@@ -10,12 +10,14 @@ explicitly so it works inside a long-lived process (FastAPI lifespan).
     predictor.py      ``CaptionPredictor`` — singleton wrapper for the API
 """
 
+from captioning.inference.beam import generate_caption_beam
 from captioning.inference.greedy import generate_caption_greedy
 from captioning.inference.image_loader import load_image_from_path
 from captioning.inference.predictor import CaptionPredictor
 
 __all__ = [
     "CaptionPredictor",
+    "generate_caption_beam",
     "generate_caption_greedy",
     "load_image_from_path",
 ]

src/captioning/inference/beam.py

@@ -0,0 +1,240 @@
+"""Beam-search caption generation.
+
+Greedy decoding (the only Phase 1 option) routinely produces generic captions
+because the model's most-likely-next-token at every step rarely lines up with
+the most-likely-*sequence*. Beam search explores multiple partial captions in
+parallel and ranks them by total log-probability, lifting BLEU-4 by 2-5
+points on most transformer captioners without retraining.
+
+Algorithm (standard beam search with length and repetition controls):
+    * Maintain ``beam_width`` active beams, each a (token-id sequence, score).
+    * At each step, batch every active beam through the decoder once, take the
+      log-softmax at the current position, apply the repetition penalty and
+      the optional no-repeat-ngram block, and pick the global top-K
+      candidates across (beam, vocab) pairs.
+    * Beams that emit ``[end]`` move into the finished list (their score is
+      already final at that point); the search ends when ``beam_width`` beams
+      have finished or we hit the max-length budget.
+    * Final ranking divides each finished beam's score by
+      ``len(seq) ** length_penalty`` so the search isn't biased toward very
+      short sequences (the classic length problem in beam search).
+
+This implementation is intentionally kept *callable* — the same predictor
+class dispatches between :func:`generate_caption_greedy` and this one based
+on ``decode_strategy``. Phase 3 model wrappers (BLIP, ViT-GPT2) can reuse
+the same dispatcher.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+from captioning.preprocessing.caption import END_TOKEN, START_TOKEN
+from captioning.preprocessing.tokenizer import CaptionTokenizer
+
+_LOG_EPSILON = 1e-12
+
+
+@dataclass
+class _Beam:
+    """One partial caption under exploration."""
+
+    token_ids: list[int]
+    score: float
+    finished: bool = False
+    history: set[int] = field(default_factory=set)
+
+    def length(self) -> int:
+        """Number of generated tokens (excludes the seed [start] token)."""
+        return max(len(self.token_ids) - 1, 1)
+
+
+def _apply_repetition_penalty(
+    log_probs,
+    history_ids: set[int],
+    penalty: float,
+):
+    """Subtract ``log(penalty)`` from already-seen tokens' log-probabilities.
+
+    HuggingFace's repetition_penalty (Keskar et al. 2019) divides logits by
+    ``penalty`` (>1) for tokens already in the context. We work with log-
+    probabilities here, so the equivalent operation is to *subtract*
+    ``log(penalty)`` for positive log-probabilities and add it for negative
+    ones — but log-probabilities are always non-positive, so we always make
+    seen tokens less likely. That is the correct direction (we want to
+    discourage repetition).
+    """
+    if penalty <= 1.0 or not history_ids:
+        return log_probs
+    log_pen = math.log(penalty)
+    for tid in history_ids:
+        if 0 <= tid < log_probs.shape[-1]:
+            log_probs[tid] -= log_pen
+    return log_probs
+
+
+def _blocks_repeat_ngram(seq: list[int], candidate: int, n: int) -> bool:
+    """Return True if appending ``candidate`` would repeat an n-gram in ``seq``."""
+    if n <= 0 or len(seq) < n - 1:
+        return False
+    tail = tuple(seq[-(n - 1) :] + [candidate]) if n > 1 else (candidate,)
+    return any(tuple(seq[i : i + n]) == tail for i in range(len(seq) - n + 1))
+
+
+def generate_caption_beam(  # — beam search has many knobs by nature
+    model,
+    tokenizer: CaptionTokenizer,
+    image_tensor,
+    max_length: int,
+    *,
+    beam_width: int = 3,
+    length_penalty: float = 1.0,
+    repetition_penalty: float = 1.0,
+    no_repeat_ngram_size: int = 0,
+) -> str:
+    """Generate a caption using beam search with optional length / repetition control.
+
+    Args:
+        model: An ``ImageCaptioningModel`` whose weights have been loaded.
+        tokenizer: Fitted :class:`CaptionTokenizer`.
+        image_tensor: ``[299, 299, 3]`` float tensor as produced by
+            ``inference.load_image_from_path``.
+        max_length: Same budget as greedy (``config.model.max_length``); the
+            search stops at the first of (all beams finished, length exhausted).
+        beam_width: Number of parallel hypotheses. ``1`` reduces to greedy.
+        length_penalty: GNMT-style penalty exponent. ``score / len ** alpha``.
+            ``0.0`` disables it; ``0.6-1.0`` is the common range. Higher values
+            favour longer captions.
+        repetition_penalty: HuggingFace's CTRL-style penalty. ``1.0`` disables
+            it; ``>1.0`` penalises tokens already in the partial caption.
+        no_repeat_ngram_size: If ``> 0``, forbids emitting any token that
+            would complete an n-gram already present in the partial caption.
+            ``3`` is a common choice for captioning.
+
+    Returns:
+        The best-scoring caption (sentinels stripped, same convention as
+        :func:`generate_caption_greedy`).
+    """
+    import numpy as np
+    import tensorflow as tf
+
+    # 1. Encode the image once. Beams share the encoded features.
+    img = tf.expand_dims(image_tensor, axis=0)
+    img_embed = model.cnn_model(img)
+    img_encoded = model.encoder(img_embed, training=False)
+
+    start_id = tokenizer.word_to_id(START_TOKEN)
+    end_id = tokenizer.word_to_id(END_TOKEN)
+
+    # 2. Initialise a single seed beam containing only the [start] token.
+    beams: list[_Beam] = [_Beam(token_ids=[start_id], score=0.0, history={start_id})]
+    finished: list[_Beam] = []
+
+    decode_steps = max_length - 1  # decoder is fed sequences of length max_length-1
+
+    for step in range(decode_steps):
+        if not beams:
+            break
+
+        # 3. Batch every active beam into a single decoder forward pass.
+        token_batch = np.zeros((len(beams), decode_steps), dtype=np.int64)
+        for i, beam in enumerate(beams):
+            seq = beam.token_ids[:decode_steps]
+            token_batch[i, : len(seq)] = seq
+
+        token_tensor = tf.convert_to_tensor(token_batch)
+        mask = tf.cast(token_tensor != 0, tf.int32)
+        # Encoded features must be broadcast to match the beam batch dimension.
+        encoded_batch = tf.repeat(img_encoded, repeats=len(beams), axis=0)
+        preds = model.decoder(token_tensor, encoded_batch, training=False, mask=mask)
+        # preds is [B, T, V]; we read position `step` for each beam.
+        step_probs = preds.numpy()[:, step, :]
+        step_log_probs = np.log(step_probs + _LOG_EPSILON)
+
+        # 4. Expand every beam, then keep the global top-K.
+        candidates: list[_Beam] = []
+        vocab_size = step_log_probs.shape[-1]
+        for i, beam in enumerate(beams):
+            lp = step_log_probs[i].copy()
+            lp = _apply_repetition_penalty(lp, beam.history, repetition_penalty)
+
+            # Pick a wider candidate pool than beam_width per beam — when most
+            # beams want the same token, expansion needs slack to remain diverse.
+            pool = min(beam_width * 2, vocab_size)
+            top_ids = np.argpartition(-lp, pool - 1)[:pool]
+            top_ids = top_ids[np.argsort(-lp[top_ids])]
+
+            for tid in top_ids:
+                tid_int = int(tid)
+                if no_repeat_ngram_size > 0 and _blocks_repeat_ngram(
+                    beam.token_ids, tid_int, no_repeat_ngram_size
+                ):
+                    continue
+                new_seq = [*beam.token_ids, tid_int]
+                new_score = beam.score + float(lp[tid_int])
+                new_history = beam.history | {tid_int}
+                candidates.append(
+                    _Beam(
+                        token_ids=new_seq,
+                        score=new_score,
+                        finished=(tid_int == end_id),
+                        history=new_history,
+                    )
+                )
+
+        # 5. Sort candidates by score and keep the top ``beam_width`` actives.
+        candidates.sort(key=lambda b: b.score, reverse=True)
+        next_beams: list[_Beam] = []
+        for cand in candidates:
+            if cand.finished:
+                finished.append(cand)
+                continue
+            next_beams.append(cand)
+            if len(next_beams) >= beam_width:
+                break
+        beams = next_beams
+
+        # 6. Early termination — we already have enough finished beams and
+        # none of the active ones can beat the best finished score (their
+        # best-case future log-prob is 0, so length-normalised score won't
+        # beat the current top).
+        if len(finished) >= beam_width and beams:
+            best_finished = max(_length_normalised(b, length_penalty) for b in finished)
+            best_active_upper_bound = max(_length_normalised(b, length_penalty) for b in beams)
+            if best_active_upper_bound <= best_finished:
+                break
+
+    # 7. Anything still active at the budget cap counts as finished.
+    finished.extend(beams)
+    if not finished:
+        return ""
+
+    finished.sort(key=lambda b: _length_normalised(b, length_penalty), reverse=True)
+    best = finished[0]
+    return _detokenize(best.token_ids, tokenizer, end_id)
+
+
+def _length_normalised(beam: _Beam, alpha: float) -> float:
+    """Apply length penalty to a beam score (higher == better)."""
+    if alpha == 0.0:
+        return beam.score
+    return beam.score / (beam.length() ** alpha)
+
+
+def _detokenize(
+    token_ids: list[int],
+    tokenizer: CaptionTokenizer,
+    end_id: int,
+) -> str:
+    """Convert beam token ids back to a clean caption string."""
+    words: list[str] = []
+    for tid in token_ids:
+        if tid == end_id:
+            break
+        word = tokenizer.decode_id(tid)
+        # Skip [start], padding, and OOV ids that decode to empty strings.
+        if word in {"", START_TOKEN, END_TOKEN, "[UNK]"}:
+            continue
+        words.append(word)
+    return " ".join(words)

src/captioning/inference/predictor.py

@@ -5,8 +5,9 @@ Why a class around the existing functions:
       model across every request. A predictor object is the natural home for
       "loaded model + loaded tokenizer + decoded config".
     * Tests can construct one with stub objects without monkey-patching globals.
-    * Phase 1b adds beam search; Phase 3 adds a model registry. Both extend
-      this class, not the functional callsites.
+    * Multiple decode strategies (greedy, beam) live behind the same
+      ``predict_tensor`` / ``predict_path`` API — callers do not need to know
+      which one is active.
 
 Construction is *not* the same as readiness: ``CaptionPredictor.warmup()``
 runs one inference on a dummy tensor so the first real request doesn't pay
@@ -19,6 +20,7 @@ from pathlib import Path
 from typing import Literal
 
 from captioning.config.schema import AppConfig
+from captioning.inference.beam import generate_caption_beam
 from captioning.inference.greedy import generate_caption_greedy
 from captioning.inference.image_loader import load_image_from_path
 from captioning.preprocessing.tokenizer import CaptionTokenizer
@@ -26,6 +28,8 @@ from captioning.utils.logging import get_logger
 
 log = get_logger(__name__)
 
+DecodeStrategy = Literal["greedy", "beam"]
+
 
 class CaptionPredictor:
     """Thin wrapper exposing ``predict_path`` / ``predict_tensor`` / ``warmup``."""
@@ -36,24 +40,41 @@ class CaptionPredictor:
         tokenizer: CaptionTokenizer,
         config: AppConfig,
         *,
-        decode_strategy: Literal["greedy"] = "greedy",
+        decode_strategy: DecodeStrategy = "greedy",
+        beam_width: int = 3,
+        length_penalty: float = 1.0,
+        repetition_penalty: float = 1.0,
+        no_repeat_ngram_size: int = 0,
     ) -> None:
         """Args:
         model: Loaded ``ImageCaptioningModel``. Caller is responsible for
             having called ``model.load_weights(...)`` already.
         tokenizer: Fitted ``CaptionTokenizer``.
         config: Validated ``AppConfig`` — ``model.max_length`` is consumed.
-        decode_strategy: Phase 1 supports only ``"greedy"``. Phase 1b adds
-            ``"beam"``; this argument is here so the signature is stable.
+        decode_strategy: ``"greedy"`` (argmax per step, byte-for-byte parity
+            with the IEEE notebook) or ``"beam"`` (beam search with length
+            and repetition controls).
+        beam_width: Beam width when ``decode_strategy == "beam"``. Ignored
+            for greedy.
+        length_penalty: GNMT length penalty; ``0.0`` disables, ``0.6-1.0`` is
+            the common range.
+        repetition_penalty: HF-style multiplicative penalty on already-seen
+            tokens; ``1.0`` disables.
+        no_repeat_ngram_size: If > 0, blocks any token that would repeat an
+            n-gram already in the partial caption.
         """
-        if decode_strategy != "greedy":
-            raise NotImplementedError(
-                f"Phase 1 supports decode_strategy='greedy' only, got {decode_strategy!r}"
-            )
+        if decode_strategy not in {"greedy", "beam"}:
+            raise ValueError(f"decode_strategy must be 'greedy' or 'beam', got {decode_strategy!r}")
+        if beam_width < 1:
+            raise ValueError(f"beam_width must be >= 1, got {beam_width}")
         self.model = model
         self.tokenizer = tokenizer
         self.config = config
-        self.decode_strategy = decode_strategy
+        self.decode_strategy: DecodeStrategy = decode_strategy
+        self.beam_width = beam_width
+        self.length_penalty = length_penalty
+        self.repetition_penalty = repetition_penalty
+        self.no_repeat_ngram_size = no_repeat_ngram_size
 
     @classmethod
     def from_artifacts(
@@ -61,17 +82,18 @@ class CaptionPredictor:
         weights_path: str | Path,
         tokenizer_dir: str | Path,
         config: AppConfig,
+        *,
+        decode_strategy: DecodeStrategy | None = None,
+        beam_width: int | None = None,
+        length_penalty: float | None = None,
+        repetition_penalty: float | None = None,
+        no_repeat_ngram_size: int | None = None,
     ) -> CaptionPredictor:
         """Load weights and tokenizer from disk and return a ready predictor.
 
-        Args:
-            weights_path: Path to ``model.h5`` (notebook cell 30 saved this).
-            tokenizer_dir: Directory containing ``vocab.pkl`` (and ``vocab.json``).
-            config: Validated ``AppConfig``. ``model.max_length`` and
-                ``model.vocabulary_size`` must match the trained weights.
-
-        Returns:
-            A ``CaptionPredictor`` ready for inference.
+        Decoding knobs fall back to :class:`ServeConfig` defaults when not
+        passed explicitly — keeping CLI flags overridable while still letting
+        deploy-time YAML drive the production behaviour.
         """
         from captioning.models.factory import build_caption_model
 
@@ -86,19 +108,56 @@ class CaptionPredictor:
         cls._dummy_pass(model, config)
         model.load_weights(str(weights_path))
 
-        log.info("predictor_loaded", weights=str(weights_path))
-        return cls(model=model, tokenizer=tokenizer, config=config)
+        resolved_strategy: DecodeStrategy = (
+            decode_strategy or config.serve.decode_strategy  # type: ignore[assignment]
+        )
+        log.info(
+            "predictor_loaded",
+            weights=str(weights_path),
+            decode_strategy=resolved_strategy,
+        )
+        return cls(
+            model=model,
+            tokenizer=tokenizer,
+            config=config,
+            decode_strategy=resolved_strategy,
+            beam_width=beam_width if beam_width is not None else config.serve.beam_width,
+            length_penalty=(
+                length_penalty if length_penalty is not None else config.serve.length_penalty
+            ),
+            repetition_penalty=(
+                repetition_penalty
+                if repetition_penalty is not None
+                else config.serve.repetition_penalty
+            ),
+            no_repeat_ngram_size=(
+                no_repeat_ngram_size
+                if no_repeat_ngram_size is not None
+                else config.serve.no_repeat_ngram_size
+            ),
+        )
 
     def warmup(self) -> None:
         """Run one dummy inference so the first real request is fast."""
         import tensorflow as tf
 
         dummy = tf.zeros((299, 299, 3), dtype=tf.float32)
-        _ = generate_caption_greedy(self.model, self.tokenizer, dummy, self.config.model.max_length)
-        log.info("predictor_warmed_up")
+        _ = self.predict_tensor(dummy)
+        log.info("predictor_warmed_up", decode_strategy=self.decode_strategy)
 
     def predict_tensor(self, image_tensor) -> str:
         """Generate a caption from an already-preprocessed image tensor."""
+        if self.decode_strategy == "beam":
+            return generate_caption_beam(
+                self.model,
+                self.tokenizer,
+                image_tensor,
+                self.config.model.max_length,
+                beam_width=self.beam_width,
+                length_penalty=self.length_penalty,
+                repetition_penalty=self.repetition_penalty,
+                no_repeat_ngram_size=self.no_repeat_ngram_size,
+            )
         return generate_caption_greedy(
             self.model,
             self.tokenizer,

src/captioning/models/captioning_model.py

@@ -1,16 +1,30 @@
 """``ImageCaptioningModel`` — top-level Keras model with custom train/test step.
 
-Mirrors notebook cell 20 verbatim. The model owns its own loss & accuracy
-trackers (rather than using compile-time metrics) because the masked
+Mirrors notebook cell 20 verbatim by default. The model owns its own loss &
+accuracy trackers (rather than using compile-time metrics) because the masked
 arithmetic in ``calculate_loss`` / ``calculate_accuracy`` depends on the
 caption padding mask, which Keras's standard metric API can't see.
 
-Behavioural quirk preserved for parity (NOT a bug in our code):
+Two opt-in fixes layered on top of the baseline, both controlled by the
+constructor (defaults preserve the IEEE notebook quirk exactly):
+
+* ``honour_training_flag_in_test_step``
     The notebook's ``compute_loss_and_acc`` hardcodes ``training=True`` on
     both the encoder and decoder calls, even when invoked from ``test_step``.
     That means dropout is active during validation in the IEEE results.
-    We preserve this so BLEU matches the paper. Phase 1b will fix it in a
-    deliberate, clearly-marked commit.
+    Setting this flag to True restores the conventional behaviour — dropout
+    off in test_step — so val_loss reflects deployment behaviour and early
+    stopping fires on a clean signal.
+
+* ``correct_masked_accuracy``
+    The baseline's accuracy tracker is a ``Mean`` of per-batch ratios, which
+    weights small batches the same as large ones. Setting this flag to True
+    feeds the per-batch token count as ``sample_weight`` so the reported
+    metric is a true global token-level masked accuracy.
+
+Both knobs are off by default to keep numeric parity with the notebook; the
+trainer flips them on automatically when the user opts in via
+``train.honour_training_flag_in_test_step``.
 """
 
 from __future__ import annotations
@@ -22,12 +36,23 @@ def _build_captioning_model_class():
     class ImageCaptioningModel(tf.keras.Model):
         """Stitches CNN encoder + Transformer encoder + Transformer decoder."""
 
-        def __init__(self, cnn_model, encoder, decoder, image_aug=None) -> None:
+        def __init__(
+            self,
+            cnn_model,
+            encoder,
+            decoder,
+            image_aug=None,
+            *,
+            honour_training_flag_in_test_step: bool = False,
+            correct_masked_accuracy: bool = False,
+        ) -> None:
             super().__init__()
             self.cnn_model = cnn_model
             self.encoder = encoder
             self.decoder = decoder
             self.image_aug = image_aug
+            self.honour_training_flag_in_test_step = honour_training_flag_in_test_step
+            self.correct_masked_accuracy = correct_masked_accuracy
             self.loss_tracker = tf.keras.metrics.Mean(name="loss")
             self.acc_tracker = tf.keras.metrics.Mean(name="accuracy")
 
@@ -49,17 +74,20 @@ def _build_captioning_model_class():
         # --- shared loss/acc step (parity quirk: training=True hardcoded) --
 
         def compute_loss_and_acc(self, img_embed, captions, training=True):
-            # Notebook quirk preserved: encoder/decoder always called with
-            # training=True. The `training` parameter is intentionally unused.
-            del training  # silence linters: this is deliberate
-            encoder_output = self.encoder(img_embed, training=True)
+            # The IEEE notebook hardcoded `training=True` on encoder/decoder
+            # calls even from `test_step`, which means dropout is on during
+            # validation. Honouring the flag (opt-in) restores the standard
+            # behaviour and gives a cleaner val_loss signal.
+            effective_training = bool(training) if self.honour_training_flag_in_test_step else True
+            encoder_output = self.encoder(img_embed, training=effective_training)
             y_input = captions[:, :-1]
             y_true = captions[:, 1:]
             mask = y_true != 0
-            y_pred = self.decoder(y_input, encoder_output, training=True, mask=mask)
+            y_pred = self.decoder(y_input, encoder_output, training=effective_training, mask=mask)
             loss = self.calculate_loss(y_true, y_pred, mask)
             acc = self.calculate_accuracy(y_true, y_pred, mask)
-            return loss, acc
+            mask_count = tf.reduce_sum(tf.cast(mask, tf.float32))
+            return loss, acc, mask_count
 
         # --- Keras hooks ---------------------------------------------------
 
@@ -70,22 +98,30 @@ def _build_captioning_model_class():
             img_embed = self.cnn_model(imgs)
 
             with tf.GradientTape() as tape:
-                loss, acc = self.compute_loss_and_acc(img_embed, captions)
+                loss, acc, mask_count = self.compute_loss_and_acc(img_embed, captions)
 
             train_vars = self.encoder.trainable_variables + self.decoder.trainable_variables
             grads = tape.gradient(loss, train_vars)
             self.optimizer.apply_gradients(zip(grads, train_vars, strict=False))
             self.loss_tracker.update_state(loss)
-            self.acc_tracker.update_state(acc)
+            if self.correct_masked_accuracy:
+                # Weight per-batch accuracy by token count so the epoch
+                # average is a true global accuracy, not a mean of ratios.
+                self.acc_tracker.update_state(acc, sample_weight=mask_count)
+            else:
+                self.acc_tracker.update_state(acc)
 
             return {"loss": self.loss_tracker.result(), "acc": self.acc_tracker.result()}
 
         def test_step(self, batch):
             imgs, captions = batch
             img_embed = self.cnn_model(imgs)
-            loss, acc = self.compute_loss_and_acc(img_embed, captions, training=False)
+            loss, acc, mask_count = self.compute_loss_and_acc(img_embed, captions, training=False)
             self.loss_tracker.update_state(loss)
-            self.acc_tracker.update_state(acc)
+            if self.correct_masked_accuracy:
+                self.acc_tracker.update_state(acc, sample_weight=mask_count)
+            else:
+                self.acc_tracker.update_state(acc)
             return {"loss": self.loss_tracker.result(), "acc": self.acc_tracker.result()}
 
         @property

src/captioning/models/factory.py

@@ -63,4 +63,20 @@ def build_caption_model(
     )
     cnn = build_cnn_encoder()
     aug = default_image_augmentation() if use_augmentation else None
-    return ImageCaptioningModel(cnn_model=cnn, encoder=encoder, decoder=decoder, image_aug=aug)
+
+    # ``honour_training_flag_in_test_step`` and ``correct_masked_accuracy``
+    # default to False so this factory keeps producing notebook-parity models
+    # unless the user opts in by flipping the corresponding YAML flag.
+    honour_flag = bool(config.train.honour_training_flag_in_test_step)
+    # The masked-accuracy correction is harmless under parity (it's a
+    # better-weighted average of the same per-batch numbers), so we tie it to
+    # the same opt-in flag rather than adding a separate one — keeps the
+    # YAML surface minimal.
+    return ImageCaptioningModel(
+        cnn_model=cnn,
+        encoder=encoder,
+        decoder=decoder,
+        image_aug=aug,
+        honour_training_flag_in_test_step=honour_flag,
+        correct_masked_accuracy=honour_flag,
+    )

src/captioning/preprocessing/tokenizer.py

@@ -116,6 +116,16 @@ class CaptionTokenizer:
         word = self._idx2word(idx)
         return word.numpy().decode("utf-8")
 
+    def word_to_id(self, word: str) -> int:
+        """Look up a single word's integer id, returning 1 (the OOV id) if absent.
+
+        Used by beam search to seed beams with the ``[start]`` token without
+        going through ``TextVectorization``'s padded-string path.
+        """
+        self._require_fit()
+        assert self._word2idx is not None
+        return int(self._word2idx(word).numpy())
+
     # ---------------------------------------------------------- persistence ---
 
     def save(self, directory: str | Path) -> None:

src/captioning/training/__init__.py

@@ -1,21 +1,31 @@
-"""Training — losses, callbacks, and the trainer that orchestrates ``model.fit``.
+"""Training — losses, schedules, callbacks, and the trainer.
 
 The notebook computes loss + masked accuracy inside the model's ``train_step``;
 we keep that structure for parity but expose the loss function and callbacks
-as standalone modules so they can be unit-tested and reused (e.g. by Phase 1b
-beam-search evaluators).
+as standalone modules so they can be unit-tested and reused.
 
-    losses.py      ``masked_sparse_categorical_crossentropy`` — the same loss the notebook uses
-    callbacks.py   ``default_callbacks(config)`` — early stopping (and Phase 4 checkpoint hooks)
+    losses.py      ``masked_sparse_categorical_crossentropy`` (baseline) +
+                   ``label_smoothed_crossentropy`` + ``build_loss``
+    schedules.py   ``WarmupCosineDecay`` + ``build_learning_rate``
+    callbacks.py   ``default_callbacks(config)`` — early stopping + checkpoint
     trainer.py     ``Trainer.fit()`` — wraps compile + fit + history serialization
 """
 
 from captioning.training.callbacks import default_callbacks
-from captioning.training.losses import masked_sparse_categorical_crossentropy
+from captioning.training.losses import (
+    build_loss,
+    label_smoothed_crossentropy,
+    masked_sparse_categorical_crossentropy,
+)
+from captioning.training.schedules import WarmupCosineDecay, build_learning_rate
 from captioning.training.trainer import Trainer
 
 __all__ = [
     "Trainer",
+    "WarmupCosineDecay",
+    "build_learning_rate",
+    "build_loss",
     "default_callbacks",
+    "label_smoothed_crossentropy",
     "masked_sparse_categorical_crossentropy",
 ]

src/captioning/training/losses.py

@@ -8,8 +8,17 @@ Why ``reduction="none"``: the model's ``calculate_loss`` (cell 20) does the
 reduction itself, multiplying by the padding mask before averaging. A built-in
 reduction would average over the padded tokens too, biasing the loss.
 
-We expose the loss via a tiny factory rather than a constant so callers don't
-have to import TF themselves to get it.
+For the stabilisation phase we also support label-smoothed cross-entropy.
+Label smoothing replaces the one-hot target ``y_true`` with a mixture of the
+true label and a uniform distribution over the vocabulary:
+
+    target = (1 - eps) * one_hot(y) + eps / vocab_size
+
+The decoder's output is already softmaxed (`Dense(..., activation='softmax')`),
+so the loss reduces to ``-sum(target * log(p), axis=-1)``. Smoothing
+discourages the decoder from collapsing to a few high-probability tokens —
+the most common failure mode of cross-entropy-trained captioners and a
+likely root cause of the generic captions we're trying to fix.
 """
 
 from __future__ import annotations
@@ -25,3 +34,50 @@ def masked_sparse_categorical_crossentropy():
     import tensorflow as tf
 
     return tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False, reduction="none")
+
+
+def label_smoothed_crossentropy(label_smoothing: float, vocab_size: int):
+    """Per-token cross-entropy with uniform label smoothing.
+
+    Returned callable has the same signature as the sparse loss above
+    (``loss(y_true, y_pred) -> [B, T]``) so the model's masking machinery in
+    ``ImageCaptioningModel.calculate_loss`` works unchanged.
+
+    Args:
+        label_smoothing: Smoothing strength in ``[0, 1)``. ``0.0`` reduces to
+            the sparse-categorical baseline.
+        vocab_size: Size of the output distribution (matches the decoder's
+            final ``Dense`` units). Used to compute the uniform component.
+    """
+    import tensorflow as tf
+
+    if label_smoothing == 0.0:
+        return masked_sparse_categorical_crossentropy()
+
+    eps = float(label_smoothing)
+    log_eps = tf.constant(1e-12, dtype=tf.float32)
+    vocab = int(vocab_size)
+    uniform = eps / float(vocab)
+
+    def loss_fn(y_true, y_pred):
+        # y_true: [B, T] int ids; y_pred: [B, T, V] softmax probabilities.
+        y_pred = tf.cast(y_pred, tf.float32)
+        one_hot = tf.one_hot(tf.cast(y_true, tf.int32), depth=vocab, dtype=tf.float32)
+        target = one_hot * (1.0 - eps) + uniform
+        # Standard cross-entropy on softmax probs. Add log_eps to avoid log(0)
+        # on padding columns where the model would otherwise emit 0.
+        return -tf.reduce_sum(target * tf.math.log(y_pred + log_eps), axis=-1)
+
+    return loss_fn
+
+
+def build_loss(label_smoothing: float, vocab_size: int):
+    """Pick the right loss based on ``label_smoothing``.
+
+    Convenience wrapper so the trainer never has to branch on the smoothing
+    value itself — it always calls ``build_loss(...)`` and the right
+    implementation comes back.
+    """
+    if label_smoothing == 0.0:
+        return masked_sparse_categorical_crossentropy()
+    return label_smoothed_crossentropy(label_smoothing, vocab_size)

src/captioning/training/schedules.py

@@ -0,0 +1,107 @@
+"""Learning-rate schedules.
+
+The baseline pipeline uses a constant Adam LR (matching the IEEE notebook),
+which is fine for short fine-tuning runs but tends to leave Transformer
+captioners in a mediocre local minimum: the LR is too aggressive at start
+(decoder weights are still random) and too high near convergence (the model
+oscillates around a flat basin instead of settling).
+
+The fix the literature converged on is linear warmup followed by cosine
+decay (the GPT/BERT/ViT recipe):
+
+    lr(step) = peak_lr * step / warmup_steps                if step < warmup
+    lr(step) = min_lr + (peak_lr - min_lr) * 0.5 *
+               (1 + cos(pi * (step - warmup) / decay_steps)) otherwise
+
+We implement it as a ``LearningRateSchedule`` so the optimizer can call it
+per-step automatically, without us having to track step counts manually.
+"""
+
+from __future__ import annotations
+
+
+def _build_warmup_cosine_class():
+    """Lazy-build the schedule class to keep TF off the package import path."""
+    import tensorflow as tf
+
+    class WarmupCosineDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
+        """Linear warmup followed by cosine decay to ``min_learning_rate``.
+
+        Args:
+            peak_learning_rate: Maximum LR reached at the end of warmup.
+            warmup_steps: Number of steps to linearly ramp from 0 to peak.
+                ``0`` disables warmup (starts directly at ``peak``).
+            decay_steps: Number of steps over which to cosine-decay from
+                ``peak`` to ``min_learning_rate`` after warmup.
+            min_learning_rate: Floor reached at the end of decay (and held
+                thereafter).
+        """
+
+        def __init__(
+            self,
+            peak_learning_rate: float,
+            warmup_steps: int,
+            decay_steps: int,
+            min_learning_rate: float = 0.0,
+        ) -> None:
+            super().__init__()
+            self.peak_learning_rate = float(peak_learning_rate)
+            self.warmup_steps = int(warmup_steps)
+            self.decay_steps = max(int(decay_steps), 1)
+            self.min_learning_rate = float(min_learning_rate)
+
+        def __call__(self, step):
+            step = tf.cast(step, tf.float32)
+            peak = tf.constant(self.peak_learning_rate, dtype=tf.float32)
+            floor = tf.constant(self.min_learning_rate, dtype=tf.float32)
+            warmup = tf.constant(float(self.warmup_steps), dtype=tf.float32)
+            decay = tf.constant(float(self.decay_steps), dtype=tf.float32)
+
+            # During warmup: linear ramp 0 -> peak.
+            warmup_lr = peak * tf.math.divide_no_nan(step, warmup)
+
+            # After warmup: cosine decay peak -> floor over decay_steps.
+            progress = tf.minimum(1.0, tf.math.divide_no_nan(step - warmup, decay))
+            cosine = 0.5 * (1.0 + tf.cos(tf.constant(3.141592653589793) * progress))
+            decay_lr = floor + (peak - floor) * cosine
+
+            return tf.where(step < warmup, warmup_lr, decay_lr)
+
+        def get_config(self) -> dict[str, float | int]:
+            return {
+                "peak_learning_rate": self.peak_learning_rate,
+                "warmup_steps": self.warmup_steps,
+                "decay_steps": self.decay_steps,
+                "min_learning_rate": self.min_learning_rate,
+            }
+
+    return WarmupCosineDecay
+
+
+WarmupCosineDecay = _build_warmup_cosine_class()
+
+
+def build_learning_rate(
+    *,
+    schedule: str,
+    peak_learning_rate: float,
+    warmup_steps: int,
+    decay_steps: int,
+    min_learning_rate: float,
+):
+    """Return either a float (constant LR) or a :class:`WarmupCosineDecay`.
+
+    The optimizer treats a float as a fixed LR and a ``LearningRateSchedule``
+    as a per-step callable — we hide that asymmetry behind this factory so
+    the trainer only ever passes ``learning_rate=build_learning_rate(...)``.
+    """
+    if schedule == "constant":
+        return peak_learning_rate
+    if schedule == "cosine":
+        return WarmupCosineDecay(
+            peak_learning_rate=peak_learning_rate,
+            warmup_steps=warmup_steps,
+            decay_steps=decay_steps,
+            min_learning_rate=min_learning_rate,
+        )
+    raise ValueError(f"unsupported lr_schedule: {schedule!r}")

src/captioning/training/trainer.py

@@ -6,8 +6,11 @@ Wraps notebook cells 22 and 23 in a class so:
     * Phase 4 can replace the trainer with a CLI-driven main loop without
       changing the notebook-equivalent behaviour.
 
-The trainer is intentionally thin — no MLflow integration yet (Phase 2
-adds it), no distributed strategy (out of scope for the IEEE notebook).
+The trainer reads the optional training-stability fields off ``TrainConfig``
+(``label_smoothing``, ``lr_schedule``, ``warmup_steps``, ...). With defaults
+in place every existing config produces a byte-identical compile call to the
+notebook; flipping one YAML flag opts a run into the modern recipe without
+touching code.
 """
 
 from __future__ import annotations
@@ -17,12 +20,31 @@ from pathlib import Path
 
 from captioning.config.schema import AppConfig
 from captioning.training.callbacks import default_callbacks
-from captioning.training.losses import masked_sparse_categorical_crossentropy
+from captioning.training.losses import build_loss
+from captioning.training.schedules import build_learning_rate
 from captioning.utils.logging import get_logger
 
 log = get_logger(__name__)
 
 
+def _infer_steps_per_epoch(dataset) -> int | None:
+    """Best-effort cardinality probe for a ``tf.data.Dataset``.
+
+    Returns ``None`` when the dataset's cardinality is unknown or infinite.
+    Used only to derive ``cosine_decay_steps`` when the user didn't pin it
+    explicitly.
+    """
+    try:
+        import tensorflow as tf
+
+        card = int(tf.data.experimental.cardinality(dataset).numpy())
+    except Exception:  # — cardinality probing is best-effort
+        return None
+    if card in (-1, -2):  # UNKNOWN, INFINITE
+        return None
+    return card
+
+
 class Trainer:
     """Thin orchestration layer around an ``ImageCaptioningModel``."""
 
@@ -35,16 +57,47 @@ class Trainer:
         self.config = config
         self._compiled = False
 
-    def compile(self) -> None:
-        """Apply the same ``compile`` call the notebook makes (cell 22)."""
+    def compile(self, *, steps_per_epoch: int | None = None) -> None:
+        """Build optimizer + loss from config and call ``model.compile``.
+
+        Args:
+            steps_per_epoch: Used to derive ``cosine_decay_steps`` when the
+                config doesn't pin it explicitly. Passing ``None`` falls back
+                to the config value (or 1 if neither is set — degenerates to
+                immediate floor LR, but still a well-defined schedule).
+        """
         import tensorflow as tf
 
+        train = self.config.train
+        # Vocab size lives on the decoder's final Dense layer; pulling it here
+        # avoids threading the tokenizer through the trainer just for loss.
+        vocab_size = int(self.model.decoder.out.units)
+        loss = build_loss(train.label_smoothing, vocab_size)
+
+        cosine_steps = train.cosine_decay_steps or (
+            (steps_per_epoch or 1) * max(train.epochs - 0, 1)
+        )
+        learning_rate = build_learning_rate(
+            schedule=train.lr_schedule,
+            peak_learning_rate=train.learning_rate,
+            warmup_steps=train.warmup_steps,
+            decay_steps=cosine_steps,
+            min_learning_rate=train.min_learning_rate,
+        )
+
         self.model.compile(
-            optimizer=tf.keras.optimizers.Adam(learning_rate=self.config.train.learning_rate),
-            loss=masked_sparse_categorical_crossentropy(),
+            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
+            loss=loss,
         )
         self._compiled = True
-        log.info("model_compiled", learning_rate=self.config.train.learning_rate)
+        log.info(
+            "model_compiled",
+            lr_schedule=train.lr_schedule,
+            peak_learning_rate=train.learning_rate,
+            warmup_steps=train.warmup_steps,
+            cosine_decay_steps=cosine_steps,
+            label_smoothing=train.label_smoothing,
+        )
 
     def fit(
         self,
@@ -68,7 +121,7 @@ class Trainer:
             ``history.history`` as a ``dict[str, list[float]]``.
         """
         if not self._compiled:
-            self.compile()
+            self.compile(steps_per_epoch=_infer_steps_per_epoch(train_dataset))
 
         callbacks = default_callbacks(self.config, output_dir=output_dir)
         log.info("fit_start", epochs=self.config.train.epochs)

tests/unit/test_beam_decoder.py

@@ -0,0 +1,176 @@
+"""Beam-search decoder unit tests.
+
+The full TF decoder forward path is exercised by the parity audit and the
+smoke test in ``scripts/predict.py``. Here we test the *algorithmic* pieces
+of beam search in isolation:
+
+    * Length penalty correctly rescales scores.
+    * Repetition penalty downweights seen tokens.
+    * n-gram blocker forbids exact-repeat n-grams.
+    * Detokeniser strips ``[start]`` / ``[end]`` and stops at ``[end]``.
+
+A small fake model is used to verify end-to-end search behaviour without
+loading TensorFlow weights.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from captioning.inference.beam import (
+    _apply_repetition_penalty,
+    _Beam,
+    _blocks_repeat_ngram,
+    _detokenize,
+    _length_normalised,
+    generate_caption_beam,
+)
+
+
+def test_length_penalty_zero_returns_raw_score() -> None:
+    b = _Beam(token_ids=[1, 2, 3], score=-5.0)
+    assert _length_normalised(b, 0.0) == -5.0
+
+
+def test_length_penalty_one_divides_by_length() -> None:
+    b = _Beam(token_ids=[1, 2, 3, 4], score=-6.0)  # length=3
+    assert _length_normalised(b, 1.0) == pytest.approx(-2.0)
+
+
+def test_repetition_penalty_downweights_seen_tokens() -> None:
+    log_probs = np.array([-1.0, -2.0, -3.0, -4.0])
+    out = _apply_repetition_penalty(log_probs.copy(), history_ids={1, 3}, penalty=2.0)
+    # Penalty subtracts log(2) ~ 0.693 from seen-token log-probs.
+    assert out[0] == pytest.approx(-1.0)
+    assert out[1] == pytest.approx(-2.0 - np.log(2.0))
+    assert out[2] == pytest.approx(-3.0)
+    assert out[3] == pytest.approx(-4.0 - np.log(2.0))
+
+
+def test_repetition_penalty_one_is_noop() -> None:
+    log_probs = np.array([-1.0, -2.0, -3.0])
+    out = _apply_repetition_penalty(log_probs.copy(), history_ids={0, 1}, penalty=1.0)
+    np.testing.assert_array_equal(out, log_probs)
+
+
+def test_blocks_repeat_ngram_detects_repeat() -> None:
+    # seq ends with [4, 5]; appending 6 forms trigram [4, 5, 6] not present.
+    assert not _blocks_repeat_ngram([1, 2, 3, 4, 5], 6, n=3)
+    # Now seq contains [4, 5, 6]; appending 6 still wouldn't form a repeat.
+    assert not _blocks_repeat_ngram([4, 5, 6, 4, 5], 7, n=3)
+    # seq has [4, 5, 6] AND ends with [4, 5]; appending 6 repeats [4, 5, 6].
+    assert _blocks_repeat_ngram([4, 5, 6, 1, 4, 5], 6, n=3)
+
+
+def test_blocks_repeat_ngram_zero_size_disables() -> None:
+    assert not _blocks_repeat_ngram([1, 1, 1], 1, n=0)
+
+
+def test_detokenize_stops_at_end_and_skips_special_tokens() -> None:
+    tokenizer = MagicMock()
+    # ids: [start]=1, "a"=2, "man"=3, [end]=4
+    table = {1: "[start]", 2: "a", 3: "man", 4: "[end]"}
+    tokenizer.decode_id = lambda i: table[i]
+    out = _detokenize([1, 2, 3, 4, 99], tokenizer, end_id=4)
+    assert out == "a man"
+
+
+# ---- End-to-end beam search with a fake model -----------------------------
+
+
+class _FakeModel:
+    """Decoder fixture that always assigns the highest probability to ``best_id``.
+
+    The decoder output is the only piece beam search cares about; we stub the
+    CNN / encoder to identity-like behaviour so the whole inference pass runs
+    without TF being loaded.
+    """
+
+    def __init__(self, vocab_size: int, best_id: int) -> None:
+        self.vocab_size = vocab_size
+        self.best_id = best_id
+
+        self.cnn_model = MagicMock(side_effect=self._identity_image)
+        self.encoder = MagicMock(side_effect=self._identity_encoder)
+        self.decoder = MagicMock(side_effect=self._decoder_step)
+
+    def _identity_image(self, img):
+        return img
+
+    def _identity_encoder(self, x, training):
+        return x
+
+    def _decoder_step(self, tokens, encoded, training, mask):
+        import tensorflow as tf
+
+        batch = int(tf.shape(tokens)[0])
+        seq_len = int(tf.shape(tokens)[1])
+        probs = np.full((batch, seq_len, self.vocab_size), 1e-3, dtype=np.float32)
+        probs[:, :, self.best_id] = 0.999
+        # Normalise so each row over vocab sums to ~1.
+        probs /= probs.sum(axis=-1, keepdims=True)
+        return tf.convert_to_tensor(probs)
+
+
+def test_beam_search_emits_caption_when_model_prefers_end_token() -> None:
+    import tensorflow as tf
+
+    tokenizer = MagicMock()
+    # vocab: 0=pad, 1=[start], 2=[end], 3="dog"
+    word_to_id_table = {"[start]": 1, "[end]": 2}
+    decode_id_table = {0: "", 1: "[start]", 2: "[end]", 3: "dog"}
+    tokenizer.word_to_id = lambda w: word_to_id_table[w]
+    tokenizer.decode_id = lambda i: decode_id_table[i]
+
+    # Model always predicts "dog" (id=3).
+    model = _FakeModel(vocab_size=4, best_id=3)
+    image = tf.zeros((299, 299, 3), dtype=tf.float32)
+
+    caption = generate_caption_beam(
+        model,
+        tokenizer,
+        image,
+        max_length=6,
+        beam_width=2,
+        length_penalty=0.0,
+    )
+    # With no length penalty and no repetition penalty, the greedy-ish path
+    # outputs repeated "dog" until max_length. We just assert it produced
+    # *something* and didn't crash.
+    assert caption.startswith("dog")
+
+
+def test_beam_search_terminates_on_eos() -> None:
+    """Beam search must produce a clean caption when the model emits [end]."""
+    import tensorflow as tf
+
+    tokenizer = MagicMock()
+    word_to_id_table = {"[start]": 1, "[end]": 2}
+    decode_id_table = {0: "", 1: "[start]", 2: "[end]", 3: "dog"}
+    tokenizer.word_to_id = lambda w: word_to_id_table[w]
+    tokenizer.decode_id = lambda i: decode_id_table[i]
+
+    # Step 0: prefer "dog"; step 1+: prefer [end].
+    class _EosFakeModel(_FakeModel):
+        def _decoder_step(self, tokens, encoded, training, mask):
+            batch = int(tf.shape(tokens)[0])
+            seq_len = int(tf.shape(tokens)[1])
+            probs = np.full((batch, seq_len, self.vocab_size), 1e-3, dtype=np.float32)
+            probs[:, 0, 3] = 0.99  # at position 0 prefer "dog"
+            for pos in range(1, seq_len):
+                probs[:, pos, 2] = 0.99  # afterwards prefer [end]
+            probs /= probs.sum(axis=-1, keepdims=True)
+            return tf.convert_to_tensor(probs)
+
+    model = _EosFakeModel(vocab_size=4, best_id=3)
+    caption = generate_caption_beam(
+        model,
+        tokenizer,
+        tf.zeros((299, 299, 3), dtype=tf.float32),
+        max_length=6,
+        beam_width=2,
+    )
+    assert caption == "dog"

tests/unit/test_config.py

@@ -87,3 +87,44 @@ def test_modelconfig_independent_of_other_sections() -> None:
     assert m.vocabulary_size == 500
     # Defaults preserved
     assert m.max_length == 40
+
+
+# ---- Opt-in stability flags ------------------------------------------------
+
+
+def test_train_stability_defaults_preserve_notebook_parity() -> None:
+    t = TrainConfig()
+    assert t.label_smoothing == 0.0
+    assert t.lr_schedule == "constant"
+    assert t.warmup_steps == 0
+    assert t.honour_training_flag_in_test_step is False
+
+
+def test_label_smoothing_rejects_out_of_range() -> None:
+    with pytest.raises(ValidationError):
+        TrainConfig(label_smoothing=1.0)
+    with pytest.raises(ValidationError):
+        TrainConfig(label_smoothing=-0.1)
+
+
+def test_lr_schedule_rejects_unknown() -> None:
+    with pytest.raises(ValidationError):
+        TrainConfig(lr_schedule="square_wave")
+
+
+def test_decode_strategy_validates() -> None:
+    from captioning.config.schema import ServeConfig
+
+    with pytest.raises(ValidationError):
+        ServeConfig(decode_strategy="nucleus")
+    s = ServeConfig(decode_strategy="beam", beam_width=4)
+    assert s.beam_width == 4
+
+
+def test_beam_width_and_repetition_penalty_rejected_out_of_range() -> None:
+    from captioning.config.schema import ServeConfig
+
+    with pytest.raises(ValidationError):
+        ServeConfig(beam_width=0)
+    with pytest.raises(ValidationError):
+        ServeConfig(repetition_penalty=0.5)

tests/unit/test_evaluation_metrics.py

@@ -0,0 +1,210 @@
+"""Tests for ROUGE-L, CIDEr, METEOR adapters and the unified runner.
+
+We don't validate the upstream implementations — they have their own test
+suites. We *do* validate our adapters: sentinel stripping, ragged references,
+the perfect-prediction bound, and that the unified ``compute_all_metrics``
+correctly records partial failures in ``errors`` rather than crashing the
+whole pass.
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from captioning.evaluation import (
+    MIN_SAMPLES_FOR_CIDER,
+    BleuBreakdown,
+    RunMeta,
+    compute_all_metrics,
+    corpus_bleu_breakdown,
+    corpus_bleu_score,
+    corpus_cider_score,
+    corpus_rouge_l_score,
+    diagnose_many,
+    diagnose_sample,
+    write_diagnostics_jsonl,
+    write_run_artifacts,
+)
+
+# ---- BLEU ------------------------------------------------------------------
+
+
+def test_bleu_breakdown_returns_all_four_orders() -> None:
+    refs = [["a man riding a bike"], ["a dog in the park"]]
+    preds = ["a man riding a bike", "a dog in the park"]
+    result = corpus_bleu_breakdown(preds, refs)
+    assert isinstance(result, BleuBreakdown)
+    assert result.bleu1 == pytest.approx(100.0)
+    assert result.bleu2 == pytest.approx(100.0)
+    assert result.bleu4 == pytest.approx(100.0)
+    assert corpus_bleu_score(preds, refs) == pytest.approx(result.bleu4)
+
+
+def test_bleu_strips_sentinels_before_scoring() -> None:
+    refs = [["[start] a man riding a bike [end]"]]
+    preds = ["[start] a man riding a bike [end]"]
+    assert corpus_bleu_score(preds, refs) == pytest.approx(100.0)
+
+
+# ---- ROUGE-L ---------------------------------------------------------------
+
+
+def test_rouge_l_perfect_matches_score_100() -> None:
+    refs = [["a man riding a bike"], ["a dog in the park"]]
+    preds = ["a man riding a bike", "a dog in the park"]
+    score = corpus_rouge_l_score(preds, refs)
+    assert score == pytest.approx(100.0)
+
+
+def test_rouge_l_partial_overlap_scores_in_range() -> None:
+    refs = [["a man riding a bike on a road"]]
+    preds = ["a man on a road"]
+    score = corpus_rouge_l_score(preds, refs)
+    # Reference has 7 tokens, prediction has 5 tokens, LCS=5
+    # P = 5/5 = 1.0, R = 5/7 ≈ 0.71, F ≈ 0.83
+    assert 70.0 < score < 90.0
+
+
+def test_rouge_l_picks_best_reference() -> None:
+    refs = [["xyz qrs nothing matches", "a man riding a bike"]]
+    preds = ["a man riding a bike"]
+    score = corpus_rouge_l_score(preds, refs)
+    assert score == pytest.approx(100.0)
+
+
+def test_rouge_l_length_mismatch_raises() -> None:
+    with pytest.raises(ValueError):
+        corpus_rouge_l_score(["a"], [["a"], ["b"]])
+
+
+# ---- CIDEr -----------------------------------------------------------------
+
+
+def test_cider_requires_minimum_samples() -> None:
+    with pytest.raises(ValueError, match="degenerate"):
+        corpus_cider_score(["a man"], [["a man"]])
+
+
+def test_cider_returns_positive_for_good_predictions() -> None:
+    refs = [
+        ["a man riding a bike"],
+        ["a dog in the park"],
+        ["two children playing"],
+    ]
+    preds = ["a man riding a bike", "a dog in the park", "two children playing"]
+    score = corpus_cider_score(preds, refs)
+    assert score > 0.0
+
+
+# ---- Runner ----------------------------------------------------------------
+
+
+def test_compute_all_metrics_returns_every_field() -> None:
+    refs = [["a man riding a bike"], ["a dog in the park"]]
+    preds = ["a man riding a bike", "a dog in the park"]
+    report = compute_all_metrics(preds, refs, include_meteor=False, include_cider=False)
+    assert report.n_examples == 2
+    assert report.bleu1 is not None
+    assert report.bleu4 is not None
+    assert report.rouge_l is not None
+    assert report.meteor is None  # explicitly skipped
+    assert report.cider is None  # explicitly skipped
+
+
+def test_compute_all_metrics_skips_cider_on_tiny_corpus() -> None:
+    refs = [["a man riding a bike"]]
+    preds = ["a man riding a bike"]
+    report = compute_all_metrics(preds, refs, include_meteor=False)
+    assert report.cider is None
+    assert "cider" in report.errors
+
+
+def test_compute_all_metrics_serialises_to_dict() -> None:
+    refs = [["a man riding a bike"], ["a dog in the park"]]
+    preds = ["a man riding a bike", "a dog in the park"]
+    report = compute_all_metrics(preds, refs, include_meteor=False, include_cider=False)
+    payload = report.to_dict()
+    # JSON-roundtrip must not lose information.
+    assert json.loads(json.dumps(payload)) == payload
+
+
+# ---- Inspection -----------------------------------------------------------
+
+
+def test_diagnose_sample_flags_empty_prediction() -> None:
+    d = diagnose_sample("img.jpg", "", ["a man riding a bike"])
+    assert "empty" in d.flags
+    assert d.length_tokens == 0
+
+
+def test_diagnose_sample_flags_repetitive_prediction() -> None:
+    d = diagnose_sample("img.jpg", "a a a a man", ["a man riding a bike"])
+    assert "repetitive" in d.flags
+    assert d.longest_repeat_run == 4
+
+
+def test_diagnose_sample_flags_very_short_prediction() -> None:
+    d = diagnose_sample("img.jpg", "a man", ["a man riding a bike"])
+    assert "very_short" in d.flags
+
+
+def test_diagnose_many_writes_jsonl(tmp_path) -> None:
+    images = ["a.jpg", "b.jpg"]
+    preds = ["a man riding a bike", ""]
+    refs = [["a man on a bicycle"], ["a dog in the park"]]
+    diags = diagnose_many(images, preds, refs)
+    out = tmp_path / "diag.jsonl"
+    write_diagnostics_jsonl(diags, out)
+    lines = out.read_text(encoding="utf-8").splitlines()
+    assert len(lines) == 2
+    parsed = [json.loads(line) for line in lines]
+    assert parsed[0]["image"] == "a.jpg"
+    # Empty prediction also flags as ``very_short`` because it has 0 tokens.
+    assert "empty" in parsed[1]["flags"]
+
+
+# ---- Benchmark scaffolding ------------------------------------------------
+
+
+def test_write_run_artifacts_emits_expected_files(tmp_path) -> None:
+    images = ["a.jpg", "b.jpg"]
+    preds = ["a man riding a bike", "a dog in the park"]
+    refs = [["a man on a bicycle"], ["a dog in the park"]]
+    diags = diagnose_many(images, preds, refs)
+    report = compute_all_metrics(preds, refs, include_meteor=False, include_cider=False)
+    meta = RunMeta(
+        model_id="test-model",
+        decode_strategy="greedy",
+        weights_path="nowhere",
+        tokenizer_dir="nowhere",
+        n_samples=len(preds),
+        max_length=40,
+    )
+
+    out_dir = write_run_artifacts(
+        tmp_path / "runX",
+        metrics=report,
+        meta=meta,
+        images=images,
+        predictions=preds,
+        references=refs,
+        diagnostics=diags,
+    )
+
+    assert (out_dir / "metrics.json").is_file()
+    assert (out_dir / "run_meta.json").is_file()
+    assert (out_dir / "predictions.jsonl").is_file()
+    assert (out_dir / "diagnostics.jsonl").is_file()
+    assert (out_dir / "report.md").is_file()
+
+    predictions_lines = (out_dir / "predictions.jsonl").read_text(encoding="utf-8").splitlines()
+    assert len(predictions_lines) == 2
+
+    metadata = json.loads((out_dir / "run_meta.json").read_text(encoding="utf-8"))
+    assert metadata["model_id"] == "test-model"
+    assert metadata["n_samples"] == 2
+
+
+_ = MIN_SAMPLES_FOR_CIDER  # — exposed re-export, exercised by import

tests/unit/test_tokenizer.py

@@ -65,3 +65,15 @@ def test_max_length_is_respected(tiny_caption_corpus: list[str]) -> None:
     long_caption = " ".join(["[start]"] + ["word"] * 30 + ["[end]"])
     ids = tok.encode([long_caption])
     assert ids.shape == (1, 10)
+
+
+def test_word_to_id_round_trips_with_decode(tiny_caption_corpus: list[str]) -> None:
+    """``word_to_id`` is the inverse of ``decode_id`` for in-vocabulary tokens."""
+    tok = CaptionTokenizer(vocab_size=200, max_length=20)
+    tok.fit(tiny_caption_corpus)
+    start_id = tok.word_to_id("[start]")
+    end_id = tok.word_to_id("[end]")
+    assert isinstance(start_id, int)
+    assert start_id != end_id
+    assert tok.decode_id(start_id) == "[start]"
+    assert tok.decode_id(end_id) == "[end]"

tests/unit/test_training_stability.py

@@ -0,0 +1,136 @@
+"""Tests for the opt-in training-stability primitives.
+
+Covers:
+    * ``label_smoothed_crossentropy`` returns a per-token loss tensor with the
+      same shape as the baseline sparse loss, and reduces to it at smoothing=0.
+    * ``WarmupCosineDecay`` produces the expected piecewise schedule.
+    * ``build_loss`` / ``build_learning_rate`` dispatch correctly on config.
+"""
+
+from __future__ import annotations
+
+from itertools import pairwise
+
+import numpy as np
+import pytest
+
+from captioning.training.losses import build_loss, label_smoothed_crossentropy
+from captioning.training.schedules import WarmupCosineDecay, build_learning_rate
+
+# ---- Label smoothing -------------------------------------------------------
+
+
+def test_label_smoothed_loss_returns_per_token_shape() -> None:
+    import tensorflow as tf
+
+    vocab = 5
+    loss_fn = label_smoothed_crossentropy(0.1, vocab)
+    y_true = tf.constant([[1, 2, 0]], dtype=tf.int32)
+    y_pred = tf.constant(
+        [
+            [
+                [0.05, 0.85, 0.05, 0.025, 0.025],
+                [0.05, 0.05, 0.85, 0.025, 0.025],
+                [0.85, 0.05, 0.05, 0.025, 0.025],
+            ]
+        ],
+        dtype=tf.float32,
+    )
+    loss = loss_fn(y_true, y_pred).numpy()
+    assert loss.shape == (1, 3)
+    # The first two tokens are confidently correct → low loss.
+    assert loss[0, 0] < 1.0
+    assert loss[0, 1] < 1.0
+
+
+def test_label_smoothing_with_zero_returns_baseline_loss() -> None:
+    loss = build_loss(0.0, vocab_size=10)
+    # The baseline SparseCategoricalCrossentropy is an instance, not a function.
+    import tensorflow as tf
+
+    assert isinstance(loss, tf.keras.losses.SparseCategoricalCrossentropy)
+
+
+def test_label_smoothing_is_higher_than_unsmoothed_on_perfect_prediction() -> None:
+    """Smoothing punishes overconfidence — perfect one-hot prediction gets a
+    higher per-token loss with smoothing > 0 than without."""
+    import tensorflow as tf
+
+    vocab = 5
+    y_true = tf.constant([[1]], dtype=tf.int32)
+    one_hot_pred = tf.constant([[[0.0, 1.0, 0.0, 0.0, 0.0]]], dtype=tf.float32)
+
+    smoothed = label_smoothed_crossentropy(0.1, vocab)(y_true, one_hot_pred).numpy()
+    unsmoothed = -np.log(1.0)  # sparse cross-entropy on argmax==y_true is 0
+    assert smoothed[0, 0] > unsmoothed + 1e-3
+
+
+# ---- Learning-rate schedule -----------------------------------------------
+
+
+def test_warmup_cosine_zero_at_step_zero() -> None:
+    import tensorflow as tf
+
+    schedule = WarmupCosineDecay(peak_learning_rate=1.0, warmup_steps=10, decay_steps=100)
+    assert float(schedule(tf.constant(0, dtype=tf.int64))) == pytest.approx(0.0)
+
+
+def test_warmup_cosine_peaks_at_end_of_warmup() -> None:
+    import tensorflow as tf
+
+    schedule = WarmupCosineDecay(peak_learning_rate=1.0, warmup_steps=10, decay_steps=100)
+    assert float(schedule(tf.constant(10, dtype=tf.int64))) == pytest.approx(1.0, abs=1e-3)
+
+
+def test_warmup_cosine_floors_at_end_of_decay() -> None:
+    import tensorflow as tf
+
+    schedule = WarmupCosineDecay(
+        peak_learning_rate=1.0,
+        warmup_steps=10,
+        decay_steps=100,
+        min_learning_rate=0.1,
+    )
+    final = float(schedule(tf.constant(110, dtype=tf.int64)))
+    assert final == pytest.approx(0.1, abs=1e-3)
+
+
+def test_warmup_cosine_is_monotone_during_warmup() -> None:
+    import tensorflow as tf
+
+    schedule = WarmupCosineDecay(peak_learning_rate=1.0, warmup_steps=10, decay_steps=100)
+    values = [float(schedule(tf.constant(s, dtype=tf.int64))) for s in range(11)]
+    assert all(b >= a for a, b in pairwise(values))
+
+
+def test_build_learning_rate_returns_float_for_constant() -> None:
+    lr = build_learning_rate(
+        schedule="constant",
+        peak_learning_rate=1e-3,
+        warmup_steps=0,
+        decay_steps=10,
+        min_learning_rate=0.0,
+    )
+    assert lr == 1e-3
+
+
+def test_build_learning_rate_returns_schedule_for_cosine() -> None:
+    lr = build_learning_rate(
+        schedule="cosine",
+        peak_learning_rate=1e-3,
+        warmup_steps=5,
+        decay_steps=50,
+        min_learning_rate=0.0,
+    )
+    assert isinstance(lr, WarmupCosineDecay)
+
+
+def test_build_learning_rate_rejects_unknown_schedule() -> None:
+    with pytest.raises(ValueError, match="unsupported"):
+        build_learning_rate(
+            schedule="square_wave",
+            peak_learning_rate=1.0,
+            warmup_steps=0,
+            decay_steps=10,
+            min_learning_rate=0.0,
+        )