holoscript-trait-inference 0.1.0__py3-none-any.whl

trait_inference/model/decoder.py

@@ -0,0 +1,192 @@
+"""Constrained-decoding trait set generator — Paper 19 contribution.
+
+Per `phase-1-spec.md` §3.1: small LLM decoder with constrained
+generation over the trait composition lattice. The decoder samples ONLY
+JSON arrays of valid trait names, in canonical sorted order, with no
+duplicates. This bakes gate item 4 (≥90% validity) into the
+ARCHITECTURE per F.031.
+
+Approach:
+- HF transformers LLM (default Qwen2.5-0.5B; ≤1B sweep target).
+- Prompt template: {description} → {trait JSON}.
+- Constrained decoding via outlines.RegexFSM with a regex enumerating
+  the label space.
+- Post-decode: parse JSON, deduplicate, return as sorted set
+  (defense-in-depth).
+
+Heavy deps (torch, transformers, outlines) imported only when
+TraitDecoder is constructed.
+"""
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any, Sequence
+
+
+# ----------------------------------------------------------------------------
+# Regex builder for the trait JSON list (no heavy deps)
+# ----------------------------------------------------------------------------
+
+def build_trait_regex(label_space: Sequence[str]) -> str:
+    """Build a regex that matches a JSON list of traits drawn from the
+    label space, in any order, with arbitrary length 0..N.
+
+    Returns a Python regex string suitable for outlines.RegexFSM.
+
+    NOTE: This regex enforces MEMBERSHIP in the label space and JSON
+    syntax, but does NOT enforce sorted-unique-set semantics — those
+    are applied post-decode via parse_trait_set(). Stateful
+    sort-unique constraint is possible via outlines CFG but adds
+    significant complexity; defense-in-depth post-processing is
+    sufficient for the spec.
+
+    Examples (with label_space = ["@grabbable", "@physics", "@rigid"]):
+      []                                       — valid (empty)
+      ["@grabbable"]                           — valid
+      ["@grabbable", "@physics"]               — valid
+      ["@physics", "@grabbable"]               — valid (post-dedup will sort)
+      ["@unknown"]                             — REJECTED by regex
+    """
+    if not label_space:
+        return r"\[\]"
+    # Each trait literal — escaped for regex, JSON-quoted.
+    trait_alts = "|".join(re.escape(t) for t in sorted(set(label_space)))
+    one_trait = rf'"({trait_alts})"'
+    # `[]` OR `[T(, T)*]`
+    return rf'\[(?:{one_trait}(?:, {one_trait})*)?\]'
+
+
+def parse_trait_set(text: str, label_space: Sequence[str] | None = None) -> list[str]:
+    """Parse a JSON-list string from the decoder, return canonical
+    sorted-unique trait list.
+
+    Tolerant: returns [] on parse failure rather than raising — generator
+    output should be assumed well-formed (constrained decoding) but
+    defense-in-depth means we never crash the eval harness on a single
+    bad output.
+
+    If label_space is provided, traits not in label_space are dropped.
+    """
+    text = text.strip()
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        return []
+    if not isinstance(parsed, list):
+        return []
+    out: set[str] = set()
+    for item in parsed:
+        if not isinstance(item, str):
+            continue
+        if label_space is not None and item not in label_space:
+            continue
+        out.add(item)
+    return sorted(out)
+
+
+# ----------------------------------------------------------------------------
+# TraitDecoder — model wrapper
+# ----------------------------------------------------------------------------
+
+@dataclass
+class TraitDecoderConfig:
+    """Config for the constrained-decoding model."""
+    model_name: str = "Qwen/Qwen2.5-0.5B"   # ≤1B per spec; sweep over Qwen-0.5B/1B + Llama-3.2-1B
+    label_space: tuple[str, ...] = ()
+    prompt_template: str = (
+        "Given a description of a virtual object, return the minimal "
+        "set of HoloScript traits that apply, as a JSON array.\n"
+        "Description: {description}\n"
+        "Traits: "
+    )
+    max_new_tokens: int = 256
+    temperature: float = 0.0    # deterministic; required for Paper 19 reproducibility
+    device: str = "cuda"        # override "cpu" for local smoke-test
+    encoder_name: str | None = None  # if set, use sentence-transformer for input embedding (Phase 2.1 ablation axis)
+
+
+class TraitDecoder:
+    """Wrapper over HF transformers + outlines constrained generation.
+
+    Construction triggers heavy imports + model load. For CPU smoke
+    tests, pass `device="cpu"` AND a tiny model_name like
+    "sshleifer/tiny-gpt2".
+
+    Usage:
+        decoder = TraitDecoder(TraitDecoderConfig(label_space=(...,)))
+        traits = decoder.predict("a ball that bounces")
+        traits  # ["@collidable", "@physics", "@rigid"]
+    """
+
+    def __init__(self, config: TraitDecoderConfig):
+        if not config.label_space:
+            raise ValueError("TraitDecoderConfig.label_space must be non-empty")
+        # Heavy imports happen here, not at module load.
+        try:
+            import torch  # noqa: F401
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:
+            raise ImportError(
+                "TraitDecoder requires the [model] extra. "
+                "Install via: pip install -e '.[model]'"
+            ) from exc
+
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.model = AutoModelForCausalLM.from_pretrained(config.model_name)
+        self.model.to(config.device)
+        self.model.eval()
+
+        self._regex = build_trait_regex(config.label_space)
+        self._fsm: Any = None  # lazy-built on first predict (outlines optional at construction)
+
+    def _build_fsm(self) -> Any:
+        """Build the outlines FSM for constrained decoding."""
+        if self._fsm is not None:
+            return self._fsm
+        try:
+            import outlines  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "TraitDecoder.predict() requires outlines. "
+                "Install via: pip install -e '.[model]'"
+            ) from exc
+        # outlines API surface has changed across versions; we use the
+        # generate.regex helper which is stable as of 0.0.40+.
+        from outlines.models.transformers import Transformers
+        from outlines.generate import regex as outlines_regex_gen
+        wrapped = Transformers(self.model, self.tokenizer)
+        self._fsm = outlines_regex_gen(wrapped, self._regex)
+        return self._fsm
+
+    def predict(self, description: str) -> list[str]:
+        """Generate constrained trait set for a single description."""
+        return self.predict_batch([description])[0]
+
+    def predict_batch(self, descriptions: list[str]) -> list[list[str]]:
+        """Batch predict. Returns list of canonical sorted trait lists."""
+        fsm = self._build_fsm()
+        prompts = [self.config.prompt_template.format(description=d) for d in descriptions]
+        # outlines `generate.regex` returns a callable; arg is prompt string(s).
+        results = fsm(prompts, max_tokens=self.config.max_new_tokens)
+        if isinstance(results, str):
+            results = [results]
+        out: list[list[str]] = []
+        for text in results:
+            out.append(parse_trait_set(text, label_space=self.config.label_space))
+        return out
+
+    @classmethod
+    def for_smoke_test(cls, label_space: Sequence[str]) -> "TraitDecoder":
+        """Construct a CPU-only tiny model for unit/smoke tests."""
+        return cls(TraitDecoderConfig(
+            model_name="sshleifer/tiny-gpt2",
+            label_space=tuple(label_space),
+            device="cpu",
+            max_new_tokens=64,
+        ))

trait_inference/model/sweep.py

@@ -0,0 +1,242 @@
+"""Hyperparameter sweep runner — Paper 19.
+
+Per `phase-1-spec.md` §3.3 + §4.4: 162-cell hyperparameter sweep
+fractional-factorial-pruned to ~30 cells. Per `preregistration.md`
+§1: N=5 reseed per headline cell.
+
+Sweep axes:
+  - decoder model:        Qwen-0.5B / Qwen-1B / Llama-3.2-1B
+  - encoder name:         all-MiniLM-L6 / mpnet-base / None (LLM-only)
+  - beam width:           1 / 4 / 8 (inference-time)
+  - constraint timing:    per-token / per-emission
+  - training set size:    500 / 1000 / 2000 (resolves seed doc Q1)
+
+Each cell = ~6 GPU-hours on RTX 4090; full sweep at 30 cells × N=5
+reseed = 900 GPU-hours = ~$240 at $0.27/hr × 30 parallel GPUs ≈ 30hr
+wall-clock if fully parallel.
+
+This module designs the sweep — the runner is the entry point invoked
+by trait-inference CLI `model sweep`.
+"""
+from __future__ import annotations
+
+import itertools
+import json
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Iterator
+
+
+@dataclass(frozen=True)
+class SweepCell:
+    """One cell in the hyperparameter sweep."""
+    cell_id: str
+    decoder_model: str
+    encoder_name: str | None
+    beam_width: int
+    constraint_timing: str   # "per_token" | "per_emission"
+    train_size: int
+    reseed_index: int        # 0..N-1 per preregistration N=5
+
+
+@dataclass
+class SweepConfig:
+    """Sweep configuration — defaults match `phase-1-spec.md` §3.3."""
+    decoder_models: tuple[str, ...] = (
+        "Qwen/Qwen2.5-0.5B",
+        "Qwen/Qwen2.5-1.5B",      # 1B-class
+        "meta-llama/Llama-3.2-1B",
+    )
+    encoder_names: tuple[str | None, ...] = (
+        None,                                      # LLM-only baseline
+        "sentence-transformers/all-MiniLM-L6-v2",
+        "sentence-transformers/mpnet-base-v2",
+    )
+    beam_widths: tuple[int, ...] = (1, 4, 8)
+    constraint_timings: tuple[str, ...] = ("per_token", "per_emission")
+    train_sizes: tuple[int, ...] = (500, 1000, 2000)
+    reseed_n: int = 5    # per preregistration §1
+    fractional_factorial_pruned: bool = True
+    pruned_target: int = 30   # spec target for ~30 cells
+
+    def full_grid(self) -> list[SweepCell]:
+        """Generate the full Cartesian-product grid (162 cells before
+        pruning, 162 × N=5 = 810 with reseed)."""
+        cells = []
+        for i, (dm, en, bw, ct, ts) in enumerate(itertools.product(
+            self.decoder_models,
+            self.encoder_names,
+            self.beam_widths,
+            self.constraint_timings,
+            self.train_sizes,
+        )):
+            for r in range(self.reseed_n):
+                cells.append(SweepCell(
+                    cell_id=f"cell-{i:03d}-r{r}",
+                    decoder_model=dm,
+                    encoder_name=en,
+                    beam_width=bw,
+                    constraint_timing=ct,
+                    train_size=ts,
+                    reseed_index=r,
+                ))
+        return cells
+
+    def pruned_grid(self) -> list[SweepCell]:
+        """Generate the fractional-factorial-pruned grid.
+
+        Strategy: for each axis, vary one-at-a-time around a center
+        point (Plackett-Burman-style). Center = (Qwen-0.5B, mpnet-base,
+        beam=4, per_emission, train=1000). For headline measurement,
+        all reseeds (N=5) of the BEST cell from the pruned grid.
+
+        This produces ~30 cells matching `phase-1-spec.md` §3.3 target.
+        """
+        center = SweepCell(
+            cell_id="center",
+            decoder_model="Qwen/Qwen2.5-0.5B",
+            encoder_name="sentence-transformers/mpnet-base-v2",
+            beam_width=4,
+            constraint_timing="per_emission",
+            train_size=1000,
+            reseed_index=0,
+        )
+        cells: list[SweepCell] = [center]
+
+        # One-at-a-time perturbations
+        for dm in self.decoder_models:
+            if dm == center.decoder_model:
+                continue
+            cells.append(SweepCell(
+                cell_id=f"axis-decoder-{dm.split('/')[-1]}",
+                decoder_model=dm,
+                encoder_name=center.encoder_name,
+                beam_width=center.beam_width,
+                constraint_timing=center.constraint_timing,
+                train_size=center.train_size,
+                reseed_index=0,
+            ))
+        for en in self.encoder_names:
+            if en == center.encoder_name:
+                continue
+            cells.append(SweepCell(
+                cell_id=f"axis-encoder-{en if en else 'none'}".replace("/", "_"),
+                decoder_model=center.decoder_model,
+                encoder_name=en,
+                beam_width=center.beam_width,
+                constraint_timing=center.constraint_timing,
+                train_size=center.train_size,
+                reseed_index=0,
+            ))
+        for bw in self.beam_widths:
+            if bw == center.beam_width:
+                continue
+            cells.append(SweepCell(
+                cell_id=f"axis-beam-{bw}",
+                decoder_model=center.decoder_model,
+                encoder_name=center.encoder_name,
+                beam_width=bw,
+                constraint_timing=center.constraint_timing,
+                train_size=center.train_size,
+                reseed_index=0,
+            ))
+        for ct in self.constraint_timings:
+            if ct == center.constraint_timing:
+                continue
+            cells.append(SweepCell(
+                cell_id=f"axis-constraint-{ct}",
+                decoder_model=center.decoder_model,
+                encoder_name=center.encoder_name,
+                beam_width=center.beam_width,
+                constraint_timing=ct,
+                train_size=center.train_size,
+                reseed_index=0,
+            ))
+        for ts in self.train_sizes:
+            if ts == center.train_size:
+                continue
+            cells.append(SweepCell(
+                cell_id=f"axis-trainsize-{ts}",
+                decoder_model=center.decoder_model,
+                encoder_name=center.encoder_name,
+                beam_width=center.beam_width,
+                constraint_timing=center.constraint_timing,
+                train_size=ts,
+                reseed_index=0,
+            ))
+
+        # Reseed all cells N=5 times for headline measurement
+        full = []
+        for cell in cells:
+            for r in range(self.reseed_n):
+                full.append(SweepCell(
+                    cell_id=f"{cell.cell_id}-r{r}",
+                    decoder_model=cell.decoder_model,
+                    encoder_name=cell.encoder_name,
+                    beam_width=cell.beam_width,
+                    constraint_timing=cell.constraint_timing,
+                    train_size=cell.train_size,
+                    reseed_index=r,
+                ))
+        return full
+
+    def to_dict(self) -> dict:
+        return {
+            "decoder_models": list(self.decoder_models),
+            "encoder_names": [str(e) for e in self.encoder_names],
+            "beam_widths": list(self.beam_widths),
+            "constraint_timings": list(self.constraint_timings),
+            "train_sizes": list(self.train_sizes),
+            "reseed_n": self.reseed_n,
+            "fractional_factorial_pruned": self.fractional_factorial_pruned,
+            "pruned_target": self.pruned_target,
+        }
+
+
+class TraitSweep:
+    """Sweep runner — emits one cell config per JSON file under
+    `output_dir/cells/`, ready to be claimed by parallel GPU agents.
+
+    Does NOT run training itself — that's the per-cell trainer's job
+    (a GPU agent picks up a cell config + runs trainer + emits
+    measurement JSON).
+
+    The split (sweep generates work units, trainer consumes them) lets
+    30 GPUs run in parallel without coordination — each GPU claims one
+    or more cells from `output_dir/cells/`.
+    """
+
+    def __init__(self, config: SweepConfig | None = None):
+        self.config = config or SweepConfig()
+
+    def emit_cells(self, output_dir: str | Path) -> dict:
+        """Write one JSON file per cell to `output_dir/cells/`. Returns
+        summary dict."""
+        output_dir = Path(output_dir)
+        cells_dir = output_dir / "cells"
+        cells_dir.mkdir(parents=True, exist_ok=True)
+
+        cells = (
+            self.config.pruned_grid()
+            if self.config.fractional_factorial_pruned
+            else self.config.full_grid()
+        )
+
+        for cell in cells:
+            cell_path = cells_dir / f"{cell.cell_id}.json"
+            cell_path.write_text(
+                json.dumps(asdict(cell), indent=2, sort_keys=False),
+                encoding="utf-8",
+            )
+
+        summary = {
+            "sweep_config": self.config.to_dict(),
+            "cell_count": len(cells),
+            "cells_dir": str(cells_dir),
+            "emitted_at": datetime.now(timezone.utc).isoformat(),
+        }
+        (output_dir / "sweep_summary.json").write_text(
+            json.dumps(summary, indent=2), encoding="utf-8"
+        )
+        return summary

trait_inference/model/trainer.py

@@ -0,0 +1,237 @@
+"""Training loop for the constrained-decoding trait model — Paper 19.
+
+Per `phase-1-spec.md` §3.3: HF Trainer fine-tune over (description,
+trait JSON) pairs with cross-entropy loss + early stopping on
+validation F1 macro.
+
+Heavy deps imported lazily.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+from trait_inference.dataset import Pair, Split
+
+
+@dataclass
+class TrainConfig:
+    """Frozen-spec hyperparameters per `phase-1-spec.md` §3.3.
+
+    Sweep axes (per §3.3):
+      - model_name (Qwen-0.5B / Qwen-1B / Llama-3.2-1B)
+      - learning_rate (encoder vs decoder)
+      - batch_size, num_epochs (compute-budget bounded)
+
+    Frozen per `preregistration.md`:
+      - reseed N=5
+      - early_stopping_patience=3 on val F1 macro
+      - bootstrap_b=1000 on headline measurement (eval only)
+    """
+    model_name: str = "Qwen/Qwen2.5-0.5B"
+    label_space: tuple[str, ...] = ()
+    output_dir: str = "checkpoints/trait_decoder_v0"
+    num_epochs: int = 20
+    train_batch_size: int = 32
+    eval_batch_size: int = 32
+    learning_rate: float = 5e-5
+    warmup_steps: int = 500
+    weight_decay: float = 0.01
+    max_seq_len: int = 512
+    logging_steps: int = 50
+    eval_steps: int = 500
+    save_steps: int = 500
+    early_stopping_patience: int = 3
+    seed: int = 42
+    fp16: bool = True             # mixed-precision on supported GPUs
+    grad_accumulation_steps: int = 1
+    prompt_template: str = (
+        "Given a description of a virtual object, return the minimal "
+        "set of HoloScript traits that apply, as a JSON array.\n"
+        "Description: {description}\n"
+        "Traits: {trait_json}"
+    )
+
+    def to_dict(self) -> dict:
+        return {
+            "model_name": self.model_name,
+            "label_space_size": len(self.label_space),
+            "output_dir": self.output_dir,
+            "num_epochs": self.num_epochs,
+            "train_batch_size": self.train_batch_size,
+            "eval_batch_size": self.eval_batch_size,
+            "learning_rate": self.learning_rate,
+            "warmup_steps": self.warmup_steps,
+            "weight_decay": self.weight_decay,
+            "max_seq_len": self.max_seq_len,
+            "logging_steps": self.logging_steps,
+            "eval_steps": self.eval_steps,
+            "save_steps": self.save_steps,
+            "early_stopping_patience": self.early_stopping_patience,
+            "seed": self.seed,
+            "fp16": self.fp16,
+            "grad_accumulation_steps": self.grad_accumulation_steps,
+        }
+
+
+def _format_example(pair: Pair, prompt_template: str) -> dict[str, str]:
+    """Format a Pair into a {prompt, completion} dict for HF Trainer."""
+    trait_json = json.dumps(list(pair.trait_set), separators=(", ", ": "))
+    text = prompt_template.format(description=pair.description, trait_json=trait_json)
+    return {"text": text, "trait_set": list(pair.trait_set)}
+
+
+class TraitTrainer:
+    """HF Trainer wrapper for the trait-decoder fine-tune.
+
+    Construction triggers heavy imports + tokenizer load. Train +
+    eval are compute-heavy and require GPU for non-trivial models.
+    """
+
+    def __init__(self, config: TrainConfig):
+        if not config.label_space:
+            raise ValueError("TrainConfig.label_space must be non-empty")
+        try:
+            import torch  # noqa: F401
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:
+            raise ImportError(
+                "TraitTrainer requires the [model] extra. "
+                "Install via: pip install -e '.[model]'"
+            ) from exc
+
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.model = AutoModelForCausalLM.from_pretrained(config.model_name)
+
+    def _build_dataset(self, pairs: list[Pair]) -> Any:
+        """Build a HF Dataset from Pairs."""
+        from datasets import Dataset
+
+        formatted = [_format_example(p, self.config.prompt_template) for p in pairs]
+        ds = Dataset.from_list(formatted)
+
+        def tokenize(batch: dict[str, list]) -> dict[str, list]:
+            tokens = self.tokenizer(
+                batch["text"],
+                truncation=True,
+                max_length=self.config.max_seq_len,
+                padding=False,
+            )
+            tokens["labels"] = tokens["input_ids"].copy()
+            return tokens
+
+        return ds.map(tokenize, batched=True, remove_columns=["text", "trait_set"])
+
+    def _build_eval_callback(
+        self, val_pairs: list[Pair], label_space: tuple[str, ...]
+    ) -> Callable[..., dict[str, float]]:
+        """Return a compute_metrics function for HF Trainer that runs the
+        decoder against val pairs and computes F1 macro."""
+        from trait_inference.metrics import f1_macro
+        from trait_inference.model.decoder import TraitDecoder, TraitDecoderConfig
+
+        def compute_metrics(eval_pred: Any) -> dict[str, float]:
+            # NOTE: eval_pred contains predictions from the trainer's
+            # default forward pass — but we want to evaluate via
+            # constrained decoding, which is a separate inference path.
+            # So we run the decoder ourselves here over val_pairs.
+            decoder_cfg = TraitDecoderConfig(
+                model_name=self.config.model_name,
+                label_space=label_space,
+                device="cuda" if self.config.fp16 else "cpu",
+            )
+            decoder = TraitDecoder(decoder_cfg)
+            descriptions = [p.description for p in val_pairs]
+            preds = decoder.predict_batch(descriptions)
+            gold = [list(p.trait_set) for p in val_pairs]
+            return {
+                "val_f1_macro": f1_macro(gold, preds, label_space=label_space),
+            }
+
+        return compute_metrics
+
+    def train(
+        self,
+        train_pairs: list[Pair],
+        val_pairs: list[Pair],
+    ) -> dict[str, Any]:
+        """Run the full training loop. Returns summary dict.
+
+        Saves checkpoint to `config.output_dir` + structured training
+        log to `output_dir/train_log.json`.
+        """
+        from transformers import (
+            DataCollatorForLanguageModeling,
+            EarlyStoppingCallback,
+            Trainer,
+            TrainingArguments,
+        )
+
+        train_ds = self._build_dataset(train_pairs)
+        val_ds = self._build_dataset(val_pairs)
+
+        training_args = TrainingArguments(
+            output_dir=self.config.output_dir,
+            num_train_epochs=self.config.num_epochs,
+            per_device_train_batch_size=self.config.train_batch_size,
+            per_device_eval_batch_size=self.config.eval_batch_size,
+            learning_rate=self.config.learning_rate,
+            warmup_steps=self.config.warmup_steps,
+            weight_decay=self.config.weight_decay,
+            logging_steps=self.config.logging_steps,
+            eval_strategy="steps",
+            eval_steps=self.config.eval_steps,
+            save_strategy="steps",
+            save_steps=self.config.save_steps,
+            save_total_limit=2,
+            load_best_model_at_end=True,
+            metric_for_best_model="val_f1_macro",
+            greater_is_better=True,
+            fp16=self.config.fp16,
+            seed=self.config.seed,
+            gradient_accumulation_steps=self.config.grad_accumulation_steps,
+            report_to="none",
+        )
+
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.tokenizer, mlm=False
+        )
+
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_ds,
+            eval_dataset=val_ds,
+            data_collator=data_collator,
+            compute_metrics=self._build_eval_callback(val_pairs, self.config.label_space),
+            callbacks=[EarlyStoppingCallback(
+                early_stopping_patience=self.config.early_stopping_patience
+            )],
+        )
+
+        train_output = trainer.train()
+        eval_metrics = trainer.evaluate()
+
+        # Persist structured training log + config
+        log_path = Path(self.config.output_dir) / "train_log.json"
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        log_path.write_text(json.dumps({
+            "config": self.config.to_dict(),
+            "train_runtime_seconds": train_output.metrics.get("train_runtime"),
+            "train_loss": train_output.metrics.get("train_loss"),
+            "best_val_f1_macro": eval_metrics.get("eval_val_f1_macro"),
+            "completed_at": datetime.now(timezone.utc).isoformat(),
+        }, indent=2), encoding="utf-8")
+
+        return {
+            "train_runtime_seconds": train_output.metrics.get("train_runtime"),
+            "train_loss": train_output.metrics.get("train_loss"),
+            "best_val_f1_macro": eval_metrics.get("eval_val_f1_macro"),
+            "checkpoint_dir": self.config.output_dir,
+        }