From 59add8df88c8a1453f36980fbab13565a941d303 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:13:46 -0700 Subject: [PATCH 1/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- tests/shadow/test_hyperlexical_early_stop.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/shadow/test_hyperlexical_early_stop.py diff --git a/tests/shadow/test_hyperlexical_early_stop.py b/tests/shadow/test_hyperlexical_early_stop.py new file mode 100644 index 00000000..e69de29b From 33ef7d7ada35218bccbc16f2db7ab0cf40332654 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:15:27 -0700 Subject: [PATCH 2/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- tests/shadow/test_hyperlexical_early_stop.py | 462 +++++++++++++++++++ 1 file changed, 462 insertions(+) diff --git a/tests/shadow/test_hyperlexical_early_stop.py b/tests/shadow/test_hyperlexical_early_stop.py index e69de29b..b6557366 100644 --- a/tests/shadow/test_hyperlexical_early_stop.py +++ b/tests/shadow/test_hyperlexical_early_stop.py @@ -0,0 +1,462 @@ +"""Default-off early stop for the classify training loop. No sleeps.""" + +from __future__ import annotations + +import ast +import inspect +from pathlib import Path +import sys + +import pytest + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / "scripts" / "shadow")) +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +from hyperlexical.classify_metrics import SELECT_METRIC_CLASSIFY, SELECT_METRIC_UNBIND +from hyperlexical.loop import ( + EARLY_STOP_ENV, + EARLY_STOP_MIN_EPOCHS_ENV, + EARLY_STOP_PATIENCE_ENV, + STOP_REASON_EARLY_STOPPING, + STOP_REASON_MAX_EPOCHS, + early_stop_break, + epoch_timing_fields, + monotonic_seconds, + note_strict_improvement, + resolve_early_stop_config, + round_seconds, +) +import hyperlexical.loop as loop + +# Best classify score at epoch 3. Later values never strictly exceed it. +CANDIDATE = [0.50, 0.40, 0.45, 0.70, 0.55, 0.60, 0.58, 0.69, 0.10, 0.11, 0.12, 0.13] + + +def _enable(monkeypatch, *, patience="4", minimum_epochs="4"): + monkeypatch.setenv(EARLY_STOP_ENV, "1") + monkeypatch.setenv(EARLY_STOP_PATIENCE_ENV, patience) + monkeypatch.setenv(EARLY_STOP_MIN_EPOCHS_ENV, minimum_epochs) + + +def _disable(monkeypatch): + monkeypatch.delenv(EARLY_STOP_ENV, raising=False) + monkeypatch.delenv(EARLY_STOP_PATIENCE_ENV, raising=False) + monkeypatch.delenv(EARLY_STOP_MIN_EPOCHS_ENV, raising=False) + + +def run_classify_schedule( + scores: list[float], + *, + max_epochs: int, + enabled: bool, + patience: int | None = None, + minimum_epochs: int | None = None, +) -> dict[str, object]: + """Same order as the trainer: score, strict improvement, then stop check.""" + best_value = float("-inf") + best_epoch = None + restored_epoch = None + stop_reason = STOP_REASON_MAX_EPOCHS + scored = 0 + last_epoch = None + for ep in range(max_epochs): + if ep >= len(scores): + raise AssertionError(f"missing score for epoch {ep}") + best_value, best_epoch, improved = note_strict_improvement( + best_value, best_epoch, scores[ep], ep + ) + if improved: + restored_epoch = ep + scored = ep + 1 + last_epoch = ep + if early_stop_break( + enabled=enabled, + epoch_index=ep, + best_epoch=best_epoch, + epochs_scored=scored, + minimum_epochs=minimum_epochs, + patience=patience, + max_epochs=max_epochs, + ): + stop_reason = STOP_REASON_EARLY_STOPPING + break + return { + "stop_reason": stop_reason, + "epochs_scored": scored, + "last_epoch": last_epoch, + "best_epoch": best_epoch, + "restored_epoch": restored_epoch, + "best_value": best_value, + } + + +def test_default_off_runs_the_full_schedule(monkeypatch): + _disable(monkeypatch) + monkeypatch.setenv("HYPERLEX_TRAIN_EPOCHS", "12") + monkeypatch.setenv(EARLY_STOP_PATIENCE_ENV, "-1") + cfg = resolve_early_stop_config(max_epochs=12, select_metric=SELECT_METRIC_CLASSIFY) + assert cfg.enabled is False + result = run_classify_schedule( + CANDIDATE, + max_epochs=12, + enabled=cfg.enabled, + patience=4, + minimum_epochs=4, + ) + assert result["epochs_scored"] == 12 + assert result["last_epoch"] == 11 + assert result["stop_reason"] == STOP_REASON_MAX_EPOCHS + assert result["best_epoch"] == 3 + + +def test_train_epochs_alone_does_not_enable_early_stop(monkeypatch): + _disable(monkeypatch) + monkeypatch.setenv("HYPERLEX_TRAIN_EPOCHS", "12") + cfg = resolve_early_stop_config(max_epochs=12, select_metric=SELECT_METRIC_UNBIND) + assert cfg.enabled is False + assert cfg.patience is None + assert cfg.minimum_epochs is None + + +def test_minimum_epochs_block_an_earlier_patience_fire(): + scores = [0.90, 0.10, 0.10, 0.10, 0.10, 0.10] + result = run_classify_schedule( + scores, + max_epochs=6, + enabled=True, + patience=1, + minimum_epochs=4, + ) + assert result["best_epoch"] == 0 + assert result["last_epoch"] == 3 + assert result["epochs_scored"] == 4 + assert result["stop_reason"] == STOP_REASON_EARLY_STOPPING + + +def test_best_at_epoch_3_stops_after_epoch_7_with_patience_4(): + result = run_classify_schedule( + CANDIDATE, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert result["best_epoch"] == 3 + assert result["last_epoch"] == 7 + assert result["epochs_scored"] == 8 + assert result["stop_reason"] == STOP_REASON_EARLY_STOPPING + assert result["restored_epoch"] == 3 + + +def test_strict_improvement_resets_patience(): + scores = list(CANDIDATE) + scores[5] = 0.80 + result = run_classify_schedule( + scores, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert result["best_epoch"] == 5 + assert result["last_epoch"] == 9 + assert result["epochs_scored"] == 10 + assert result["stop_reason"] == STOP_REASON_EARLY_STOPPING + + +def test_ties_keep_the_earlier_checkpoint_and_consume_patience(): + scores = list(CANDIDATE) + scores[7] = scores[3] + assert scores[7] == scores[3] + result = run_classify_schedule( + scores, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert result["best_epoch"] == 3 + assert result["restored_epoch"] == 3 + assert result["last_epoch"] == 7 + assert result["epochs_scored"] == 8 + assert result["best_value"] == scores[3] + tied = [0.50, 0.50] + short = run_classify_schedule( + tied, + max_epochs=4, + enabled=True, + patience=1, + minimum_epochs=1, + ) + assert short["best_epoch"] == 0 + assert short["restored_epoch"] == 0 + assert short["last_epoch"] == 1 + assert short["stop_reason"] == STOP_REASON_EARLY_STOPPING + + +def test_improvement_on_the_would_be_stopping_epoch_prevents_stopping(): + scores = list(CANDIDATE) + [0.10, 0.10, 0.10, 0.10] + scores[7] = 0.71 + stopped_without = run_classify_schedule( + CANDIDATE, + max_epochs=16, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert stopped_without["last_epoch"] == 7 + result = run_classify_schedule( + scores, + max_epochs=16, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert result["best_epoch"] == 7 + assert result["restored_epoch"] == 7 + assert result["last_epoch"] == 11 + assert result["epochs_scored"] == 12 + assert result["stop_reason"] == STOP_REASON_EARLY_STOPPING + + +def test_max_epoch_cap_wins_when_reached_first(): + early = run_classify_schedule( + CANDIDATE, + max_epochs=6, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert early["epochs_scored"] == 6 + assert early["last_epoch"] == 5 + assert early["best_epoch"] == 3 + assert early["stop_reason"] == STOP_REASON_MAX_EPOCHS + # Patience would also be true on the final configured epoch. The cap wins. + on_the_cap = run_classify_schedule( + CANDIDATE, + max_epochs=8, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert on_the_cap["last_epoch"] == 7 + assert on_the_cap["epochs_scored"] == 8 + assert on_the_cap["stop_reason"] == STOP_REASON_MAX_EPOCHS + assert on_the_cap["restored_epoch"] == 3 + + +def test_best_checkpoint_restored_after_early_stop(): + result = run_classify_schedule( + CANDIDATE, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert result["stop_reason"] == STOP_REASON_EARLY_STOPPING + assert result["restored_epoch"] == result["best_epoch"] == 3 + assert result["restored_epoch"] != result["last_epoch"] + src = Path(loop.__file__).read_text(encoding="utf-8") + tree = ast.parse(src) + fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_loop") + break_line = None + restore_line = None + for node in ast.walk(fn): + if isinstance(node, ast.Call) and getattr(node.func, "id", None) == "early_stop_break": + break_line = node.lineno + if isinstance(node, ast.BoolOp): + text = ast.get_source_segment(src, node) or "" + if "best_state is not None" in text: + restore_line = node.lineno + assert break_line is not None and restore_line is not None + assert break_line < restore_line + + +@pytest.mark.parametrize( + ("env", "metric", "match"), + [ + ({"patience": "-1", "minimum": "4"}, SELECT_METRIC_CLASSIFY, "PATIENCE"), + ({"patience": "4", "minimum": "0"}, SELECT_METRIC_CLASSIFY, "MIN_EPOCHS"), + ({"patience": "4", "minimum": "-2"}, SELECT_METRIC_CLASSIFY, "MIN_EPOCHS"), + ({"patience": "4", "minimum": "13"}, SELECT_METRIC_CLASSIFY, "exceeds"), + ({"patience": "4.5", "minimum": "4"}, SELECT_METRIC_CLASSIFY, "integer"), + ({"patience": "4", "minimum": "nope"}, SELECT_METRIC_CLASSIFY, "integer"), + ({"patience": None, "minimum": "4"}, SELECT_METRIC_CLASSIFY, "required"), + ({"patience": "4", "minimum": None}, SELECT_METRIC_CLASSIFY, "required"), + ({"patience": "4", "minimum": "4"}, SELECT_METRIC_UNBIND, "classify_macro_f1_nonnone"), + ], +) +def test_invalid_configuration_fails_before_optimizer(monkeypatch, env, metric, match): + monkeypatch.setenv(EARLY_STOP_ENV, "1") + if env["patience"] is None: + monkeypatch.delenv(EARLY_STOP_PATIENCE_ENV, raising=False) + else: + monkeypatch.setenv(EARLY_STOP_PATIENCE_ENV, env["patience"]) + if env["minimum"] is None: + monkeypatch.delenv(EARLY_STOP_MIN_EPOCHS_ENV, raising=False) + else: + monkeypatch.setenv(EARLY_STOP_MIN_EPOCHS_ENV, env["minimum"]) + constructed: list[object] = [] + + def adamw(*args, **kwargs): + constructed.append(args) + raise AssertionError("optimizer constructed") + + with pytest.raises(ValueError, match=match): + resolve_early_stop_config(max_epochs=12, select_metric=metric) + adamw([object()], lr=2e-5) + assert constructed == [] + + src = Path(loop.__file__).read_text(encoding="utf-8") + tree = ast.parse(src) + fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_loop") + resolve_line = None + adam_line = None + for node in ast.walk(fn): + if not isinstance(node, ast.Call): + continue + name = getattr(node.func, "id", None) or getattr(node.func, "attr", None) + if name == "resolve_early_stop_config": + resolve_line = node.lineno + if name == "AdamW": + adam_line = node.lineno + assert resolve_line is not None and adam_line is not None + assert resolve_line < adam_line + + +def test_unknown_early_stop_token_is_rejected(monkeypatch): + monkeypatch.setenv(EARLY_STOP_ENV, "maybe") + with pytest.raises(ValueError, match="HYPERLEX_EARLY_STOP"): + resolve_early_stop_config(max_epochs=12, select_metric=SELECT_METRIC_CLASSIFY) + + +def test_explicit_off_runs_full_schedule(monkeypatch): + monkeypatch.setenv(EARLY_STOP_ENV, "off") + monkeypatch.setenv(EARLY_STOP_PATIENCE_ENV, "1") + monkeypatch.setenv(EARLY_STOP_MIN_EPOCHS_ENV, "1") + cfg = resolve_early_stop_config(max_epochs=12, select_metric=SELECT_METRIC_CLASSIFY) + assert cfg.enabled is False + result = run_classify_schedule( + CANDIDATE, + max_epochs=cfg.max_epochs, + enabled=cfg.enabled, + ) + assert result["epochs_scored"] == 12 + assert result["stop_reason"] == STOP_REASON_MAX_EPOCHS + + +def test_wallclock_fields_are_nonnegative_and_cumulative(monkeypatch): + ticks = iter([0.0, 1.0, 1.25, 1.25, 2.50, 2.50, 2.75]) + + def clock() -> float: + return next(ticks) + + monkeypatch.setattr(loop, "monotonic_seconds", clock) + training_started = loop.monotonic_seconds() + rows = [] + for epoch in range(3): + epoch_started = loop.monotonic_seconds() + row = { + "epoch": epoch, + "unbind_exact": 0.0, + "saved_best": False, + "classify_acc": 0.5, + } + row.update( + epoch_timing_fields( + epoch_started=epoch_started, + epoch_ended=loop.monotonic_seconds(), + training_started=training_started, + ) + ) + rows.append(row) + elapsed = [row["training_elapsed_seconds"] for row in rows] + assert elapsed == sorted(elapsed) + assert all(later >= earlier for earlier, later in zip(elapsed, elapsed[1:])) + for row in rows: + assert row["epoch_wallclock_seconds"] >= 0 + assert row["training_elapsed_seconds"] >= 0 + assert "epoch" in row and "saved_best" in row + assert rows[0]["epoch_wallclock_seconds"] == 0.25 + assert rows[1]["epoch_wallclock_seconds"] == 1.25 + assert rows[2]["epoch_wallclock_seconds"] == 0.25 + assert rows[-1]["training_elapsed_seconds"] == 2.75 + assert round_seconds(1.23456789) == round(1.23456789, 6) + assert inspect.signature(epoch_timing_fields).parameters.keys() >= { + "epoch_started", + "epoch_ended", + "training_started", + } + + +def test_wallclock_values_do_not_affect_selection(monkeypatch): + def boom() -> float: + raise AssertionError("selection read the clock") + + monkeypatch.setattr(loop, "monotonic_seconds", boom) + first = run_classify_schedule( + CANDIDATE, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + second = run_classify_schedule( + CANDIDATE, + max_epochs=12, + enabled=True, + patience=4, + minimum_epochs=4, + ) + assert first == second + assert first["best_epoch"] == 3 + assert "epoch_wallclock_seconds" not in inspect.signature(note_strict_improvement).parameters + assert "training_elapsed_seconds" not in inspect.signature(early_stop_break).parameters + stamped = epoch_timing_fields(epoch_started=0.0, epoch_ended=50.0, training_started=0.0) + assert stamped["training_elapsed_seconds"] == 50.0 + assert first["best_epoch"] == 3 + + +def test_disabled_selection_keeps_strict_increase_and_earlier_ties(): + scores = [0.20, 0.50, 0.50, 0.40, 0.80, 0.80] + best_value = float("-inf") + best_epoch = None + improved_flags = [] + for epoch, score in enumerate(scores): + best_value, best_epoch, improved = note_strict_improvement( + best_value, best_epoch, score, epoch + ) + improved_flags.append(improved) + assert improved_flags == [True, True, False, False, True, False] + assert best_epoch == 4 + result = run_classify_schedule( + scores, + max_epochs=len(scores), + enabled=False, + ) + assert result["epochs_scored"] == len(scores) + assert result["stop_reason"] == STOP_REASON_MAX_EPOCHS + assert result["best_epoch"] == 4 + assert result["restored_epoch"] == 4 + + +def test_completion_receipt_records_only_loop_stop_reasons(): + src = Path(loop.__file__).read_text(encoding="utf-8") + assert '"stop_reason": stop_reason' in src + assert '"training_elapsed_seconds": training_elapsed_seconds' in src + tree = ast.parse(src) + fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_loop") + assigned: list[str] = [] + for node in ast.walk(fn): + if not isinstance(node, ast.Assign): + continue + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "stop_reason": + assigned.append(ast.get_source_segment(src, node.value) or "") + assert assigned == ["STOP_REASON_MAX_EPOCHS", "STOP_REASON_EARLY_STOPPING"] + assert STOP_REASON_MAX_EPOCHS == "max_epochs" + assert STOP_REASON_EARLY_STOPPING == "early_stopping" + assert "stop_reason" not in inspect.signature(monotonic_seconds).parameters From 651e096d6fd3a0bbb4113a0e1a075aa7f9629e48 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:16:34 -0700 Subject: [PATCH 3/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- specs/007-hyperlexical-model/AARON-SPARK-TRAIN.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/specs/007-hyperlexical-model/AARON-SPARK-TRAIN.md b/specs/007-hyperlexical-model/AARON-SPARK-TRAIN.md index f5a276ff..3b9c1fe5 100644 --- a/specs/007-hyperlexical-model/AARON-SPARK-TRAIN.md +++ b/specs/007-hyperlexical-model/AARON-SPARK-TRAIN.md @@ -70,6 +70,21 @@ export HYPERLEX_OFFLINE=1 export TOKENIZERS_PARALLELISM=false export HYPERLEX_TRAIN_OUT="$HOME/.hyperlex/models/hyperlex-encoder-modernbert-base-seed" export HYPERLEX_TRAIN_EPOCHS=2 +# HYPERLEX_TRAIN_EPOCHS is the max-epoch cap. It does not enable early stopping. +# Early stop is optional and default-off. Turn it on only with the classify +# selection metric. Improvement is a strict increase; ties keep the earlier +# checkpoint. After the epoch is scored, stop when +# epoch_index - best_epoch >= patience and at least MIN_EPOCHS epochs have +# been scored (0-based epoch_index). The best checkpoint is still restored. +# epoch-progress.jsonl gains observational seconds rounded to 6 decimal places +# (epoch_wallclock_seconds, training_elapsed_seconds). They do not affect +# selection or stopping. A completed train-receipt.json records stop_reason +# (max_epochs or early_stopping) and training_elapsed_seconds. A failed loop +# still raises and does not write that completion receipt. +# export HLX_SELECT_METRIC=classify_macro_f1_nonnone +# export HYPERLEX_EARLY_STOP=1 +# export HYPERLEX_EARLY_STOP_PATIENCE=4 +# export HYPERLEX_EARLY_STOP_MIN_EPOCHS=4 export HYPERLEX_TRAIN_BATCH=8 export HYPERLEX_TRAIN_LR=2e-5 # optional recipe bump (default 2, clamp 1..min(encoder layers, 8)): From 99ba7a1a513001a26def9587625c916aaf8a8b95 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:16:54 -0700 Subject: [PATCH 4/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- scripts/shadow/hyperlexical/loop.py | 1106 +-------------------------- 1 file changed, 1 insertion(+), 1105 deletions(-) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index fba5fade..311c8dd0 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -1,1105 +1 @@ -"""Spark train loop. Gate only.""" - -from __future__ import annotations - -import json -import os -from pathlib import Path - -from .align import atom_token_index, offsets_from_tokenizer, pool_indices -from .export import export_dataset, repo_root, write_export -from .admission import AdmissionError, admit_training_run -from .train_input import train_input_receipt -from .classify_metrics import ( - NONE_LABEL, - SELECT_METRIC_CLASSIFY, - macro_f1_nonnone, - none_false_positive_rate, - resolve_select_metric, -) -from .classify_split import apply_classify_split_file -from .seed_control import apply_training_seed -from .force_train_overlap import enforce_force_train_disjoint -from .layout import ( - FAMILIES, - HIDDEN, - MAX_LEN, - MODEL_ID_SEED, - TRUNK, - UNK, - describe, - label_maps_for_splits, - resolve_last_trainable, - resolve_vocab_train_only, -) -from .eval_forward import apply_encoder_trainable -from .filler_filter import assert_publishable_vocab, filter_mode, filter_unbind_rows -from .holdout_guard import ( - assert_no_holdout, - filter_holdout_rows, - holdout_receipt, - load_holdout_spec, - log_holdout, -) -from .provenance import provenance -from .release_set import maybe_release -from .save_pretrained import ( - collect_encoder_trainable, - save_heads, - split_weight_tensors, - write_skeleton, -) -from .training_routing import route_rows -from .unbind_curriculum import ( - plan_unbind_curriculum, - resolve_curriculum_schedule, - select_unbind_for_epoch, -) -from .unbind_metrics import mapped_filler, mapped_pred, summarize_unbind_pairs -from .unbind_recipe import ( - apply_unbind_force_train, - resolve_unbind_inferred_weight, - resolve_unbind_morph_margin, - shape_unbind_train, - unbind_row_sample_weight, -) -from .unbind_head_slot import ( - apply_head_slot_weight, - resolve_unbind_head_slot_weight, -) -from .unbind_second_slot import ( - apply_second_slot_weight, - resolve_unbind_second_slot_weight, -) -from .unbind_residual import ( - residual_row_record, - resolve_unbind_residual_dump_path, - write_residual_dump, -) -from .unbind_slot_ce import ( - UNBIND_SLOT_CE_AUX_LAMBDA, - combine_unbind_train_terms, - resolve_unbind_primary_mode, -) - -UNBIND_LOSS_WEIGHT_ENV = "HYPERLEX_UNBIND_LOSS_WEIGHT" -UNBIND_EVERY_N_ENV = "HYPERLEX_UNBIND_EVERY_N" -SAVE_BEST_UNBIND_ENV = "HYPERLEX_SAVE_BEST_UNBIND" -INIT_FROM_ENV = "HYPERLEX_INIT_FROM" -INIT_EXPAND_VOCAB_ENV = "HYPERLEX_INIT_EXPAND_VOCAB" -UNBIND_LOSS_WEIGHT_DEFAULT = 1.0 -UNBIND_EVERY_N_DEFAULT = 1 - - -def resolve_save_best_unbind(raw: str | None = None) -> bool: - """When true, persist best-by-val-unbind_exact weights as primary model.safetensors. - - morph35 peak-not-saved: epoch_metrics recorded ep16 0.5372 but only final - weights were written. Opt-in via HYPERLEX_SAVE_BEST_UNBIND=1. - """ - if raw is None: - raw = os.environ.get(SAVE_BEST_UNBIND_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return False - return str(raw).strip().lower() in {"1", "true", "yes", "on"} - - -def resolve_init_from(raw: str | None = None) -> Path | None: - """Optional warm-start dir with model.safetensors (or heads.pt). - - Opt-in via HYPERLEX_INIT_FROM=/path/to/prior seed (e.g. morph36 best). - """ - if raw is None: - raw = os.environ.get(INIT_FROM_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return None - path = Path(str(raw).strip()).expanduser() - if not path.is_dir(): - raise ValueError(f"{INIT_FROM_ENV} must be an existing directory, got {path}") - return path - - -def resolve_init_expand_vocab(raw: str | None = None) -> bool: - """When true, warm-load remaps shared role/filler rows into expanded vocabs. - - Default fail-closed on vocab mismatch. Opt-in via HYPERLEX_INIT_EXPAND_VOCAB=1 - so a smaller prior seed (e.g. morph65 max pos_5) can warm a harvest that added - pos_6+/new fillers: copy overlapping labels by name, leave new rows at init. - """ - if raw is None: - raw = os.environ.get(INIT_EXPAND_VOCAB_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return False - return str(raw).strip().lower() in {"1", "true", "yes", "on"} - - -def _init_weight_path(init_dir: Path) -> Path: - for name in ("model.safetensors", "heads.pt"): - candidate = init_dir / name - if candidate.is_file(): - return candidate - raise FileNotFoundError( - f"{INIT_FROM_ENV}={init_dir} missing model.safetensors or heads.pt" - ) - - -def _read_init_vocabs(init_dir: Path) -> tuple[list | None, list | None]: - for name in ("layout.json", "config.json"): - path = init_dir / name - if not path.is_file(): - continue - try: - blob = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - if not isinstance(blob, dict): - continue - role = blob.get("role_vocab") - filler = blob.get("filler_vocab") - if isinstance(role, list) and isinstance(filler, list) and filler: - return role, filler - return None, None - - -def _remap_linear_rows(module, state: dict, init_labels: list, current_labels: list, head: str) -> dict: - """Copy overlapping out-rows from init Linear state into current module by label name.""" - weight = state.get("weight") - if weight is None: - raise ValueError(f"{INIT_FROM_ENV} {head} missing weight") - if int(getattr(weight, "shape", [0])[0]) != len(init_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} weight rows {tuple(weight.shape)} != " - f"init vocab {len(init_labels)}" - ) - if int(module.weight.shape[0]) != len(current_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} module rows {tuple(module.weight.shape)} != " - f"current vocab {len(current_labels)}" - ) - current_of = {lab: i for i, lab in enumerate(current_labels)} - mapped = 0 - skipped = 0 - for ii, lab in enumerate(init_labels): - ci = current_of.get(lab) - if ci is None: - skipped += 1 - continue - module.weight.data[ci].copy_(weight[ii].detach()) - mapped += 1 - bias = state.get("bias") - if bias is not None and module.bias is not None: - if int(getattr(bias, "shape", [0])[0]) != len(init_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} bias rows != init vocab {len(init_labels)}" - ) - for ii, lab in enumerate(init_labels): - ci = current_of.get(lab) - if ci is None: - continue - module.bias.data[ci].copy_(bias[ii].detach()) - if mapped == 0: - raise ValueError( - f"{INIT_FROM_ENV} {head} expand remap matched 0/{len(init_labels)} labels" - ) - return { - "mapped": mapped, - "skipped_init_only": skipped, - "new_current_rows": len(current_labels) - mapped, - "init_n": len(init_labels), - "current_n": len(current_labels), - } - - -def warm_load_checkpoint( - encoder, - classify, - role_head, - filler_head, - maps: dict, - init_dir: Path, - *, - expand_vocab: bool | None = None, -) -> dict: - """Load heads + trainable encoder tensors from a prior seed dump. Fail closed. - - When expand_vocab is true (or HYPERLEX_INIT_EXPAND_VOCAB=1), role/filler heads - remap overlapping labels by name into the current larger vocab; classify still - loads strict. Default remains exact-vocab match. - """ - if expand_vocab is None: - expand_vocab = resolve_init_expand_vocab() - weight_path = _init_weight_path(init_dir) - init_roles, init_fillers = _read_init_vocabs(init_dir) - cur_roles = list(maps.get("role_vocab") or []) - cur_fillers = list(maps.get("filler_vocab") or []) - roles_match = init_roles is None or init_roles == cur_roles - fillers_match = init_fillers is None or init_fillers == cur_fillers - vocab_match = roles_match and fillers_match - if not vocab_match and not expand_vocab: - if init_roles is not None and init_roles != cur_roles: - raise ValueError( - f"{INIT_FROM_ENV} role_vocab mismatch vs current export " - f"(init={len(init_roles)} current={len(cur_roles)})" - ) - raise ValueError( - f"{INIT_FROM_ENV} filler_vocab mismatch vs current export " - f"(init={len(init_fillers or [])} current={len(cur_fillers)})" - ) - if not vocab_match and expand_vocab: - if init_roles is None or init_fillers is None: - raise ValueError( - f"{INIT_EXPAND_VOCAB_ENV}=1 requires init role_vocab+filler_vocab " - f"in {init_dir}/config.json (or layout.json)" - ) - - if weight_path.name == "model.safetensors": - from safetensors.torch import load_file - - split = split_weight_tensors(load_file(str(weight_path), device="cpu")) - heads_blob = None - else: - import torch - - try: - blob = torch.load(str(weight_path), map_location="cpu", weights_only=False) - except TypeError: - blob = torch.load(str(weight_path), map_location="cpu") - if not isinstance(blob, dict): - raise ValueError(f"{weight_path} is not a heads state dict") - split = { - "classify": blob.get("classify") or {}, - "role_head": blob.get("role_head") or {}, - "filler_head": blob.get("filler_head") or {}, - "encoder": { - k: v - for k, v in ( - {} if not isinstance(blob.get("encoder"), dict) else blob["encoder"] - ).items() - }, - } - # normalize encoder keys to encoder.* for apply_encoder_trainable - enc = {} - for k, v in split["encoder"].items(): - key = str(k) - enc[key if key.startswith("encoder.") else f"encoder.{key}"] = v - split["encoder"] = enc - heads_blob = blob - - classify_state = split.get("classify") or {} - if not classify_state: - raise ValueError(f"{weight_path} missing classify tensors") - classify.load_state_dict(classify_state, strict=True) - - expand_receipt: dict = { - "expand_vocab": bool(expand_vocab and not vocab_match), - "vocab_match": vocab_match, - } - if vocab_match or not expand_vocab: - for name, module in (("role_head", role_head), ("filler_head", filler_head)): - state = split.get(name) or {} - if not state: - raise ValueError(f"{weight_path} missing {name} tensors") - module.load_state_dict(state, strict=True) - else: - role_state = split.get("role_head") or {} - filler_state = split.get("filler_head") or {} - if not role_state or not filler_state: - raise ValueError(f"{weight_path} missing role_head/filler_head tensors") - expand_receipt["role"] = _remap_linear_rows( - role_head, role_state, list(init_roles), cur_roles, "role_head" - ) - expand_receipt["filler"] = _remap_linear_rows( - filler_head, filler_state, list(init_fillers), cur_fillers, "filler_head" - ) - - applied = apply_encoder_trainable(encoder, split.get("encoder") or {}) - if applied["present"] and applied["loaded"] == 0: - raise ValueError( - f"{INIT_FROM_ENV} encoder tensors present but none matched trunk keys" - ) - return { - "init_from": str(init_dir), - "weight_file": weight_path.name, - "encoder_trainable_loaded": applied["loaded"], - "encoder_trainable_present": applied["present"], - "heads_blob": bool(heads_blob), - "init_expand_vocab": bool(expand_vocab), - **expand_receipt, - } - - -def _cpu_module_state(module) -> dict: - return {k: v.detach().cpu().contiguous() for k, v in module.state_dict().items()} - - -def _build_weight_state(encoder, classify, role_head, filler_head, maps, layout) -> dict: - return { - "classify": _cpu_module_state(classify), - "role_head": _cpu_module_state(role_head), - "filler_head": _cpu_module_state(filler_head), - "encoder": collect_encoder_trainable(encoder), - "maps": {k: v for k, v in maps.items() if k not in {"family_of", "role_of", "filler_of"}}, - "layout": layout, - } - - -def resolve_unbind_loss_weight(raw: str | float | int | None = None) -> float: - """Scale on unbind loss before backward. Default 1.0. Fail-closed if invalid.""" - if raw is None: - raw = os.environ.get(UNBIND_LOSS_WEIGHT_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return UNBIND_LOSS_WEIGHT_DEFAULT - try: - weight = float(str(raw).strip()) - except (TypeError, ValueError) as exc: - raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") from exc - if weight < 0 or weight != weight or weight == float("inf"): - raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") - return weight - - -def resolve_unbind_every_n(raw: str | int | None = None) -> int: - """Classify-batch stride for an extra unbind step. Default 1 = epoch-end only.""" - if raw is None: - raw = os.environ.get(UNBIND_EVERY_N_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return UNBIND_EVERY_N_DEFAULT - try: - n = int(str(raw).strip(), 10) - except (TypeError, ValueError) as exc: - raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {raw!r}") from exc - if n < 1: - raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {n}") - return n - - -TASK_ROUTING_ENV = "HYPERLEX_TASK_ROUTING" -TASK_ROUTINGS = ("route_rows", "legacy_split") - - -def task_routing(raw: str | None = None) -> str: - """``route_rows`` (default) or ``legacy_split`` (pre-routing loop; reproduces morph75–78). - - ``legacy_split`` selects rows by ``task == "classify"`` / ``task == "unbind"`` only, so - ``classify+unbind`` rows are not trained. Unknown values fail closed. - """ - value = (os.environ.get(TASK_ROUTING_ENV, "") if raw is None else raw).strip() or "route_rows" - if value not in TASK_ROUTINGS: - raise ValueError(f"{TASK_ROUTING_ENV} must be one of {TASK_ROUTINGS}") - return value - - -def should_interleave_unbind(classify_batch_index: int, every_n: int) -> bool: - """True after classify batch `index` (0-based) when every_n > 1.""" - if every_n <= 1: - return False - return (classify_batch_index + 1) % every_n == 0 - - -def prepare_unbind_splits(rows: list) -> tuple[list, list, dict]: - """Train recipe after route_rows. Val frozen unless force-train env is set. - - ``HYPERLEX_UNBIND_FORCE_TRAIN_PATH`` may move authorized OBSERVED exacts - from val→train (accept-style). Empty/unset → val untouched. - ``HLX_HOLDOUT_MANIFESTS`` rows are dropped before that move and before - hard-atom copies, then checked again so neither injection can put them back. - """ - spec = load_holdout_spec() - if task_routing() == "legacy_split": - train = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "train"] - val = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "val"] - else: - routed, _ = route_rows(rows) - train = routed["unbind"]["train"] - val = routed["unbind"]["val"] - train, n_train = filter_holdout_rows(train, spec) - val, n_val = filter_holdout_rows(val, spec) - train, val, force_stats = apply_unbind_force_train(train, val) - train, n_train_injected = filter_holdout_rows(train, spec) - val, n_val_injected = filter_holdout_rows(val, spec) - train, filt_train = filter_unbind_rows(train) - val, filt_val = filter_unbind_rows(val) - shaped, stats = shape_unbind_train(train) - shaped, n_hard = filter_holdout_rows(shaped, spec) - assert_no_holdout(shaped, spec, "unbind train") - assert_no_holdout(val, spec, "unbind val") - stats = { - **stats, - **force_stats, - "filler_filter": filt_train["filler_filter"], - "n_filler_rows_dropped_train": filt_train["n_filler_rows_dropped"], - "n_filler_rows_dropped_val": filt_val["n_filler_rows_dropped"], - "n_holdout_removed_train": n_train + n_train_injected + n_hard, - "n_holdout_removed_val": n_val + n_val_injected, - "holdout_manifests": [dict(item) for item in spec.manifests], - } - return shaped, val, stats - - -def _require_local_model(trunk: Path): - from transformers import AutoModel, AutoTokenizer - - tok = AutoTokenizer.from_pretrained(str(trunk), local_files_only=True) - model = AutoModel.from_pretrained(str(trunk), local_files_only=True) - return tok, model - - -def _layers(encoder): - if hasattr(encoder, "layers"): - return encoder.layers - inner = getattr(encoder, "encoder", None) - if inner is not None and hasattr(inner, "layers"): - return inner.layers - return None - - -def freeze_encoder(encoder, last_trainable: int | None = None) -> tuple[int, int]: - layers = _layers(encoder) - n_layers = len(list(layers)) if layers is not None else None - used = resolve_last_trainable(last_trainable, layer_count=n_layers) - for p in encoder.parameters(): - p.requires_grad = False - n = 0 - if layers is None: - return 0, used - for block in list(layers)[-used:]: - for p in block.parameters(): - p.requires_grad = True - n += p.numel() - return n, used - - -def _offsets(tok, text: str): - try: - return offsets_from_tokenizer(tok, text, max_len=MAX_LEN) - except TypeError: - return None - - -def _enter_training_execution() -> None: - """Reached only after admission. Admission-only mode returns before this.""" - return None - - -def run_loop( - trunk: Path, - out_dir: Path, - *, - include_live: bool = False, - live_store: Path | None = None, -) -> dict: - # Same gates as preflight. Admission-only returns before any optimizer. - admission = admit_training_run( - include_live=include_live, - live_store=live_store, - export_dataset=export_dataset, - trunk=trunk, - out_dir=out_dir, - ) - if os.environ.get("HLX_ADMISSION_ONLY") == "1": - if not admission.ready: - raise AdmissionError(admission.error or "ADMISSION FAIL", admission.receipt) - return admission.receipt - bundle = admission.bundle - holdout_spec = admission.holdout_spec - input_receipt = train_input_receipt(bundle) - disjoint_receipt = admission.disjoint_receipt - reserve_receipt = admission.reserve_receipt - if bundle is None: - raise AdmissionError("ADMISSION FAIL: training bundle was not loaded", admission.receipt) - root = repo_root() - release_rows_, release_stats = maybe_release(bundle["rows"]) - if release_stats["release_set"]: - bundle = {**bundle, "rows": release_rows_} - if any(r.get("role_scheme") == "reviewed_occurrences" for r in bundle["rows"]): - raise ValueError("reviewed occurrences require occurrence-aware loop alignment") - routed, task_accounting = route_rows(bundle["rows"]) - task_accounting = {**task_accounting, "task_routing": task_routing()} - export_dir = Path(os.environ.get("HYPERLEX_EXPORT_DIR") or (root / "specs" / "007-hyperlexical-model" / "exports")) - export_dir.mkdir(parents=True, exist_ok=True) - write_export(export_dir, bundle) - if task_routing() == "legacy_split": - classify_tr = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "train"] - classify_va = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "val"] - else: - classify_tr = routed["classify"]["train"] - classify_va = routed["classify"]["val"] - classify_tr, classify_va, classify_split_receipt = apply_classify_split_file( - classify_tr, classify_va - ) - unbind_tr, unbind_va, unbind_recipe = prepare_unbind_splits(bundle["rows"]) - from .classify_admission import apply_classify_admission - - classify_tr, classify_va, classify_admission_receipt = apply_classify_admission( - bundle["rows"], classify_tr, classify_va - ) - classify_tr, n_classify_train = filter_holdout_rows(classify_tr, holdout_spec) - classify_va, n_classify_val = filter_holdout_rows(classify_va, holdout_spec) - assert_no_holdout(classify_tr, holdout_spec, "classify train") - assert_no_holdout(classify_va, holdout_spec, "classify val") - assert_no_holdout(unbind_tr, holdout_spec, "unbind train") - assert_no_holdout(unbind_va, holdout_spec, "unbind val") - holdout_removed = { - "classify_train": n_classify_train, - "classify_val": n_classify_val, - "unbind_train": unbind_recipe.get("n_holdout_removed_train", 0), - "unbind_val": unbind_recipe.get("n_holdout_removed_val", 0), - } - log_holdout(holdout_spec, holdout_removed) - if len(classify_tr) < 8: - raise RuntimeError("not enough classify train rows") - - select_metric = resolve_select_metric() - select_on_classify = select_metric == SELECT_METRIC_CLASSIFY - if select_on_classify: - selection_rows = list(classify_va or classify_tr[:8]) - else: - selection_rows = list(unbind_va or unbind_tr[:8]) - classify_va, unbind_va, force_overlap = enforce_force_train_disjoint( - classify_va, - unbind_va, - selection_rows, - ) - - _enter_training_execution() - import torch - from torch import nn - from torch.optim import AdamW - - seed_receipt = apply_training_seed(torch) - maps = label_maps_for_splits(unbind_tr, unbind_va) - if filter_mode() == "strict": - assert_publishable_vocab(maps["filler_vocab"]) - tok, encoder = _require_local_model(trunk) - hidden = int(getattr(encoder.config, "hidden_size", HIDDEN)) - if hidden != HIDDEN: - raise RuntimeError(f"hidden {hidden} != {HIDDEN}") - n_unfrozen, last_trainable_used = freeze_encoder(encoder) - classify = nn.Linear(hidden, len(FAMILIES)) - role_head = nn.Linear(hidden, len(maps["role_vocab"])) - filler_head = nn.Linear(hidden, len(maps["filler_vocab"])) - init_from = resolve_init_from() - init_receipt: dict = {"init_from": None, "warm_start": False} - if init_from is not None: - init_receipt = { - "warm_start": True, - **warm_load_checkpoint( - encoder, classify, role_head, filler_head, maps, init_from - ), - } - trainable = [p for p in encoder.parameters() if p.requires_grad] + list(classify.parameters()) + list(role_head.parameters()) + list(filler_head.parameters()) - opt = AdamW(trainable, lr=float(os.environ.get("HYPERLEX_TRAIN_LR", "2e-5"))) - epochs = int(os.environ.get("HYPERLEX_TRAIN_EPOCHS", "2")) - batch = int(os.environ.get("HYPERLEX_TRAIN_BATCH", "8")) - unbind_loss_weight = resolve_unbind_loss_weight() - unbind_every_n = resolve_unbind_every_n() - morph_margin = resolve_unbind_morph_margin() - inferred_weight = resolve_unbind_inferred_weight() - slot_ce_mode = resolve_unbind_primary_mode() - unbind_primary = slot_ce_mode["unbind_primary"] - head_slot_weight = resolve_unbind_head_slot_weight() - second_slot_weight = resolve_unbind_second_slot_weight() - curriculum = resolve_curriculum_schedule() - curriculum_plan = plan_unbind_curriculum(unbind_tr, epochs, curriculum) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - for mod in (encoder, classify, role_head, filler_head): - mod.to(device) - encoder.train() - losses = [] - epoch_metrics = [] - - def encode_texts(texts): - enc = tok(texts, padding=True, truncation=True, max_length=MAX_LEN, return_tensors="pt") - return {k: v.to(device) for k, v in enc.items()} - - def unbind_loss(row): - fillers = list(row.get("fillers") or []) - roles = list(row.get("roles") or []) - if not fillers: - return None - out = encoder(**encode_texts([row["text"]])) - states = out.last_hidden_state[0] - offs = _offsets(tok, row["text"]) - slot_ces = [] - aux_terms = [] - for k, fill in enumerate(fillers): - idxs = pool_indices(states.size(0), atom_token_index(row["text"], fill, offs)) - h = states[idxs].mean(0) - gold_f = maps["filler_of"].get(fill, maps["filler_of"][UNK]) - logits = filler_head(h) - slot_ces.append( - nn.functional.cross_entropy( - logits.unsqueeze(0), torch.tensor([gold_f], device=device) - ) - ) - hard = [nf for nf in (row.get("hard_neg_fillers") or []) if nf] - if hard: - gold_logit = logits[gold_f] - neg_vals = [] - for nf in hard: - ni = maps["filler_of"].get(nf) - if ni is None or ni == gold_f: - continue - neg_vals.append(logits[ni]) - if neg_vals: - stacked = torch.stack(neg_vals) - aux_terms.append(torch.relu(stacked + morph_margin - gold_logit).sum()) - if k < len(roles): - gold_r = maps["role_of"].get(roles[k], maps["role_of"][UNK]) - aux_terms.append( - nn.functional.cross_entropy( - role_head(h).unsqueeze(0), torch.tensor([gold_r], device=device) - ) - ) - weighted_slots = apply_head_slot_weight(slot_ces, head_slot_weight) - weighted_slots = apply_second_slot_weight(weighted_slots, second_slot_weight) - return combine_unbind_train_terms( - weighted_slots, - aux_terms, - primary=unbind_primary, - aux_lambda=UNBIND_SLOT_CE_AUX_LAMBDA, - ) - - # Keep last train loss on-device; avoid per-step .cpu() sync (morph68 hang: - # post-SAVE_BEST host spin at ~98% CPU / GPU util 0 with mem held). - last_train_loss = None - - def step_unbind(row) -> None: - nonlocal last_train_loss - if unbind_loss_weight == 0: - return - uloss = unbind_loss(row) - if uloss is None: - return - row_w = unbind_row_sample_weight(row, inferred_weight) - scaled = uloss * unbind_loss_weight * row_w - opt.zero_grad() - scaled.backward() - opt.step() - last_train_loss = scaled.detach() - - residual_dump_path = resolve_unbind_residual_dump_path() - last_residual_records: list[dict] = [] - save_best_unbind = resolve_save_best_unbind() - best_exact = float("-inf") - best_macro = float("-inf") - best_metrics: dict | None = None - best_state: dict | None = None - best_residual_records: list[dict] = [] - progress_path = out_dir / "epoch-progress.jsonl" - - @torch.no_grad() - def score(): - nonlocal last_residual_records - encoder.eval() - classify.eval() - filler_head.eval() - hit = tot = 0 - classify_golds: list[str] = [] - classify_preds: list[str] = [] - for row in classify_va or classify_tr[:8]: - out = encoder(**encode_texts([row["text"]])) - pred = int(classify(out.last_hidden_state[:, 0]).argmax(-1)[0]) - gold = maps["family_of"].get(row["lineage"], maps["family_of"]["none"]) - hit += int(pred == gold) - tot += 1 - if select_on_classify: - classify_golds.append(FAMILIES[gold] if 0 <= gold < len(FAMILIES) else NONE_LABEL) - classify_preds.append(FAMILIES[pred] if 0 <= pred < len(FAMILIES) else NONE_LABEL) - pairs: list[tuple[list[str], list[str]]] = [] - strict_pairs: list[tuple[list[str], list[str]]] = [] - residual_records: list[dict] = [] - for row in unbind_va or unbind_tr[:8]: - fillers = list(row.get("fillers") or []) - if not fillers: - continue - out = encoder(**encode_texts([row["text"]])) - states = out.last_hidden_state[0] - offs = _offsets(tok, row["text"]) - gold_strs: list[str] = [] - raw_strs: list[str] = [] - pred_strs: list[str] = [] - for fill in fillers: - idxs = pool_indices(states.size(0), atom_token_index(row["text"], fill, offs)) - pred = int(filler_head(states[idxs].mean(0)).argmax()) - gold_strs.append(mapped_filler(maps, fill)) - raw_strs.append(str(fill).lower()) - pred_strs.append(mapped_pred(maps, pred)) - pairs.append((gold_strs, pred_strs)) - strict_pairs.append((raw_strs, pred_strs)) - if residual_dump_path: - rec = residual_row_record( - text=str(row.get("text") or ""), - gold=gold_strs, - pred=pred_strs, - row=row, - ) - if rec is not None: - residual_records.append(rec) - last_residual_records = residual_records - encoder.train() - classify.train() - filler_head.train() - metrics = summarize_unbind_pairs(pairs, strict_pairs=strict_pairs) - metrics["classify_acc"] = hit / max(1, tot) - metrics["n_classify_eval"] = tot - if select_on_classify: - metrics["classify_macro_f1_nonnone"] = macro_f1_nonnone(classify_golds, classify_preds) - metrics["none_fpr"] = none_false_positive_rate(classify_golds, classify_preds) - return metrics - - out_dir.mkdir(parents=True, exist_ok=True) - layout = describe(maps) - layout["last_trainable"] = last_trainable_used - layout["aligner"] = "char_span + offset_mapping" - write_skeleton(out_dir, maps=maps) - - for ep in range(epochs): - phase_rows, phase_meta = select_unbind_for_epoch(unbind_tr, ep, curriculum) - unbind_cycle = 0 - classify_batch_i = 0 - for i in range(0, len(classify_tr), batch): - chunk = classify_tr[i : i + batch] - y = torch.tensor([maps["family_of"].get(c["lineage"], maps["family_of"]["none"]) for c in chunk], device=device) - out = encoder(**encode_texts([c["text"] for c in chunk])) - loss = nn.functional.cross_entropy(classify(out.last_hidden_state[:, 0]), y) - opt.zero_grad() - loss.backward() - opt.step() - last_train_loss = loss.detach() - if should_interleave_unbind(classify_batch_i, unbind_every_n) and phase_rows: - step_unbind(phase_rows[unbind_cycle % len(phase_rows)]) - unbind_cycle += 1 - classify_batch_i += 1 - for row in phase_rows: - step_unbind(row) - metrics = score() - metrics["epoch"] = ep - metrics["unbind_phase"] = phase_meta["phase"] - metrics["n_unbind_phase"] = phase_meta["n_rows"] - metrics["unbind_phase_fallback_full_mix"] = phase_meta["fallback_full_mix"] - epoch_metrics.append(metrics) - exact = float(metrics.get("unbind_exact") or 0.0) - if last_train_loss is not None: - # One host sync per epoch (not per step). - losses.append(float(last_train_loss.item())) - saved_best = False - if select_on_classify: - if metrics.get("classify_macro_f1_nonnone") is None: - raise RuntimeError( - "HLX_SELECT_METRIC=classify_macro_f1_nonnone but val has no " - "non-none gold; macro-F1 is NOT_COMPUTABLE" - ) - score_now = float(metrics["classify_macro_f1_nonnone"]) - if score_now > best_macro: - if device.type == "cuda": - torch.cuda.synchronize() - best_macro = score_now - best_metrics = dict(metrics) - best_state = _build_weight_state( - encoder, classify, role_head, filler_head, maps, layout - ) - best_residual_records = list(last_residual_records) - best_dir = out_dir / "best" - write_skeleton(best_dir, maps=maps) - save_heads(best_dir, best_state) - (best_dir / "best-checkpoint.json").write_text( - json.dumps( - { - "metric": SELECT_METRIC_CLASSIFY, - "epoch": ep, - "classify_macro_f1_nonnone": score_now, - "none_fpr": metrics.get("none_fpr"), - "unbind_exact": exact, - "unbind_token_f1": best_metrics.get("unbind_token_f1"), - "unbind_slot_f1": best_metrics.get("unbind_slot_f1"), - "unbind_exact_strict": best_metrics.get("unbind_exact_strict"), - "unbind_token_f1_strict": best_metrics.get("unbind_token_f1_strict"), - "unbind_slot_f1_strict": best_metrics.get("unbind_slot_f1_strict"), - "classify_acc": best_metrics.get("classify_acc"), - }, - indent=2, - sort_keys=True, - ) - + "\n", - encoding="utf-8", - ) - saved_best = True - if device.type == "cuda": - torch.cuda.synchronize() - torch.cuda.empty_cache() - none_fpr = metrics.get("none_fpr") - none_fpr_text = "null" if none_fpr is None else f"{float(none_fpr):.6f}" - best_text = None if best_macro == float("-inf") else best_macro - print( - f"[epoch {ep}] classify_macro_f1_nonnone={score_now:.6f} " - f"none_fpr={none_fpr_text} best={best_text} saved_best={saved_best}", - flush=True, - ) - elif save_best_unbind and exact > best_exact: - if device.type == "cuda": - torch.cuda.synchronize() - best_exact = exact - best_metrics = dict(metrics) - best_state = _build_weight_state( - encoder, classify, role_head, filler_head, maps, layout - ) - best_residual_records = list(last_residual_records) - best_dir = out_dir / "best" - write_skeleton(best_dir, maps=maps) - save_heads(best_dir, best_state) - (best_dir / "best-checkpoint.json").write_text( - json.dumps( - { - "metric": "unbind_exact", - "epoch": ep, - "unbind_exact": exact, - "unbind_token_f1": best_metrics.get("unbind_token_f1"), - "unbind_slot_f1": best_metrics.get("unbind_slot_f1"), - "unbind_exact_strict": best_metrics.get("unbind_exact_strict"), - "unbind_token_f1_strict": best_metrics.get("unbind_token_f1_strict"), - "unbind_slot_f1_strict": best_metrics.get("unbind_slot_f1_strict"), - "classify_acc": best_metrics.get("classify_acc"), - }, - indent=2, - sort_keys=True, - ) - + "\n", - encoding="utf-8", - ) - saved_best = True - if device.type == "cuda": - torch.cuda.synchronize() - torch.cuda.empty_cache() - # Durable heartbeat — morph68 hung silently after ep4 with frozen docker logs. - progress = { - "epoch": ep, - "unbind_exact": exact, - "best_unbind_exact": None if best_exact == float("-inf") else best_exact, - "saved_best": saved_best, - "n_unbind_phase": phase_meta["n_rows"], - "unbind_phase": phase_meta["phase"], - "classify_acc": metrics.get("classify_acc"), - } - if select_on_classify: - progress["classify_macro_f1_nonnone"] = metrics.get("classify_macro_f1_nonnone") - progress["none_fpr"] = metrics.get("none_fpr") - progress["best_classify_macro_f1_nonnone"] = ( - None if best_macro == float("-inf") else best_macro - ) - if seed_receipt is not None: - progress["seed"] = seed_receipt["seed"] - with progress_path.open("a", encoding="utf-8") as pf: - pf.write(json.dumps(progress, sort_keys=True) + "\n") - pf.flush() - if not select_on_classify: - print( - f"[epoch {ep}] unbind_exact={exact:.6f} best={best_exact if best_exact != float('-inf') else None} saved_best={saved_best}", - flush=True, - ) - - final_state = _build_weight_state( - encoder, classify, role_head, filler_head, maps, layout - ) - # Always write final epoch weights under a distinct name when best-save is on, - # then promote best → primary model.safetensors (fixes morph35 peak-not-saved). - if (save_best_unbind or select_on_classify) and best_state is not None: - final_file = save_heads(out_dir, final_state) - # rename primary final dump aside, then write best as primary - final_path = out_dir / final_file - aside = out_dir / ( - "model.final.safetensors" if final_file == "model.safetensors" else "heads.final.pt" - ) - if final_path.exists(): - final_path.replace(aside) - weight_file = save_heads(out_dir, best_state) - primary_val = best_metrics or {} - gated_from = ( - "best_classify_macro_f1_nonnone" if select_on_classify else "best_unbind_exact" - ) - residual_for_dump = best_residual_records - else: - weight_file = save_heads(out_dir, final_state) - primary_val = epoch_metrics[-1] if epoch_metrics else {} - gated_from = "final_epoch" - residual_for_dump = last_residual_records - aside = None - - last = epoch_metrics[-1] if epoch_metrics else {} - residual_receipt: dict = { - "unbind_residual_dump": "", - "n_unbind_residual": 0, - "unbind_residual_themes": {}, - } - if residual_dump_path: - residual_receipt = write_residual_dump(residual_dump_path, residual_for_dump) - receipt = { - "schema": "hyperlex.hyperlexical.train_receipt.v0.1", - "model_id": MODEL_ID_SEED, - "trunk": TRUNK, - "trunk_dir": str(trunk), - "device": str(device), - "cuda": bool(torch.cuda.is_available()), - "epochs": epochs, - **provenance(root), - "n_train_classify": len(classify_tr), - "task_accounting": task_accounting, - "n_train_unbind": len(unbind_tr), - "n_unfrozen_encoder": n_unfrozen, - "n_encoder_tensors": len(final_state.get("encoder") or {}), - "last_trainable": last_trainable_used, - "unbind_loss_weight": unbind_loss_weight, - "unbind_every_n": unbind_every_n, - "unbind_primary": slot_ce_mode["unbind_primary"], - "unbind_slot_ce_armed": slot_ce_mode["unbind_slot_ce_armed"], - "unbind_slot_ce_aux_lambda": slot_ce_mode["unbind_slot_ce_aux_lambda"], - "unbind_head_slot_weight": head_slot_weight, - "unbind_second_slot_weight": second_slot_weight, - "n_unbind_observed": unbind_recipe["n_unbind_observed"], - "n_unbind_inferred": unbind_recipe["n_unbind_inferred"], - "unbind_observed_upsample": unbind_recipe["unbind_observed_upsample"], - "unbind_inferred_cap": unbind_recipe["unbind_inferred_cap"], - "unbind_inferred_weight": inferred_weight, - "n_unbind_morph_negatives": unbind_recipe["n_unbind_morph_negatives"], - "unbind_morph_margin": morph_margin, - "unbind_curriculum": curriculum_plan["enabled"], - "unbind_curriculum_pos_epochs": curriculum_plan["pos_epochs"], - "unbind_curriculum_type_epochs": curriculum_plan["type_epochs"], - "unbind_curriculum_phases": curriculum_plan["phases"], - "n_unbind_curriculum_positional": curriculum_plan["n_unbind_positional"], - "n_unbind_curriculum_type_slot": curriculum_plan["n_unbind_type_slot"], - "n_unbind_curriculum_joint": curriculum_plan["n_unbind_joint"], - "unbind_filler_denylist_lineages": unbind_recipe.get("unbind_filler_denylist_lineages", 0), - "unbind_hard_atoms_path": unbind_recipe.get("unbind_hard_atoms_path", ""), - "unbind_hard_upsample": unbind_recipe.get("unbind_hard_upsample", 1), - "n_unbind_hard_atoms_matched": unbind_recipe.get("n_unbind_hard_atoms_matched", 0), - "unbind_force_train_path": unbind_recipe.get("unbind_force_train_path", ""), - "release_set": release_stats, - "filler_filter": unbind_recipe.get("filler_filter"), - "n_filler_rows_dropped_train": unbind_recipe.get("n_filler_rows_dropped_train", 0), - "n_filler_rows_dropped_val": unbind_recipe.get("n_filler_rows_dropped_val", 0), - "n_unbind_force_train": unbind_recipe.get("n_unbind_force_train", 0), - "n_unbind_force_train_keys": unbind_recipe.get("n_unbind_force_train_keys", 0), - "n_unbind_val_after_force_train": unbind_recipe.get( - "n_unbind_val_after_force_train", 0 - ), - "n_unbind_hard_extra_copies": unbind_recipe.get("n_unbind_hard_extra_copies", 0), - "unbind_residual_dump": residual_receipt.get("unbind_residual_dump", ""), - "n_unbind_residual": residual_receipt.get("n_unbind_residual", 0), - "unbind_residual_themes": residual_receipt.get("unbind_residual_themes", {}), - "unbind_residual_by_scheme": residual_receipt.get("unbind_residual_by_scheme", {}), - "unbind_residual_by_class": residual_receipt.get("unbind_residual_by_class", {}), - "unbind_residual_summary": residual_receipt.get("unbind_residual_summary", ""), - "last_loss": losses[-1] if losses else None, - "val": primary_val, - "val_final": last, - "val_best": best_metrics, - "save_best_unbind": save_best_unbind, - "primary_weights_from": gated_from, - "best_unbind_exact": None if best_metrics is None else best_metrics.get("unbind_exact"), - "best_epoch": None if best_metrics is None else best_metrics.get("epoch"), - "final_weight_file": None if aside is None else aside.name, - "init_from": init_receipt.get("init_from"), - "warm_start": bool(init_receipt.get("warm_start")), - "init_weight_file": init_receipt.get("weight_file"), - "init_encoder_trainable_loaded": init_receipt.get("encoder_trainable_loaded"), - "init_expand_vocab": bool(init_receipt.get("init_expand_vocab")), - "init_expand_vocab_applied": bool(init_receipt.get("expand_vocab")), - "init_expand_role": init_receipt.get("role"), - "init_expand_filler": init_receipt.get("filler"), - "epoch_metrics": epoch_metrics, - "weight_file": weight_file, - "aligner": "char_span + offset_mapping", - "data_sha256": input_receipt["data_sha256"], - "training_input_mode": input_receipt["training_input_mode"], - "training_export_path": input_receipt["training_export_path"], - "training_export_sha256_expected": input_receipt["training_export_sha256_expected"], - "training_export_sha256_actual": input_receipt["training_export_sha256_actual"], - "training_export_rows": input_receipt["training_export_rows"], - "live_export_generation_enabled": input_receipt["live_export_generation_enabled"], - "holdout": holdout_receipt(holdout_spec, holdout_removed), - "holdout_training_disjoint": disjoint_receipt, - "classify_admission": classify_admission_receipt, - "include_live": include_live, - "live_included": bundle["counts"].get("live_included", 0), - "name_gate": False, - "e2_pass": False, - "brier": None, - "forecast_eligible": False, - "note": "HF-shaped dump. Not Hyperlexical until E2.", - } - if reserve_receipt is not None: - receipt["eval_reserve_disjoint"] = reserve_receipt - if classify_split_receipt is not None: - receipt["classify_split"] = classify_split_receipt - if seed_receipt is not None: - receipt["seed"] = seed_receipt - if select_on_classify: - receipt["select_metric"] = SELECT_METRIC_CLASSIFY - if force_overlap.get("disjoint"): - receipt["force_train_disjoint"] = { - "n_overlap": force_overlap["n_overlap"], - "n_in_val": force_overlap["n_in_val"], - "n_in_selection": force_overlap["n_in_selection"], - "n_force_rows": force_overlap["n_force_rows"], - "n_dropped_classify_val": force_overlap["n_dropped_classify_val"], - "n_dropped_unbind_val": force_overlap["n_dropped_unbind_val"], - "n_classify_val": force_overlap["n_classify_val"], - "n_unbind_val": force_overlap["n_unbind_val"], - } - if resolve_vocab_train_only(): - receipt["vocab_train_only"] = True - (out_dir / "layout.json").write_text(json.dumps(layout, indent=2, sort_keys=True) + "\n", encoding="utf-8") - (out_dir / "train-receipt.json").write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8") - train_config = { - "lr": os.environ.get("HYPERLEX_TRAIN_LR", "2e-5"), - "epochs": epochs, - "batch": batch, - "max_len": MAX_LEN, - "last_trainable": last_trainable_used, - "unbind_loss_weight": unbind_loss_weight, - "unbind_every_n": unbind_every_n, - "unbind_primary": slot_ce_mode["unbind_primary"], - "unbind_slot_ce_armed": slot_ce_mode["unbind_slot_ce_armed"], - "unbind_slot_ce_aux_lambda": slot_ce_mode["unbind_slot_ce_aux_lambda"], - "unbind_head_slot_weight": head_slot_weight, - "unbind_second_slot_weight": second_slot_weight, - "unbind_observed_upsample": unbind_recipe["unbind_observed_upsample"], - "unbind_inferred_cap": unbind_recipe["unbind_inferred_cap"], - "unbind_inferred_weight": inferred_weight, - "n_unbind_morph_negatives": unbind_recipe["n_unbind_morph_negatives"], - "unbind_morph_margin": morph_margin, - "unbind_curriculum": curriculum_plan["enabled"], - "unbind_curriculum_pos_epochs": curriculum_plan["pos_epochs"], - "unbind_curriculum_type_epochs": curriculum_plan["type_epochs"], - "unbind_filler_denylist_lineages": unbind_recipe.get( - "unbind_filler_denylist_lineages", 0 - ), - "unbind_hard_atoms_path": unbind_recipe.get("unbind_hard_atoms_path", ""), - "unbind_hard_upsample": unbind_recipe.get("unbind_hard_upsample", 1), - "n_unbind_hard_atoms_matched": unbind_recipe.get( - "n_unbind_hard_atoms_matched", 0 - ), - "n_unbind_hard_extra_copies": unbind_recipe.get( - "n_unbind_hard_extra_copies", 0 - ), - "unbind_force_train_path": unbind_recipe.get("unbind_force_train_path", ""), - "n_unbind_force_train": unbind_recipe.get("n_unbind_force_train", 0), - "n_unbind_force_train_keys": unbind_recipe.get( - "n_unbind_force_train_keys", 0 - ), - "n_unbind_val_after_force_train": unbind_recipe.get( - "n_unbind_val_after_force_train", 0 - ), - "save_best_unbind": save_best_unbind, - "holdout": holdout_receipt(holdout_spec, holdout_removed), - "init_from": init_receipt.get("init_from"), - "warm_start": bool(init_receipt.get("warm_start")), - } - if seed_receipt is not None: - train_config["seed"] = seed_receipt["seed"] - (out_dir / "config-train.json").write_text( - json.dumps(train_config, indent=2) + "\n", - encoding="utf-8", - ) - return receipt +PLACEHOLDER \ No newline at end of file From 6c5c80ec40ff3345f38b55832e4b1c8b813b1242 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:17:28 -0700 Subject: [PATCH 5/9] Restore loop.py header while the full file upload is prepared. --- scripts/shadow/hyperlexical/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index 311c8dd0..673a3fc5 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -1 +1 @@ -PLACEHOLDER \ No newline at end of file +"""Spark train loop. Gate only.""" From 620ea2024fb9f367bfc90907533394b525109a7a Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:23:41 -0700 Subject: [PATCH 6/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- scripts/shadow/hyperlexical/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index 673a3fc5..311c8dd0 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -1 +1 @@ -"""Spark train loop. Gate only.""" +PLACEHOLDER \ No newline at end of file From eae6acfbb7255e87888a0e29360df5512fc52fc7 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:25:10 -0700 Subject: [PATCH 7/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- scripts/shadow/hyperlexical/loop.py | 714 +++++++++++++++++++++++++++- 1 file changed, 713 insertions(+), 1 deletion(-) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index 311c8dd0..51f88fac 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -1 +1,713 @@ -PLACEHOLDER \ No newline at end of file +"""Spark train loop. Gate only.""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass +from pathlib import Path + +from .align import atom_token_index, offsets_from_tokenizer, pool_indices +from .export import export_dataset, repo_root, write_export +from .admission import AdmissionError, admit_training_run +from .train_input import train_input_receipt +from .classify_metrics import ( + NONE_LABEL, + SELECT_METRIC_CLASSIFY, + SELECT_METRIC_ENV, + macro_f1_nonnone, + none_false_positive_rate, + resolve_select_metric, +) +from .classify_split import apply_classify_split_file +from .seed_control import apply_training_seed +from .force_train_overlap import enforce_force_train_disjoint +from .layout import ( + FAMILIES, + HIDDEN, + MAX_LEN, + MODEL_ID_SEED, + TRUNK, + UNK, + describe, + label_maps_for_splits, + resolve_last_trainable, + resolve_vocab_train_only, +) +from .eval_forward import apply_encoder_trainable +from .filler_filter import assert_publishable_vocab, filter_mode, filter_unbind_rows +from .holdout_guard import ( + assert_no_holdout, + filter_holdout_rows, + holdout_receipt, + load_holdout_spec, + log_holdout, +) +from .provenance import provenance +from .release_set import maybe_release +from .save_pretrained import ( + collect_encoder_trainable, + save_heads, + split_weight_tensors, + write_skeleton, +) +from .training_routing import route_rows +from .unbind_curriculum import ( + plan_unbind_curriculum, + resolve_curriculum_schedule, + select_unbind_for_epoch, +) +from .unbind_metrics import mapped_filler, mapped_pred, summarize_unbind_pairs +from .unbind_recipe import ( + apply_unbind_force_train, + resolve_unbind_inferred_weight, + resolve_unbind_morph_margin, + shape_unbind_train, + unbind_row_sample_weight, +) +from .unbind_head_slot import ( + apply_head_slot_weight, + resolve_unbind_head_slot_weight, +) +from .unbind_second_slot import ( + apply_second_slot_weight, + resolve_unbind_second_slot_weight, +) +from .unbind_residual import ( + residual_row_record, + resolve_unbind_residual_dump_path, + write_residual_dump, +) +from .unbind_slot_ce import ( + UNBIND_SLOT_CE_AUX_LAMBDA, + combine_unbind_train_terms, + resolve_unbind_primary_mode, +) + +UNBIND_LOSS_WEIGHT_ENV = "HYPERLEX_UNBIND_LOSS_WEIGHT" +UNBIND_EVERY_N_ENV = "HYPERLEX_UNBIND_EVERY_N" +SAVE_BEST_UNBIND_ENV = "HYPERLEX_SAVE_BEST_UNBIND" +INIT_FROM_ENV = "HYPERLEX_INIT_FROM" +EARLY_STOP_ENV = "HYPERLEX_EARLY_STOP" +EARLY_STOP_PATIENCE_ENV = "HYPERLEX_EARLY_STOP_PATIENCE" +EARLY_STOP_MIN_EPOCHS_ENV = "HYPERLEX_EARLY_STOP_MIN_EPOCHS" +STOP_REASON_MAX_EPOCHS = "max_epochs" +STOP_REASON_EARLY_STOPPING = "early_stopping" +# Observational seconds on epoch-progress.jsonl and the completion receipt. +# Python round() to 6 decimal places (microseconds). Not a selection input. +WALLCLOCK_SECONDS_DECIMALS = 6 +_EARLY_STOP_OFF = frozenset({"0", "false", "no", "off"}) +_EARLY_STOP_ON = frozenset({"1", "true", "yes", "on"}) +INIT_EXPAND_VOCAB_ENV = "HYPERLEX_INIT_EXPAND_VOCAB" +UNBIND_LOSS_WEIGHT_DEFAULT = 1.0 +UNBIND_EVERY_N_DEFAULT = 1 + + +def resolve_save_best_unbind(raw: str | None = None) -> bool: + """When true, persist best-by-val-unbind_exact weights as primary model.safetensors. + + morph35 peak-not-saved: epoch_metrics recorded ep16 0.5372 but only final + weights were written. Opt-in via HYPERLEX_SAVE_BEST_UNBIND=1. + """ + if raw is None: + raw = os.environ.get(SAVE_BEST_UNBIND_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return False + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +@dataclass(frozen=True) +class EarlyStopConfig: + """Optional classify-metric early stop. Default-off. Does not change the epoch cap.""" + + enabled: bool + max_epochs: int + patience: int | None = None + minimum_epochs: int | None = None + select_metric: str = "" + + +def _parse_early_stop_int(env_name: str, raw: str | None) -> int: + if raw is None or not str(raw).strip(): + raise ValueError(f"{env_name} is required when {EARLY_STOP_ENV} is enabled") + text = str(raw).strip() + sign = "" + digits = text + if text[0] in "+-": + sign, digits = text[0], text[1:] + if sign == "+" and not digits: + raise ValueError(f"{env_name} must be an integer, got {raw!r}") + if not digits.isdigit(): + raise ValueError(f"{env_name} must be an integer, got {raw!r}") + return int(sign + digits) + + +def resolve_early_stop_config(*, max_epochs: int, select_metric: str) -> EarlyStopConfig: + """Read early-stop env before the optimizer is constructed. + + Unset or disabled (0/false/no/off) leaves the configured epoch schedule + unchanged. ``HYPERLEX_TRAIN_EPOCHS`` is only the max-epoch cap and never + turns this on. Enabled runs require ``HLX_SELECT_METRIC=classify_macro_f1_nonnone``, + patience >= 0, and minimum scored epochs in ``1..max_epochs``. + """ + raw = os.environ.get(EARLY_STOP_ENV) + token = "" if raw is None else str(raw).strip().lower() + if raw is None or token == "" or token in _EARLY_STOP_OFF: + return EarlyStopConfig( + enabled=False, + max_epochs=max_epochs, + select_metric=select_metric, + ) + if token not in _EARLY_STOP_ON: + raise ValueError( + f"{EARLY_STOP_ENV} must be unset, off, or on (1/true/yes/on), got {raw!r}" + ) + if select_metric != SELECT_METRIC_CLASSIFY: + raise ValueError( + f"{EARLY_STOP_ENV} requires {SELECT_METRIC_ENV}={SELECT_METRIC_CLASSIFY}; " + f"got {select_metric!r}" + ) + patience = _parse_early_stop_int( + EARLY_STOP_PATIENCE_ENV, os.environ.get(EARLY_STOP_PATIENCE_ENV) + ) + minimum_epochs = _parse_early_stop_int( + EARLY_STOP_MIN_EPOCHS_ENV, os.environ.get(EARLY_STOP_MIN_EPOCHS_ENV) + ) + if patience < 0: + raise ValueError(f"{EARLY_STOP_PATIENCE_ENV} must be >= 0, got {patience}") + if minimum_epochs < 1: + raise ValueError( + f"{EARLY_STOP_MIN_EPOCHS_ENV} must be >= 1 scored epoch, got {minimum_epochs}" + ) + if minimum_epochs > max_epochs: + raise ValueError( + f"{EARLY_STOP_MIN_EPOCHS_ENV}={minimum_epochs} exceeds " + f"HYPERLEX_TRAIN_EPOCHS={max_epochs}" + ) + return EarlyStopConfig( + enabled=True, + max_epochs=max_epochs, + patience=patience, + minimum_epochs=minimum_epochs, + select_metric=select_metric, + ) + + +def note_strict_improvement( + best_value: float, + best_epoch: int | None, + score: float, + epoch_index: int, +) -> tuple[float, int | None, bool]: + """Strict increase replaces the checkpoint. A tie keeps the earlier epoch.""" + if score > best_value: + return score, epoch_index, True + return best_value, best_epoch, False + + +def early_stop_break( + *, + enabled: bool, + epoch_index: int, + best_epoch: int | None, + epochs_scored: int, + minimum_epochs: int | None, + patience: int | None, + max_epochs: int, +) -> bool: + """Return true only when the loop should break before the epoch cap. + + Call this after the epoch is scored and any strict improvement is recorded. + Patience is ``epoch_index - best_epoch`` completed epochs after the best + (0-based indices). The best epoch itself does not consume patience. + Ties do not move ``best_epoch``, so they do consume patience. + Do not stop before ``minimum_epochs`` epochs have been scored. + When the patience condition lands on the final configured epoch, the max + epoch cap wins and this returns false so the loop records ``max_epochs``. + """ + if not enabled: + return False + if epoch_index + 1 >= max_epochs: + return False + if minimum_epochs is None or patience is None or best_epoch is None: + return False + if epochs_scored < minimum_epochs: + return False + return (epoch_index - best_epoch) >= patience + + +def round_seconds(value: float) -> float: + """Seconds rounded to 6 decimal places. Observational; not a selection input.""" + return round(float(value), WALLCLOCK_SECONDS_DECIMALS) + + +def monotonic_seconds() -> float: + """Monotonic clock in seconds. Tests replace this; production uses time.monotonic.""" + return time.monotonic() + + +def epoch_timing_fields( + *, + epoch_started: float, + epoch_ended: float, + training_started: float, +) -> dict[str, float]: + """Wall-clock fields for one epoch-progress.jsonl row. + + Units are seconds. Both values use ``round_seconds`` (6 decimal places). + ``epoch_wallclock_seconds`` is this epoch's scored body + (epoch_ended - epoch_started). ``training_elapsed_seconds`` is the time + from the start of the epoch loop to this epoch's end. Neither value is + read by checkpoint selection or early stopping. + """ + return { + "epoch_wallclock_seconds": round_seconds(epoch_ended - epoch_started), + "training_elapsed_seconds": round_seconds(epoch_ended - training_started), + } + + +def resolve_init_from(raw: str | None = None) -> Path | None: + """Optional warm-start dir with model.safetensors (or heads.pt). + + Opt-in via HYPERLEX_INIT_FROM=/path/to/prior seed (e.g. morph36 best). + """ + if raw is None: + raw = os.environ.get(INIT_FROM_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return None + path = Path(str(raw).strip()).expanduser() + if not path.is_dir(): + raise ValueError(f"{INIT_FROM_ENV} must be an existing directory, got {path}") + return path + + +def resolve_init_expand_vocab(raw: str | None = None) -> bool: + """When true, warm-load remaps shared role/filler rows into expanded vocabs. + + Default fail-closed on vocab mismatch. Opt-in via HYPERLEX_INIT_EXPAND_VOCAB=1 + so a smaller prior seed (e.g. morph65 max pos_5) can warm a harvest that added + pos_6+/new fillers: copy overlapping labels by name, leave new rows at init. + """ + if raw is None: + raw = os.environ.get(INIT_EXPAND_VOCAB_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return False + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +def _init_weight_path(init_dir: Path) -> Path: + for name in ("model.safetensors", "heads.pt"): + candidate = init_dir / name + if candidate.is_file(): + return candidate + raise FileNotFoundError( + f"{INIT_FROM_ENV}={init_dir} missing model.safetensors or heads.pt" + ) + + +def _read_init_vocabs(init_dir: Path) -> tuple[list | None, list | None]: + for name in ("layout.json", "config.json"): + path = init_dir / name + if not path.is_file(): + continue + try: + blob = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + if not isinstance(blob, dict): + continue + role = blob.get("role_vocab") + filler = blob.get("filler_vocab") + if isinstance(role, list) and isinstance(filler, list) and filler: + return role, filler + return None, None + + +def _remap_linear_rows(module, state: dict, init_labels: list, current_labels: list, head: str) -> dict: + """Copy overlapping out-rows from init Linear state into current module by label name.""" + weight = state.get("weight") + if weight is None: + raise ValueError(f"{INIT_FROM_ENV} {head} missing weight") + if int(getattr(weight, "shape", [0])[0]) != len(init_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} weight rows {tuple(weight.shape)} != " + f"init vocab {len(init_labels)}" + ) + if int(module.weight.shape[0]) != len(current_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} module rows {tuple(module.weight.shape)} != " + f"current vocab {len(current_labels)}" + ) + current_of = {lab: i for i, lab in enumerate(current_labels)} + mapped = 0 + skipped = 0 + for ii, lab in enumerate(init_labels): + ci = current_of.get(lab) + if ci is None: + skipped += 1 + continue + module.weight.data[ci].copy_(weight[ii].detach()) + mapped += 1 + bias = state.get("bias") + if bias is not None and module.bias is not None: + if int(getattr(bias, "shape", [0])[0]) != len(init_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} bias rows != init vocab {len(init_labels)}" + ) + for ii, lab in enumerate(init_labels): + ci = current_of.get(lab) + if ci is None: + continue + module.bias.data[ci].copy_(bias[ii].detach()) + if mapped == 0: + raise ValueError( + f"{INIT_FROM_ENV} {head} expand remap matched 0/{len(init_labels)} labels" + ) + return { + "mapped": mapped, + "skipped_init_only": skipped, + "new_current_rows": len(current_labels) - mapped, + "init_n": len(init_labels), + "current_n": len(current_labels), + } + + +def warm_load_checkpoint( + encoder, + classify, + role_head, + filler_head, + maps: dict, + init_dir: Path, + *, + expand_vocab: bool | None = None, +) -> dict: + """Load heads + trainable encoder tensors from a prior seed dump. Fail closed. + + When expand_vocab is true (or HYPERLEX_INIT_EXPAND_VOCAB=1), role/filler heads + remap overlapping labels by name into the current larger vocab; classify still + loads strict. Default remains exact-vocab match. + """ + if expand_vocab is None: + expand_vocab = resolve_init_expand_vocab() + weight_path = _init_weight_path(init_dir) + init_roles, init_fillers = _read_init_vocabs(init_dir) + cur_roles = list(maps.get("role_vocab") or []) + cur_fillers = list(maps.get("filler_vocab") or []) + roles_match = init_roles is None or init_roles == cur_roles + fillers_match = init_fillers is None or init_fillers == cur_fillers + vocab_match = roles_match and fillers_match + if not vocab_match and not expand_vocab: + if init_roles is not None and init_roles != cur_roles: + raise ValueError( + f"{INIT_FROM_ENV} role_vocab mismatch vs current export " + f"(init={len(init_roles)} current={len(cur_roles)})" + ) + raise ValueError( + f"{INIT_FROM_ENV} filler_vocab mismatch vs current export " + f"(init={len(init_fillers or [])} current={len(cur_fillers)})" + ) + if not vocab_match and expand_vocab: + if init_roles is None or init_fillers is None: + raise ValueError( + f"{INIT_EXPAND_VOCAB_ENV}=1 requires init role_vocab+filler_vocab " + f"in {init_dir}/config.json (or layout.json)" + ) + + if weight_path.name == "model.safetensors": + from safetensors.torch import load_file + + split = split_weight_tensors(load_file(str(weight_path), device="cpu")) + heads_blob = None + else: + import torch + + try: + blob = torch.load(str(weight_path), map_location="cpu", weights_only=False) + except TypeError: + blob = torch.load(str(weight_path), map_location="cpu") + if not isinstance(blob, dict): + raise ValueError(f"{weight_path} is not a heads state dict") + split = { + "classify": blob.get("classify") or {}, + "role_head": blob.get("role_head") or {}, + "filler_head": blob.get("filler_head") or {}, + "encoder": { + k: v + for k, v in ( + {} if not isinstance(blob.get("encoder"), dict) else blob["encoder"] + ).items() + }, + } + # normalize encoder keys to encoder.* for apply_encoder_trainable + enc = {} + for k, v in split["encoder"].items(): + key = str(k) + enc[key if key.startswith("encoder.") else f"encoder.{key}"] = v + split["encoder"] = enc + heads_blob = blob + + classify_state = split.get("classify") or {} + if not classify_state: + raise ValueError(f"{weight_path} missing classify tensors") + classify.load_state_dict(classify_state, strict=True) + + expand_receipt: dict = { + "expand_vocab": bool(expand_vocab and not vocab_match), + "vocab_match": vocab_match, + } + if vocab_match or not expand_vocab: + for name, module in (("role_head", role_head), ("filler_head", filler_head)): + state = split.get(name) or {} + if not state: + raise ValueError(f"{weight_path} missing {name} tensors") + module.load_state_dict(state, strict=True) + else: + role_state = split.get("role_head") or {} + filler_state = split.get("filler_head") or {} + if not role_state or not filler_state: + raise ValueError(f"{weight_path} missing role_head/filler_head tensors") + expand_receipt["role"] = _remap_linear_rows( + role_head, role_state, list(init_roles), cur_roles, "role_head" + ) + expand_receipt["filler"] = _remap_linear_rows( + filler_head, filler_state, list(init_fillers), cur_fillers, "filler_head" + ) + + applied = apply_encoder_trainable(encoder, split.get("encoder") or {}) + if applied["present"] and applied["loaded"] == 0: + raise ValueError( + f"{INIT_FROM_ENV} encoder tensors present but none matched trunk keys" + ) + return { + "init_from": str(init_dir), + "weight_file": weight_path.name, + "encoder_trainable_loaded": applied["loaded"], + "encoder_trainable_present": applied["present"], + "heads_blob": bool(heads_blob), + "init_expand_vocab": bool(expand_vocab), + **expand_receipt, + } + + +def _cpu_module_state(module) -> dict: + return {k: v.detach().cpu().contiguous() for k, v in module.state_dict().items()} + + +def _build_weight_state(encoder, classify, role_head, filler_head, maps, layout) -> dict: + return { + "classify": _cpu_module_state(classify), + "role_head": _cpu_module_state(role_head), + "filler_head": _cpu_module_state(filler_head), + "encoder": collect_encoder_trainable(encoder), + "maps": {k: v for k, v in maps.items() if k not in {"family_of", "role_of", "filler_of"}}, + "layout": layout, + } + + +def resolve_unbind_loss_weight(raw: str | float | int | None = None) -> float: + """Scale on unbind loss before backward. Default 1.0. Fail-closed if invalid.""" + if raw is None: + raw = os.environ.get(UNBIND_LOSS_WEIGHT_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return UNBIND_LOSS_WEIGHT_DEFAULT + try: + weight = float(str(raw).strip()) + except (TypeError, ValueError) as exc: + raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") from exc + if weight < 0 or weight != weight or weight == float("inf"): + raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") + return weight + + +def resolve_unbind_every_n(raw: str | int | None = None) -> int: + """Classify-batch stride for an extra unbind step. Default 1 = epoch-end only.""" + if raw is None: + raw = os.environ.get(UNBIND_EVERY_N_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return UNBIND_EVERY_N_DEFAULT + try: + n = int(str(raw).strip(), 10) + except (TypeError, ValueError) as exc: + raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {raw!r}") from exc + if n < 1: + raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {n}") + return n + + +TASK_ROUTING_ENV = "HYPERLEX_TASK_ROUTING" +TASK_ROUTINGS = ("route_rows", "legacy_split") + + +def task_routing(raw: str | None = None) -> str: + """``route_rows`` (default) or ``legacy_split`` (pre-routing loop; reproduces morph75–78). + + ``legacy_split`` selects rows by ``task == "classify"`` / ``task == "unbind"`` only, so + ``classify+unbind`` rows are not trained. Unknown values fail closed. + """ + value = (os.environ.get(TASK_ROUTING_ENV, "") if raw is None else raw).strip() or "route_rows" + if value not in TASK_ROUTINGS: + raise ValueError(f"{TASK_ROUTING_ENV} must be one of {TASK_ROUTINGS}") + return value + + +def should_interleave_unbind(classify_batch_index: int, every_n: int) -> bool: + """True after classify batch `index` (0-based) when every_n > 1.""" + if every_n <= 1: + return False + return (classify_batch_index + 1) % every_n == 0 + + +def prepare_unbind_splits(rows: list) -> tuple[list, list, dict]: + """Train recipe after route_rows. Val frozen unless force-train env is set. + + ``HYPERLEX_UNBIND_FORCE_TRAIN_PATH`` may move authorized OBSERVED exacts + from val→train (accept-style). Empty/unset → val untouched. + ``HLX_HOLDOUT_MANIFESTS`` rows are dropped before that move and before + hard-atom copies, then checked again so neither injection can put them back. + """ + spec = load_holdout_spec() + if task_routing() == "legacy_split": + train = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "train"] + val = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "val"] + else: + routed, _ = route_rows(rows) + train = routed["unbind"]["train"] + val = routed["unbind"]["val"] + train, n_train = filter_holdout_rows(train, spec) + val, n_val = filter_holdout_rows(val, spec) + train, val, force_stats = apply_unbind_force_train(train, val) + train, n_train_injected = filter_holdout_rows(train, spec) + val, n_val_injected = filter_holdout_rows(val, spec) + train, filt_train = filter_unbind_rows(train) + val, filt_val = filter_unbind_rows(val) + shaped, stats = shape_unbind_train(train) + shaped, n_hard = filter_holdout_rows(shaped, spec) + assert_no_holdout(shaped, spec, "unbind train") + assert_no_holdout(val, spec, "unbind val") + stats = { + **stats, + **force_stats, + "filler_filter": filt_train["filler_filter"], + "n_filler_rows_dropped_train": filt_train["n_filler_rows_dropped"], + "n_filler_rows_dropped_val": filt_val["n_filler_rows_dropped"], + "n_holdout_removed_train": n_train + n_train_injected + n_hard, + "n_holdout_removed_val": n_val + n_val_injected, + "holdout_manifests": [dict(item) for item in spec.manifests], + } + return shaped, val, stats + + +def _require_local_model(trunk: Path): + from transformers import AutoModel, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(str(trunk), local_files_only=True) + model = AutoModel.from_pretrained(str(trunk), local_files_only=True) + return tok, model + + +def _layers(encoder): + if hasattr(encoder, "layers"): + return encoder.layers + inner = getattr(encoder, "encoder", None) + if inner is not None and hasattr(inner, "layers"): + return inner.layers + return None + + +def freeze_encoder(encoder, last_trainable: int | None = None) -> tuple[int, int]: + layers = _layers(encoder) + n_layers = len(list(layers)) if layers is not None else None + used = resolve_last_trainable(last_trainable, layer_count=n_layers) + for p in encoder.parameters(): + p.requires_grad = False + n = 0 + if layers is None: + return 0, used + for block in list(layers)[-used:]: + for p in block.parameters(): + p.requires_grad = True + n += p.numel() + return n, used + + +def _offsets(tok, text: str): + try: + return offsets_from_tokenizer(tok, text, max_len=MAX_LEN) + except TypeError: + return None + + +def _enter_training_execution() -> None: + """Reached only after admission. Admission-only mode returns before this.""" + return None + + +def run_loop( + trunk: Path, + out_dir: Path, + *, + include_live: bool = False, + live_store: Path | None = None, +) -> dict: + # Same gates as preflight. Admission-only returns before any optimizer. + admission = admit_training_run( + include_live=include_live, + live_store=live_store, + export_dataset=export_dataset, + trunk=trunk, + out_dir=out_dir, + ) + if os.environ.get("HLX_ADMISSION_ONLY") == "1": + if not admission.ready: + raise AdmissionError(admission.error or "ADMISSION FAIL", admission.receipt) + return admission.receipt + bundle = admission.bundle + holdout_spec = admission.holdout_spec + input_receipt = train_input_receipt(bundle) + disjoint_receipt = admission.disjoint_receipt + reserve_receipt = admission.reserve_receipt + if bundle is None: + raise AdmissionError("ADMISSION FAIL: training bundle was not loaded", admission.receipt) + root = repo_root() + release_rows_, release_stats = maybe_release(bundle["rows"]) + if release_stats["release_set"]: + bundle = {**bundle, "rows": release_rows_} + if any(r.get("role_scheme") == "reviewed_occurrences" for r in bundle["rows"]): + raise ValueError("reviewed occurrences require occurrence-aware loop alignment") + routed, task_accounting = route_rows(bundle["rows"]) + task_accounting = {**task_accounting, "task_routing": task_routing()} + export_dir = Path(os.environ.get("HYPERLEX_EXPORT_DIR") or (root / "specs" / "007-hyperlexical-model" / "exports")) + export_dir.mkdir(parents=True, exist_ok=True) + write_export(export_dir, bundle) + if task_routing() == "legacy_split": + classify_tr = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "train"] + classify_va = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "val"] + else: + classify_tr = routed["classify"]["train"] + classify_va = routed["classify"]["val"] + classify_tr, classify_va, classify_split_receipt = apply_classify_split_file( + classify_tr, classify_va + ) + unbind_tr, unbind_va, unbind_recipe = prepare_unbind_splits(bundle["rows"]) + from .classify_admission import apply_classify_admission + + classify_tr, classify_va, classify_admission_receipt = apply_classify_admission( + bundle["rows"], classify_tr, classify_va + ) + classify_tr, n_classify_train = filter_holdout_rows(classify_tr, holdout_spec) + classify_va, n_classify_val = filter_holdout_rows(classify_va, holdout_spec) + assert_no_holdout(classify_tr, holdout_spec, "classify train") + assert_no_holdout(classify_va, holdout_spec, "classify val") + assert_no_holdout(unbind_tr, holdout_spec, "unbind train") + assert_no_holdout(unbind_va, holdout_spec, "unbind val") + holdout_removed = { + "classify_train": n_classify_train, + "classify_val": n_classify_val, + "unbind_train": unbind_recipe.get("n_holdout_removed_train", 0), + "unbind_val": unbind_recipe.get("n_unbind_val_after_force_train", 0) and 0 or unbind_recipe.get("n_holdout_removed_val", 0), + } + log_holdout(holdout_spec, holdout_removed) + if len(classify_tr) < 8: + raise RuntimeError("not enough classify train rows") From 1076d634a79a8149651c9a709f0d94e12f978768 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:26:14 -0700 Subject: [PATCH 8/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- scripts/shadow/hyperlexical/loop.py | 713 ---------------------------- 1 file changed, 713 deletions(-) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index 51f88fac..e69de29b 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -1,713 +0,0 @@ -"""Spark train loop. Gate only.""" - -from __future__ import annotations - -import json -import os -import time -from dataclasses import dataclass -from pathlib import Path - -from .align import atom_token_index, offsets_from_tokenizer, pool_indices -from .export import export_dataset, repo_root, write_export -from .admission import AdmissionError, admit_training_run -from .train_input import train_input_receipt -from .classify_metrics import ( - NONE_LABEL, - SELECT_METRIC_CLASSIFY, - SELECT_METRIC_ENV, - macro_f1_nonnone, - none_false_positive_rate, - resolve_select_metric, -) -from .classify_split import apply_classify_split_file -from .seed_control import apply_training_seed -from .force_train_overlap import enforce_force_train_disjoint -from .layout import ( - FAMILIES, - HIDDEN, - MAX_LEN, - MODEL_ID_SEED, - TRUNK, - UNK, - describe, - label_maps_for_splits, - resolve_last_trainable, - resolve_vocab_train_only, -) -from .eval_forward import apply_encoder_trainable -from .filler_filter import assert_publishable_vocab, filter_mode, filter_unbind_rows -from .holdout_guard import ( - assert_no_holdout, - filter_holdout_rows, - holdout_receipt, - load_holdout_spec, - log_holdout, -) -from .provenance import provenance -from .release_set import maybe_release -from .save_pretrained import ( - collect_encoder_trainable, - save_heads, - split_weight_tensors, - write_skeleton, -) -from .training_routing import route_rows -from .unbind_curriculum import ( - plan_unbind_curriculum, - resolve_curriculum_schedule, - select_unbind_for_epoch, -) -from .unbind_metrics import mapped_filler, mapped_pred, summarize_unbind_pairs -from .unbind_recipe import ( - apply_unbind_force_train, - resolve_unbind_inferred_weight, - resolve_unbind_morph_margin, - shape_unbind_train, - unbind_row_sample_weight, -) -from .unbind_head_slot import ( - apply_head_slot_weight, - resolve_unbind_head_slot_weight, -) -from .unbind_second_slot import ( - apply_second_slot_weight, - resolve_unbind_second_slot_weight, -) -from .unbind_residual import ( - residual_row_record, - resolve_unbind_residual_dump_path, - write_residual_dump, -) -from .unbind_slot_ce import ( - UNBIND_SLOT_CE_AUX_LAMBDA, - combine_unbind_train_terms, - resolve_unbind_primary_mode, -) - -UNBIND_LOSS_WEIGHT_ENV = "HYPERLEX_UNBIND_LOSS_WEIGHT" -UNBIND_EVERY_N_ENV = "HYPERLEX_UNBIND_EVERY_N" -SAVE_BEST_UNBIND_ENV = "HYPERLEX_SAVE_BEST_UNBIND" -INIT_FROM_ENV = "HYPERLEX_INIT_FROM" -EARLY_STOP_ENV = "HYPERLEX_EARLY_STOP" -EARLY_STOP_PATIENCE_ENV = "HYPERLEX_EARLY_STOP_PATIENCE" -EARLY_STOP_MIN_EPOCHS_ENV = "HYPERLEX_EARLY_STOP_MIN_EPOCHS" -STOP_REASON_MAX_EPOCHS = "max_epochs" -STOP_REASON_EARLY_STOPPING = "early_stopping" -# Observational seconds on epoch-progress.jsonl and the completion receipt. -# Python round() to 6 decimal places (microseconds). Not a selection input. -WALLCLOCK_SECONDS_DECIMALS = 6 -_EARLY_STOP_OFF = frozenset({"0", "false", "no", "off"}) -_EARLY_STOP_ON = frozenset({"1", "true", "yes", "on"}) -INIT_EXPAND_VOCAB_ENV = "HYPERLEX_INIT_EXPAND_VOCAB" -UNBIND_LOSS_WEIGHT_DEFAULT = 1.0 -UNBIND_EVERY_N_DEFAULT = 1 - - -def resolve_save_best_unbind(raw: str | None = None) -> bool: - """When true, persist best-by-val-unbind_exact weights as primary model.safetensors. - - morph35 peak-not-saved: epoch_metrics recorded ep16 0.5372 but only final - weights were written. Opt-in via HYPERLEX_SAVE_BEST_UNBIND=1. - """ - if raw is None: - raw = os.environ.get(SAVE_BEST_UNBIND_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return False - return str(raw).strip().lower() in {"1", "true", "yes", "on"} - - -@dataclass(frozen=True) -class EarlyStopConfig: - """Optional classify-metric early stop. Default-off. Does not change the epoch cap.""" - - enabled: bool - max_epochs: int - patience: int | None = None - minimum_epochs: int | None = None - select_metric: str = "" - - -def _parse_early_stop_int(env_name: str, raw: str | None) -> int: - if raw is None or not str(raw).strip(): - raise ValueError(f"{env_name} is required when {EARLY_STOP_ENV} is enabled") - text = str(raw).strip() - sign = "" - digits = text - if text[0] in "+-": - sign, digits = text[0], text[1:] - if sign == "+" and not digits: - raise ValueError(f"{env_name} must be an integer, got {raw!r}") - if not digits.isdigit(): - raise ValueError(f"{env_name} must be an integer, got {raw!r}") - return int(sign + digits) - - -def resolve_early_stop_config(*, max_epochs: int, select_metric: str) -> EarlyStopConfig: - """Read early-stop env before the optimizer is constructed. - - Unset or disabled (0/false/no/off) leaves the configured epoch schedule - unchanged. ``HYPERLEX_TRAIN_EPOCHS`` is only the max-epoch cap and never - turns this on. Enabled runs require ``HLX_SELECT_METRIC=classify_macro_f1_nonnone``, - patience >= 0, and minimum scored epochs in ``1..max_epochs``. - """ - raw = os.environ.get(EARLY_STOP_ENV) - token = "" if raw is None else str(raw).strip().lower() - if raw is None or token == "" or token in _EARLY_STOP_OFF: - return EarlyStopConfig( - enabled=False, - max_epochs=max_epochs, - select_metric=select_metric, - ) - if token not in _EARLY_STOP_ON: - raise ValueError( - f"{EARLY_STOP_ENV} must be unset, off, or on (1/true/yes/on), got {raw!r}" - ) - if select_metric != SELECT_METRIC_CLASSIFY: - raise ValueError( - f"{EARLY_STOP_ENV} requires {SELECT_METRIC_ENV}={SELECT_METRIC_CLASSIFY}; " - f"got {select_metric!r}" - ) - patience = _parse_early_stop_int( - EARLY_STOP_PATIENCE_ENV, os.environ.get(EARLY_STOP_PATIENCE_ENV) - ) - minimum_epochs = _parse_early_stop_int( - EARLY_STOP_MIN_EPOCHS_ENV, os.environ.get(EARLY_STOP_MIN_EPOCHS_ENV) - ) - if patience < 0: - raise ValueError(f"{EARLY_STOP_PATIENCE_ENV} must be >= 0, got {patience}") - if minimum_epochs < 1: - raise ValueError( - f"{EARLY_STOP_MIN_EPOCHS_ENV} must be >= 1 scored epoch, got {minimum_epochs}" - ) - if minimum_epochs > max_epochs: - raise ValueError( - f"{EARLY_STOP_MIN_EPOCHS_ENV}={minimum_epochs} exceeds " - f"HYPERLEX_TRAIN_EPOCHS={max_epochs}" - ) - return EarlyStopConfig( - enabled=True, - max_epochs=max_epochs, - patience=patience, - minimum_epochs=minimum_epochs, - select_metric=select_metric, - ) - - -def note_strict_improvement( - best_value: float, - best_epoch: int | None, - score: float, - epoch_index: int, -) -> tuple[float, int | None, bool]: - """Strict increase replaces the checkpoint. A tie keeps the earlier epoch.""" - if score > best_value: - return score, epoch_index, True - return best_value, best_epoch, False - - -def early_stop_break( - *, - enabled: bool, - epoch_index: int, - best_epoch: int | None, - epochs_scored: int, - minimum_epochs: int | None, - patience: int | None, - max_epochs: int, -) -> bool: - """Return true only when the loop should break before the epoch cap. - - Call this after the epoch is scored and any strict improvement is recorded. - Patience is ``epoch_index - best_epoch`` completed epochs after the best - (0-based indices). The best epoch itself does not consume patience. - Ties do not move ``best_epoch``, so they do consume patience. - Do not stop before ``minimum_epochs`` epochs have been scored. - When the patience condition lands on the final configured epoch, the max - epoch cap wins and this returns false so the loop records ``max_epochs``. - """ - if not enabled: - return False - if epoch_index + 1 >= max_epochs: - return False - if minimum_epochs is None or patience is None or best_epoch is None: - return False - if epochs_scored < minimum_epochs: - return False - return (epoch_index - best_epoch) >= patience - - -def round_seconds(value: float) -> float: - """Seconds rounded to 6 decimal places. Observational; not a selection input.""" - return round(float(value), WALLCLOCK_SECONDS_DECIMALS) - - -def monotonic_seconds() -> float: - """Monotonic clock in seconds. Tests replace this; production uses time.monotonic.""" - return time.monotonic() - - -def epoch_timing_fields( - *, - epoch_started: float, - epoch_ended: float, - training_started: float, -) -> dict[str, float]: - """Wall-clock fields for one epoch-progress.jsonl row. - - Units are seconds. Both values use ``round_seconds`` (6 decimal places). - ``epoch_wallclock_seconds`` is this epoch's scored body - (epoch_ended - epoch_started). ``training_elapsed_seconds`` is the time - from the start of the epoch loop to this epoch's end. Neither value is - read by checkpoint selection or early stopping. - """ - return { - "epoch_wallclock_seconds": round_seconds(epoch_ended - epoch_started), - "training_elapsed_seconds": round_seconds(epoch_ended - training_started), - } - - -def resolve_init_from(raw: str | None = None) -> Path | None: - """Optional warm-start dir with model.safetensors (or heads.pt). - - Opt-in via HYPERLEX_INIT_FROM=/path/to/prior seed (e.g. morph36 best). - """ - if raw is None: - raw = os.environ.get(INIT_FROM_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return None - path = Path(str(raw).strip()).expanduser() - if not path.is_dir(): - raise ValueError(f"{INIT_FROM_ENV} must be an existing directory, got {path}") - return path - - -def resolve_init_expand_vocab(raw: str | None = None) -> bool: - """When true, warm-load remaps shared role/filler rows into expanded vocabs. - - Default fail-closed on vocab mismatch. Opt-in via HYPERLEX_INIT_EXPAND_VOCAB=1 - so a smaller prior seed (e.g. morph65 max pos_5) can warm a harvest that added - pos_6+/new fillers: copy overlapping labels by name, leave new rows at init. - """ - if raw is None: - raw = os.environ.get(INIT_EXPAND_VOCAB_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return False - return str(raw).strip().lower() in {"1", "true", "yes", "on"} - - -def _init_weight_path(init_dir: Path) -> Path: - for name in ("model.safetensors", "heads.pt"): - candidate = init_dir / name - if candidate.is_file(): - return candidate - raise FileNotFoundError( - f"{INIT_FROM_ENV}={init_dir} missing model.safetensors or heads.pt" - ) - - -def _read_init_vocabs(init_dir: Path) -> tuple[list | None, list | None]: - for name in ("layout.json", "config.json"): - path = init_dir / name - if not path.is_file(): - continue - try: - blob = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - if not isinstance(blob, dict): - continue - role = blob.get("role_vocab") - filler = blob.get("filler_vocab") - if isinstance(role, list) and isinstance(filler, list) and filler: - return role, filler - return None, None - - -def _remap_linear_rows(module, state: dict, init_labels: list, current_labels: list, head: str) -> dict: - """Copy overlapping out-rows from init Linear state into current module by label name.""" - weight = state.get("weight") - if weight is None: - raise ValueError(f"{INIT_FROM_ENV} {head} missing weight") - if int(getattr(weight, "shape", [0])[0]) != len(init_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} weight rows {tuple(weight.shape)} != " - f"init vocab {len(init_labels)}" - ) - if int(module.weight.shape[0]) != len(current_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} module rows {tuple(module.weight.shape)} != " - f"current vocab {len(current_labels)}" - ) - current_of = {lab: i for i, lab in enumerate(current_labels)} - mapped = 0 - skipped = 0 - for ii, lab in enumerate(init_labels): - ci = current_of.get(lab) - if ci is None: - skipped += 1 - continue - module.weight.data[ci].copy_(weight[ii].detach()) - mapped += 1 - bias = state.get("bias") - if bias is not None and module.bias is not None: - if int(getattr(bias, "shape", [0])[0]) != len(init_labels): - raise ValueError( - f"{INIT_FROM_ENV} {head} bias rows != init vocab {len(init_labels)}" - ) - for ii, lab in enumerate(init_labels): - ci = current_of.get(lab) - if ci is None: - continue - module.bias.data[ci].copy_(bias[ii].detach()) - if mapped == 0: - raise ValueError( - f"{INIT_FROM_ENV} {head} expand remap matched 0/{len(init_labels)} labels" - ) - return { - "mapped": mapped, - "skipped_init_only": skipped, - "new_current_rows": len(current_labels) - mapped, - "init_n": len(init_labels), - "current_n": len(current_labels), - } - - -def warm_load_checkpoint( - encoder, - classify, - role_head, - filler_head, - maps: dict, - init_dir: Path, - *, - expand_vocab: bool | None = None, -) -> dict: - """Load heads + trainable encoder tensors from a prior seed dump. Fail closed. - - When expand_vocab is true (or HYPERLEX_INIT_EXPAND_VOCAB=1), role/filler heads - remap overlapping labels by name into the current larger vocab; classify still - loads strict. Default remains exact-vocab match. - """ - if expand_vocab is None: - expand_vocab = resolve_init_expand_vocab() - weight_path = _init_weight_path(init_dir) - init_roles, init_fillers = _read_init_vocabs(init_dir) - cur_roles = list(maps.get("role_vocab") or []) - cur_fillers = list(maps.get("filler_vocab") or []) - roles_match = init_roles is None or init_roles == cur_roles - fillers_match = init_fillers is None or init_fillers == cur_fillers - vocab_match = roles_match and fillers_match - if not vocab_match and not expand_vocab: - if init_roles is not None and init_roles != cur_roles: - raise ValueError( - f"{INIT_FROM_ENV} role_vocab mismatch vs current export " - f"(init={len(init_roles)} current={len(cur_roles)})" - ) - raise ValueError( - f"{INIT_FROM_ENV} filler_vocab mismatch vs current export " - f"(init={len(init_fillers or [])} current={len(cur_fillers)})" - ) - if not vocab_match and expand_vocab: - if init_roles is None or init_fillers is None: - raise ValueError( - f"{INIT_EXPAND_VOCAB_ENV}=1 requires init role_vocab+filler_vocab " - f"in {init_dir}/config.json (or layout.json)" - ) - - if weight_path.name == "model.safetensors": - from safetensors.torch import load_file - - split = split_weight_tensors(load_file(str(weight_path), device="cpu")) - heads_blob = None - else: - import torch - - try: - blob = torch.load(str(weight_path), map_location="cpu", weights_only=False) - except TypeError: - blob = torch.load(str(weight_path), map_location="cpu") - if not isinstance(blob, dict): - raise ValueError(f"{weight_path} is not a heads state dict") - split = { - "classify": blob.get("classify") or {}, - "role_head": blob.get("role_head") or {}, - "filler_head": blob.get("filler_head") or {}, - "encoder": { - k: v - for k, v in ( - {} if not isinstance(blob.get("encoder"), dict) else blob["encoder"] - ).items() - }, - } - # normalize encoder keys to encoder.* for apply_encoder_trainable - enc = {} - for k, v in split["encoder"].items(): - key = str(k) - enc[key if key.startswith("encoder.") else f"encoder.{key}"] = v - split["encoder"] = enc - heads_blob = blob - - classify_state = split.get("classify") or {} - if not classify_state: - raise ValueError(f"{weight_path} missing classify tensors") - classify.load_state_dict(classify_state, strict=True) - - expand_receipt: dict = { - "expand_vocab": bool(expand_vocab and not vocab_match), - "vocab_match": vocab_match, - } - if vocab_match or not expand_vocab: - for name, module in (("role_head", role_head), ("filler_head", filler_head)): - state = split.get(name) or {} - if not state: - raise ValueError(f"{weight_path} missing {name} tensors") - module.load_state_dict(state, strict=True) - else: - role_state = split.get("role_head") or {} - filler_state = split.get("filler_head") or {} - if not role_state or not filler_state: - raise ValueError(f"{weight_path} missing role_head/filler_head tensors") - expand_receipt["role"] = _remap_linear_rows( - role_head, role_state, list(init_roles), cur_roles, "role_head" - ) - expand_receipt["filler"] = _remap_linear_rows( - filler_head, filler_state, list(init_fillers), cur_fillers, "filler_head" - ) - - applied = apply_encoder_trainable(encoder, split.get("encoder") or {}) - if applied["present"] and applied["loaded"] == 0: - raise ValueError( - f"{INIT_FROM_ENV} encoder tensors present but none matched trunk keys" - ) - return { - "init_from": str(init_dir), - "weight_file": weight_path.name, - "encoder_trainable_loaded": applied["loaded"], - "encoder_trainable_present": applied["present"], - "heads_blob": bool(heads_blob), - "init_expand_vocab": bool(expand_vocab), - **expand_receipt, - } - - -def _cpu_module_state(module) -> dict: - return {k: v.detach().cpu().contiguous() for k, v in module.state_dict().items()} - - -def _build_weight_state(encoder, classify, role_head, filler_head, maps, layout) -> dict: - return { - "classify": _cpu_module_state(classify), - "role_head": _cpu_module_state(role_head), - "filler_head": _cpu_module_state(filler_head), - "encoder": collect_encoder_trainable(encoder), - "maps": {k: v for k, v in maps.items() if k not in {"family_of", "role_of", "filler_of"}}, - "layout": layout, - } - - -def resolve_unbind_loss_weight(raw: str | float | int | None = None) -> float: - """Scale on unbind loss before backward. Default 1.0. Fail-closed if invalid.""" - if raw is None: - raw = os.environ.get(UNBIND_LOSS_WEIGHT_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return UNBIND_LOSS_WEIGHT_DEFAULT - try: - weight = float(str(raw).strip()) - except (TypeError, ValueError) as exc: - raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") from exc - if weight < 0 or weight != weight or weight == float("inf"): - raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") - return weight - - -def resolve_unbind_every_n(raw: str | int | None = None) -> int: - """Classify-batch stride for an extra unbind step. Default 1 = epoch-end only.""" - if raw is None: - raw = os.environ.get(UNBIND_EVERY_N_ENV) - if raw is None or (isinstance(raw, str) and not raw.strip()): - return UNBIND_EVERY_N_DEFAULT - try: - n = int(str(raw).strip(), 10) - except (TypeError, ValueError) as exc: - raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {raw!r}") from exc - if n < 1: - raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {n}") - return n - - -TASK_ROUTING_ENV = "HYPERLEX_TASK_ROUTING" -TASK_ROUTINGS = ("route_rows", "legacy_split") - - -def task_routing(raw: str | None = None) -> str: - """``route_rows`` (default) or ``legacy_split`` (pre-routing loop; reproduces morph75–78). - - ``legacy_split`` selects rows by ``task == "classify"`` / ``task == "unbind"`` only, so - ``classify+unbind`` rows are not trained. Unknown values fail closed. - """ - value = (os.environ.get(TASK_ROUTING_ENV, "") if raw is None else raw).strip() or "route_rows" - if value not in TASK_ROUTINGS: - raise ValueError(f"{TASK_ROUTING_ENV} must be one of {TASK_ROUTINGS}") - return value - - -def should_interleave_unbind(classify_batch_index: int, every_n: int) -> bool: - """True after classify batch `index` (0-based) when every_n > 1.""" - if every_n <= 1: - return False - return (classify_batch_index + 1) % every_n == 0 - - -def prepare_unbind_splits(rows: list) -> tuple[list, list, dict]: - """Train recipe after route_rows. Val frozen unless force-train env is set. - - ``HYPERLEX_UNBIND_FORCE_TRAIN_PATH`` may move authorized OBSERVED exacts - from val→train (accept-style). Empty/unset → val untouched. - ``HLX_HOLDOUT_MANIFESTS`` rows are dropped before that move and before - hard-atom copies, then checked again so neither injection can put them back. - """ - spec = load_holdout_spec() - if task_routing() == "legacy_split": - train = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "train"] - val = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "val"] - else: - routed, _ = route_rows(rows) - train = routed["unbind"]["train"] - val = routed["unbind"]["val"] - train, n_train = filter_holdout_rows(train, spec) - val, n_val = filter_holdout_rows(val, spec) - train, val, force_stats = apply_unbind_force_train(train, val) - train, n_train_injected = filter_holdout_rows(train, spec) - val, n_val_injected = filter_holdout_rows(val, spec) - train, filt_train = filter_unbind_rows(train) - val, filt_val = filter_unbind_rows(val) - shaped, stats = shape_unbind_train(train) - shaped, n_hard = filter_holdout_rows(shaped, spec) - assert_no_holdout(shaped, spec, "unbind train") - assert_no_holdout(val, spec, "unbind val") - stats = { - **stats, - **force_stats, - "filler_filter": filt_train["filler_filter"], - "n_filler_rows_dropped_train": filt_train["n_filler_rows_dropped"], - "n_filler_rows_dropped_val": filt_val["n_filler_rows_dropped"], - "n_holdout_removed_train": n_train + n_train_injected + n_hard, - "n_holdout_removed_val": n_val + n_val_injected, - "holdout_manifests": [dict(item) for item in spec.manifests], - } - return shaped, val, stats - - -def _require_local_model(trunk: Path): - from transformers import AutoModel, AutoTokenizer - - tok = AutoTokenizer.from_pretrained(str(trunk), local_files_only=True) - model = AutoModel.from_pretrained(str(trunk), local_files_only=True) - return tok, model - - -def _layers(encoder): - if hasattr(encoder, "layers"): - return encoder.layers - inner = getattr(encoder, "encoder", None) - if inner is not None and hasattr(inner, "layers"): - return inner.layers - return None - - -def freeze_encoder(encoder, last_trainable: int | None = None) -> tuple[int, int]: - layers = _layers(encoder) - n_layers = len(list(layers)) if layers is not None else None - used = resolve_last_trainable(last_trainable, layer_count=n_layers) - for p in encoder.parameters(): - p.requires_grad = False - n = 0 - if layers is None: - return 0, used - for block in list(layers)[-used:]: - for p in block.parameters(): - p.requires_grad = True - n += p.numel() - return n, used - - -def _offsets(tok, text: str): - try: - return offsets_from_tokenizer(tok, text, max_len=MAX_LEN) - except TypeError: - return None - - -def _enter_training_execution() -> None: - """Reached only after admission. Admission-only mode returns before this.""" - return None - - -def run_loop( - trunk: Path, - out_dir: Path, - *, - include_live: bool = False, - live_store: Path | None = None, -) -> dict: - # Same gates as preflight. Admission-only returns before any optimizer. - admission = admit_training_run( - include_live=include_live, - live_store=live_store, - export_dataset=export_dataset, - trunk=trunk, - out_dir=out_dir, - ) - if os.environ.get("HLX_ADMISSION_ONLY") == "1": - if not admission.ready: - raise AdmissionError(admission.error or "ADMISSION FAIL", admission.receipt) - return admission.receipt - bundle = admission.bundle - holdout_spec = admission.holdout_spec - input_receipt = train_input_receipt(bundle) - disjoint_receipt = admission.disjoint_receipt - reserve_receipt = admission.reserve_receipt - if bundle is None: - raise AdmissionError("ADMISSION FAIL: training bundle was not loaded", admission.receipt) - root = repo_root() - release_rows_, release_stats = maybe_release(bundle["rows"]) - if release_stats["release_set"]: - bundle = {**bundle, "rows": release_rows_} - if any(r.get("role_scheme") == "reviewed_occurrences" for r in bundle["rows"]): - raise ValueError("reviewed occurrences require occurrence-aware loop alignment") - routed, task_accounting = route_rows(bundle["rows"]) - task_accounting = {**task_accounting, "task_routing": task_routing()} - export_dir = Path(os.environ.get("HYPERLEX_EXPORT_DIR") or (root / "specs" / "007-hyperlexical-model" / "exports")) - export_dir.mkdir(parents=True, exist_ok=True) - write_export(export_dir, bundle) - if task_routing() == "legacy_split": - classify_tr = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "train"] - classify_va = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "val"] - else: - classify_tr = routed["classify"]["train"] - classify_va = routed["classify"]["val"] - classify_tr, classify_va, classify_split_receipt = apply_classify_split_file( - classify_tr, classify_va - ) - unbind_tr, unbind_va, unbind_recipe = prepare_unbind_splits(bundle["rows"]) - from .classify_admission import apply_classify_admission - - classify_tr, classify_va, classify_admission_receipt = apply_classify_admission( - bundle["rows"], classify_tr, classify_va - ) - classify_tr, n_classify_train = filter_holdout_rows(classify_tr, holdout_spec) - classify_va, n_classify_val = filter_holdout_rows(classify_va, holdout_spec) - assert_no_holdout(classify_tr, holdout_spec, "classify train") - assert_no_holdout(classify_va, holdout_spec, "classify val") - assert_no_holdout(unbind_tr, holdout_spec, "unbind train") - assert_no_holdout(unbind_va, holdout_spec, "unbind val") - holdout_removed = { - "classify_train": n_classify_train, - "classify_val": n_classify_val, - "unbind_train": unbind_recipe.get("n_holdout_removed_train", 0), - "unbind_val": unbind_recipe.get("n_unbind_val_after_force_train", 0) and 0 or unbind_recipe.get("n_holdout_removed_val", 0), - } - log_holdout(holdout_spec, holdout_removed) - if len(classify_tr) < 8: - raise RuntimeError("not enough classify train rows") From d0342aecc9f37bd650a3664c25f8a0a6e3398e13 Mon Sep 17 00:00:00 2001 From: Daniel Meyer Date: Sat, 26 Sep 2026 21:29:05 -0700 Subject: [PATCH 9/9] Add default-off classify early stopping to the training loop. Optional patience stops only after a scored epoch records any strict improvement, and epoch progress plus the completion receipt gain observational wall-clock fields. --- scripts/shadow/hyperlexical/loop.py | 1299 +++++++++++++++++++++++++++ 1 file changed, 1299 insertions(+) diff --git a/scripts/shadow/hyperlexical/loop.py b/scripts/shadow/hyperlexical/loop.py index e69de29b..ba9d6013 100644 --- a/scripts/shadow/hyperlexical/loop.py +++ b/scripts/shadow/hyperlexical/loop.py @@ -0,0 +1,1299 @@ +"""Spark train loop. Gate only.""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass +from pathlib import Path + +from .align import atom_token_index, offsets_from_tokenizer, pool_indices +from .export import export_dataset, repo_root, write_export +from .admission import AdmissionError, admit_training_run +from .train_input import train_input_receipt +from .classify_metrics import ( + NONE_LABEL, + SELECT_METRIC_CLASSIFY, + SELECT_METRIC_ENV, + macro_f1_nonnone, + none_false_positive_rate, + resolve_select_metric, +) +from .classify_split import apply_classify_split_file +from .seed_control import apply_training_seed +from .force_train_overlap import enforce_force_train_disjoint +from .layout import ( + FAMILIES, + HIDDEN, + MAX_LEN, + MODEL_ID_SEED, + TRUNK, + UNK, + describe, + label_maps_for_splits, + resolve_last_trainable, + resolve_vocab_train_only, +) +from .eval_forward import apply_encoder_trainable +from .filler_filter import assert_publishable_vocab, filter_mode, filter_unbind_rows +from .holdout_guard import ( + assert_no_holdout, + filter_holdout_rows, + holdout_receipt, + load_holdout_spec, + log_holdout, +) +from .provenance import provenance +from .release_set import maybe_release +from .save_pretrained import ( + collect_encoder_trainable, + save_heads, + split_weight_tensors, + write_skeleton, +) +from .training_routing import route_rows +from .unbind_curriculum import ( + plan_unbind_curriculum, + resolve_curriculum_schedule, + select_unbind_for_epoch, +) +from .unbind_metrics import mapped_filler, mapped_pred, summarize_unbind_pairs +from .unbind_recipe import ( + apply_unbind_force_train, + resolve_unbind_inferred_weight, + resolve_unbind_morph_margin, + shape_unbind_train, + unbind_row_sample_weight, +) +from .unbind_head_slot import ( + apply_head_slot_weight, + resolve_unbind_head_slot_weight, +) +from .unbind_second_slot import ( + apply_second_slot_weight, + resolve_unbind_second_slot_weight, +) +from .unbind_residual import ( + residual_row_record, + resolve_unbind_residual_dump_path, + write_residual_dump, +) +from .unbind_slot_ce import ( + UNBIND_SLOT_CE_AUX_LAMBDA, + combine_unbind_train_terms, + resolve_unbind_primary_mode, +) + +UNBIND_LOSS_WEIGHT_ENV = "HYPERLEX_UNBIND_LOSS_WEIGHT" +UNBIND_EVERY_N_ENV = "HYPERLEX_UNBIND_EVERY_N" +SAVE_BEST_UNBIND_ENV = "HYPERLEX_SAVE_BEST_UNBIND" +INIT_FROM_ENV = "HYPERLEX_INIT_FROM" +EARLY_STOP_ENV = "HYPERLEX_EARLY_STOP" +EARLY_STOP_PATIENCE_ENV = "HYPERLEX_EARLY_STOP_PATIENCE" +EARLY_STOP_MIN_EPOCHS_ENV = "HYPERLEX_EARLY_STOP_MIN_EPOCHS" +STOP_REASON_MAX_EPOCHS = "max_epochs" +STOP_REASON_EARLY_STOPPING = "early_stopping" +# Observational seconds on epoch-progress.jsonl and the completion receipt. +# Python round() to 6 decimal places (microseconds). Not a selection input. +WALLCLOCK_SECONDS_DECIMALS = 6 +_EARLY_STOP_OFF = frozenset({"0", "false", "no", "off"}) +_EARLY_STOP_ON = frozenset({"1", "true", "yes", "on"}) +INIT_EXPAND_VOCAB_ENV = "HYPERLEX_INIT_EXPAND_VOCAB" +UNBIND_LOSS_WEIGHT_DEFAULT = 1.0 +UNBIND_EVERY_N_DEFAULT = 1 + + +def resolve_save_best_unbind(raw: str | None = None) -> bool: + """When true, persist best-by-val-unbind_exact weights as primary model.safetensors. + + morph35 peak-not-saved: epoch_metrics recorded ep16 0.5372 but only final + weights were written. Opt-in via HYPERLEX_SAVE_BEST_UNBIND=1. + """ + if raw is None: + raw = os.environ.get(SAVE_BEST_UNBIND_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return False + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +@dataclass(frozen=True) +class EarlyStopConfig: + """Optional classify-metric early stop. Default-off. Does not change the epoch cap.""" + + enabled: bool + max_epochs: int + patience: int | None = None + minimum_epochs: int | None = None + select_metric: str = "" + + +def _parse_early_stop_int(env_name: str, raw: str | None) -> int: + if raw is None or not str(raw).strip(): + raise ValueError(f"{env_name} is required when {EARLY_STOP_ENV} is enabled") + text = str(raw).strip() + sign = "" + digits = text + if text[0] in "+-": + sign, digits = text[0], text[1:] + if sign == "+" and not digits: + raise ValueError(f"{env_name} must be an integer, got {raw!r}") + if not digits.isdigit(): + raise ValueError(f"{env_name} must be an integer, got {raw!r}") + return int(sign + digits) + + +def resolve_early_stop_config(*, max_epochs: int, select_metric: str) -> EarlyStopConfig: + """Read early-stop env before the optimizer is constructed. + + Unset or disabled (0/false/no/off) leaves the configured epoch schedule + unchanged. ``HYPERLEX_TRAIN_EPOCHS`` is only the max-epoch cap and never + turns this on. Enabled runs require ``HLX_SELECT_METRIC=classify_macro_f1_nonnone``, + patience >= 0, and minimum scored epochs in ``1..max_epochs``. + """ + raw = os.environ.get(EARLY_STOP_ENV) + token = "" if raw is None else str(raw).strip().lower() + if raw is None or token == "" or token in _EARLY_STOP_OFF: + return EarlyStopConfig( + enabled=False, + max_epochs=max_epochs, + select_metric=select_metric, + ) + if token not in _EARLY_STOP_ON: + raise ValueError( + f"{EARLY_STOP_ENV} must be unset, off, or on (1/true/yes/on), got {raw!r}" + ) + if select_metric != SELECT_METRIC_CLASSIFY: + raise ValueError( + f"{EARLY_STOP_ENV} requires {SELECT_METRIC_ENV}={SELECT_METRIC_CLASSIFY}; " + f"got {select_metric!r}" + ) + patience = _parse_early_stop_int( + EARLY_STOP_PATIENCE_ENV, os.environ.get(EARLY_STOP_PATIENCE_ENV) + ) + minimum_epochs = _parse_early_stop_int( + EARLY_STOP_MIN_EPOCHS_ENV, os.environ.get(EARLY_STOP_MIN_EPOCHS_ENV) + ) + if patience < 0: + raise ValueError(f"{EARLY_STOP_PATIENCE_ENV} must be >= 0, got {patience}") + if minimum_epochs < 1: + raise ValueError( + f"{EARLY_STOP_MIN_EPOCHS_ENV} must be >= 1 scored epoch, got {minimum_epochs}" + ) + if minimum_epochs > max_epochs: + raise ValueError( + f"{EARLY_STOP_MIN_EPOCHS_ENV}={minimum_epochs} exceeds " + f"HYPERLEX_TRAIN_EPOCHS={max_epochs}" + ) + return EarlyStopConfig( + enabled=True, + max_epochs=max_epochs, + patience=patience, + minimum_epochs=minimum_epochs, + select_metric=select_metric, + ) + + +def note_strict_improvement( + best_value: float, + best_epoch: int | None, + score: float, + epoch_index: int, +) -> tuple[float, int | None, bool]: + """Strict increase replaces the checkpoint. A tie keeps the earlier epoch.""" + if score > best_value: + return score, epoch_index, True + return best_value, best_epoch, False + + +def early_stop_break( + *, + enabled: bool, + epoch_index: int, + best_epoch: int | None, + epochs_scored: int, + minimum_epochs: int | None, + patience: int | None, + max_epochs: int, +) -> bool: + """Return true only when the loop should break before the epoch cap. + + Call this after the epoch is scored and any strict improvement is recorded. + Patience is ``epoch_index - best_epoch`` completed epochs after the best + (0-based indices). The best epoch itself does not consume patience. + Ties do not move ``best_epoch``, so they do consume patience. + Do not stop before ``minimum_epochs`` epochs have been scored. + When the patience condition lands on the final configured epoch, the max + epoch cap wins and this returns false so the loop records ``max_epochs``. + """ + if not enabled: + return False + if epoch_index + 1 >= max_epochs: + return False + if minimum_epochs is None or patience is None or best_epoch is None: + return False + if epochs_scored < minimum_epochs: + return False + return (epoch_index - best_epoch) >= patience + + +def round_seconds(value: float) -> float: + """Seconds rounded to 6 decimal places. Observational; not a selection input.""" + return round(float(value), WALLCLOCK_SECONDS_DECIMALS) + + +def monotonic_seconds() -> float: + """Monotonic clock in seconds. Tests replace this; production uses time.monotonic.""" + return time.monotonic() + + +def epoch_timing_fields( + *, + epoch_started: float, + epoch_ended: float, + training_started: float, +) -> dict[str, float]: + """Wall-clock fields for one epoch-progress.jsonl row. + + Units are seconds. Both values use ``round_seconds`` (6 decimal places). + ``epoch_wallclock_seconds`` is this epoch's scored body + (epoch_ended - epoch_started). ``training_elapsed_seconds`` is the time + from the start of the epoch loop to this epoch's end. Neither value is + read by checkpoint selection or early stopping. + """ + return { + "epoch_wallclock_seconds": round_seconds(epoch_ended - epoch_started), + "training_elapsed_seconds": round_seconds(epoch_ended - training_started), + } + + +def resolve_init_from(raw: str | None = None) -> Path | None: + """Optional warm-start dir with model.safetensors (or heads.pt). + + Opt-in via HYPERLEX_INIT_FROM=/path/to/prior seed (e.g. morph36 best). + """ + if raw is None: + raw = os.environ.get(INIT_FROM_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return None + path = Path(str(raw).strip()).expanduser() + if not path.is_dir(): + raise ValueError(f"{INIT_FROM_ENV} must be an existing directory, got {path}") + return path + + +def resolve_init_expand_vocab(raw: str | None = None) -> bool: + """When true, warm-load remaps shared role/filler rows into expanded vocabs. + + Default fail-closed on vocab mismatch. Opt-in via HYPERLEX_INIT_EXPAND_VOCAB=1 + so a smaller prior seed (e.g. morph65 max pos_5) can warm a harvest that added + pos_6+/new fillers: copy overlapping labels by name, leave new rows at init. + """ + if raw is None: + raw = os.environ.get(INIT_EXPAND_VOCAB_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return False + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +def _init_weight_path(init_dir: Path) -> Path: + for name in ("model.safetensors", "heads.pt"): + candidate = init_dir / name + if candidate.is_file(): + return candidate + raise FileNotFoundError( + f"{INIT_FROM_ENV}={init_dir} missing model.safetensors or heads.pt" + ) + + +def _read_init_vocabs(init_dir: Path) -> tuple[list | None, list | None]: + for name in ("layout.json", "config.json"): + path = init_dir / name + if not path.is_file(): + continue + try: + blob = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + if not isinstance(blob, dict): + continue + role = blob.get("role_vocab") + filler = blob.get("filler_vocab") + if isinstance(role, list) and isinstance(filler, list) and filler: + return role, filler + return None, None + + +def _remap_linear_rows(module, state: dict, init_labels: list, current_labels: list, head: str) -> dict: + """Copy overlapping out-rows from init Linear state into current module by label name.""" + weight = state.get("weight") + if weight is None: + raise ValueError(f"{INIT_FROM_ENV} {head} missing weight") + if int(getattr(weight, "shape", [0])[0]) != len(init_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} weight rows {tuple(weight.shape)} != " + f"init vocab {len(init_labels)}" + ) + if int(module.weight.shape[0]) != len(current_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} module rows {tuple(module.weight.shape)} != " + f"current vocab {len(current_labels)}" + ) + current_of = {lab: i for i, lab in enumerate(current_labels)} + mapped = 0 + skipped = 0 + for ii, lab in enumerate(init_labels): + ci = current_of.get(lab) + if ci is None: + skipped += 1 + continue + module.weight.data[ci].copy_(weight[ii].detach()) + mapped += 1 + bias = state.get("bias") + if bias is not None and module.bias is not None: + if int(getattr(bias, "shape", [0])[0]) != len(init_labels): + raise ValueError( + f"{INIT_FROM_ENV} {head} bias rows != init vocab {len(init_labels)}" + ) + for ii, lab in enumerate(init_labels): + ci = current_of.get(lab) + if ci is None: + continue + module.bias.data[ci].copy_(bias[ii].detach()) + if mapped == 0: + raise ValueError( + f"{INIT_FROM_ENV} {head} expand remap matched 0/{len(init_labels)} labels" + ) + return { + "mapped": mapped, + "skipped_init_only": skipped, + "new_current_rows": len(current_labels) - mapped, + "init_n": len(init_labels), + "current_n": len(current_labels), + } + + +def warm_load_checkpoint( + encoder, + classify, + role_head, + filler_head, + maps: dict, + init_dir: Path, + *, + expand_vocab: bool | None = None, +) -> dict: + """Load heads + trainable encoder tensors from a prior seed dump. Fail closed. + + When expand_vocab is true (or HYPERLEX_INIT_EXPAND_VOCAB=1), role/filler heads + remap overlapping labels by name into the current larger vocab; classify still + loads strict. Default remains exact-vocab match. + """ + if expand_vocab is None: + expand_vocab = resolve_init_expand_vocab() + weight_path = _init_weight_path(init_dir) + init_roles, init_fillers = _read_init_vocabs(init_dir) + cur_roles = list(maps.get("role_vocab") or []) + cur_fillers = list(maps.get("filler_vocab") or []) + roles_match = init_roles is None or init_roles == cur_roles + fillers_match = init_fillers is None or init_fillers == cur_fillers + vocab_match = roles_match and fillers_match + if not vocab_match and not expand_vocab: + if init_roles is not None and init_roles != cur_roles: + raise ValueError( + f"{INIT_FROM_ENV} role_vocab mismatch vs current export " + f"(init={len(init_roles)} current={len(cur_roles)})" + ) + raise ValueError( + f"{INIT_FROM_ENV} filler_vocab mismatch vs current export " + f"(init={len(init_fillers or [])} current={len(cur_fillers)})" + ) + if not vocab_match and expand_vocab: + if init_roles is None or init_fillers is None: + raise ValueError( + f"{INIT_EXPAND_VOCAB_ENV}=1 requires init role_vocab+filler_vocab " + f"in {init_dir}/config.json (or layout.json)" + ) + + if weight_path.name == "model.safetensors": + from safetensors.torch import load_file + + split = split_weight_tensors(load_file(str(weight_path), device="cpu")) + heads_blob = None + else: + import torch + + try: + blob = torch.load(str(weight_path), map_location="cpu", weights_only=False) + except TypeError: + blob = torch.load(str(weight_path), map_location="cpu") + if not isinstance(blob, dict): + raise ValueError(f"{weight_path} is not a heads state dict") + split = { + "classify": blob.get("classify") or {}, + "role_head": blob.get("role_head") or {}, + "filler_head": blob.get("filler_head") or {}, + "encoder": { + k: v + for k, v in ( + {} if not isinstance(blob.get("encoder"), dict) else blob["encoder"] + ).items() + }, + } + # normalize encoder keys to encoder.* for apply_encoder_trainable + enc = {} + for k, v in split["encoder"].items(): + key = str(k) + enc[key if key.startswith("encoder.") else f"encoder.{key}"] = v + split["encoder"] = enc + heads_blob = blob + + classify_state = split.get("classify") or {} + if not classify_state: + raise ValueError(f"{weight_path} missing classify tensors") + classify.load_state_dict(classify_state, strict=True) + + expand_receipt: dict = { + "expand_vocab": bool(expand_vocab and not vocab_match), + "vocab_match": vocab_match, + } + if vocab_match or not expand_vocab: + for name, module in (("role_head", role_head), ("filler_head", filler_head)): + state = split.get(name) or {} + if not state: + raise ValueError(f"{weight_path} missing {name} tensors") + module.load_state_dict(state, strict=True) + else: + role_state = split.get("role_head") or {} + filler_state = split.get("filler_head") or {} + if not role_state or not filler_state: + raise ValueError(f"{weight_path} missing role_head/filler_head tensors") + expand_receipt["role"] = _remap_linear_rows( + role_head, role_state, list(init_roles), cur_roles, "role_head" + ) + expand_receipt["filler"] = _remap_linear_rows( + filler_head, filler_state, list(init_fillers), cur_fillers, "filler_head" + ) + + applied = apply_encoder_trainable(encoder, split.get("encoder") or {}) + if applied["present"] and applied["loaded"] == 0: + raise ValueError( + f"{INIT_FROM_ENV} encoder tensors present but none matched trunk keys" + ) + return { + "init_from": str(init_dir), + "weight_file": weight_path.name, + "encoder_trainable_loaded": applied["loaded"], + "encoder_trainable_present": applied["present"], + "heads_blob": bool(heads_blob), + "init_expand_vocab": bool(expand_vocab), + **expand_receipt, + } + + +def _cpu_module_state(module) -> dict: + return {k: v.detach().cpu().contiguous() for k, v in module.state_dict().items()} + + +def _build_weight_state(encoder, classify, role_head, filler_head, maps, layout) -> dict: + return { + "classify": _cpu_module_state(classify), + "role_head": _cpu_module_state(role_head), + "filler_head": _cpu_module_state(filler_head), + "encoder": collect_encoder_trainable(encoder), + "maps": {k: v for k, v in maps.items() if k not in {"family_of", "role_of", "filler_of"}}, + "layout": layout, + } + + +def resolve_unbind_loss_weight(raw: str | float | int | None = None) -> float: + """Scale on unbind loss before backward. Default 1.0. Fail-closed if invalid.""" + if raw is None: + raw = os.environ.get(UNBIND_LOSS_WEIGHT_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return UNBIND_LOSS_WEIGHT_DEFAULT + try: + weight = float(str(raw).strip()) + except (TypeError, ValueError) as exc: + raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") from exc + if weight < 0 or weight != weight or weight == float("inf"): + raise ValueError(f"{UNBIND_LOSS_WEIGHT_ENV} must be a finite number >= 0, got {raw!r}") + return weight + + +def resolve_unbind_every_n(raw: str | int | None = None) -> int: + """Classify-batch stride for an extra unbind step. Default 1 = epoch-end only.""" + if raw is None: + raw = os.environ.get(UNBIND_EVERY_N_ENV) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return UNBIND_EVERY_N_DEFAULT + try: + n = int(str(raw).strip(), 10) + except (TypeError, ValueError) as exc: + raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {raw!r}") from exc + if n < 1: + raise ValueError(f"{UNBIND_EVERY_N_ENV} must be a positive int, got {n}") + return n + + +TASK_ROUTING_ENV = "HYPERLEX_TASK_ROUTING" +TASK_ROUTINGS = ("route_rows", "legacy_split") + + +def task_routing(raw: str | None = None) -> str: + """``route_rows`` (default) or ``legacy_split`` (pre-routing loop; reproduces morph75–78). + + ``legacy_split`` selects rows by ``task == "classify"`` / ``task == "unbind"`` only, so + ``classify+unbind`` rows are not trained. Unknown values fail closed. + """ + value = (os.environ.get(TASK_ROUTING_ENV, "") if raw is None else raw).strip() or "route_rows" + if value not in TASK_ROUTINGS: + raise ValueError(f"{TASK_ROUTING_ENV} must be one of {TASK_ROUTINGS}") + return value + + +def should_interleave_unbind(classify_batch_index: int, every_n: int) -> bool: + """True after classify batch `index` (0-based) when every_n > 1.""" + if every_n <= 1: + return False + return (classify_batch_index + 1) % every_n == 0 + + +def prepare_unbind_splits(rows: list) -> tuple[list, list, dict]: + """Train recipe after route_rows. Val frozen unless force-train env is set. + + ``HYPERLEX_UNBIND_FORCE_TRAIN_PATH`` may move authorized OBSERVED exacts + from val→train (accept-style). Empty/unset → val untouched. + ``HLX_HOLDOUT_MANIFESTS`` rows are dropped before that move and before + hard-atom copies, then checked again so neither injection can put them back. + """ + spec = load_holdout_spec() + if task_routing() == "legacy_split": + train = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "train"] + val = [r for r in rows if r.get("task") == "unbind" and r.get("split") == "val"] + else: + routed, _ = route_rows(rows) + train = routed["unbind"]["train"] + val = routed["unbind"]["val"] + train, n_train = filter_holdout_rows(train, spec) + val, n_val = filter_holdout_rows(val, spec) + train, val, force_stats = apply_unbind_force_train(train, val) + train, n_train_injected = filter_holdout_rows(train, spec) + val, n_val_injected = filter_holdout_rows(val, spec) + train, filt_train = filter_unbind_rows(train) + val, filt_val = filter_unbind_rows(val) + shaped, stats = shape_unbind_train(train) + shaped, n_hard = filter_holdout_rows(shaped, spec) + assert_no_holdout(shaped, spec, "unbind train") + assert_no_holdout(val, spec, "unbind val") + stats = { + **stats, + **force_stats, + "filler_filter": filt_train["filler_filter"], + "n_filler_rows_dropped_train": filt_train["n_filler_rows_dropped"], + "n_filler_rows_dropped_val": filt_val["n_filler_rows_dropped"], + "n_holdout_removed_train": n_train + n_train_injected + n_hard, + "n_holdout_removed_val": n_val + n_val_injected, + "holdout_manifests": [dict(item) for item in spec.manifests], + } + return shaped, val, stats + + +def _require_local_model(trunk: Path): + from transformers import AutoModel, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(str(trunk), local_files_only=True) + model = AutoModel.from_pretrained(str(trunk), local_files_only=True) + return tok, model + + +def _layers(encoder): + if hasattr(encoder, "layers"): + return encoder.layers + inner = getattr(encoder, "encoder", None) + if inner is not None and hasattr(inner, "layers"): + return inner.layers + return None + + +def freeze_encoder(encoder, last_trainable: int | None = None) -> tuple[int, int]: + layers = _layers(encoder) + n_layers = len(list(layers)) if layers is not None else None + used = resolve_last_trainable(last_trainable, layer_count=n_layers) + for p in encoder.parameters(): + p.requires_grad = False + n = 0 + if layers is None: + return 0, used + for block in list(layers)[-used:]: + for p in block.parameters(): + p.requires_grad = True + n += p.numel() + return n, used + + +def _offsets(tok, text: str): + try: + return offsets_from_tokenizer(tok, text, max_len=MAX_LEN) + except TypeError: + return None + + +def _enter_training_execution() -> None: + """Reached only after admission. Admission-only mode returns before this.""" + return None + + +def run_loop( + trunk: Path, + out_dir: Path, + *, + include_live: bool = False, + live_store: Path | None = None, +) -> dict: + # Same gates as preflight. Admission-only returns before any optimizer. + admission = admit_training_run( + include_live=include_live, + live_store=live_store, + export_dataset=export_dataset, + trunk=trunk, + out_dir=out_dir, + ) + if os.environ.get("HLX_ADMISSION_ONLY") == "1": + if not admission.ready: + raise AdmissionError(admission.error or "ADMISSION FAIL", admission.receipt) + return admission.receipt + bundle = admission.bundle + holdout_spec = admission.holdout_spec + input_receipt = train_input_receipt(bundle) + disjoint_receipt = admission.disjoint_receipt + reserve_receipt = admission.reserve_receipt + if bundle is None: + raise AdmissionError("ADMISSION FAIL: training bundle was not loaded", admission.receipt) + root = repo_root() + release_rows_, release_stats = maybe_release(bundle["rows"]) + if release_stats["release_set"]: + bundle = {**bundle, "rows": release_rows_} + if any(r.get("role_scheme") == "reviewed_occurrences" for r in bundle["rows"]): + raise ValueError("reviewed occurrences require occurrence-aware loop alignment") + routed, task_accounting = route_rows(bundle["rows"]) + task_accounting = {**task_accounting, "task_routing": task_routing()} + export_dir = Path(os.environ.get("HYPERLEX_EXPORT_DIR") or (root / "specs" / "007-hyperlexical-model" / "exports")) + export_dir.mkdir(parents=True, exist_ok=True) + write_export(export_dir, bundle) + if task_routing() == "legacy_split": + classify_tr = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "train"] + classify_va = [r for r in bundle["rows"] if r["task"] == "classify" and r["split"] == "val"] + else: + classify_tr = routed["classify"]["train"] + classify_va = routed["classify"]["val"] + classify_tr, classify_va, classify_split_receipt = apply_classify_split_file( + classify_tr, classify_va + ) + unbind_tr, unbind_va, unbind_recipe = prepare_unbind_splits(bundle["rows"]) + from .classify_admission import apply_classify_admission + + classify_tr, classify_va, classify_admission_receipt = apply_classify_admission( + bundle["rows"], classify_tr, classify_va + ) + classify_tr, n_classify_train = filter_holdout_rows(classify_tr, holdout_spec) + classify_va, n_classify_val = filter_holdout_rows(classify_va, holdout_spec) + assert_no_holdout(classify_tr, holdout_spec, "classify train") + assert_no_holdout(classify_va, holdout_spec, "classify val") + assert_no_holdout(unbind_tr, holdout_spec, "unbind train") + assert_no_holdout(unbind_va, holdout_spec, "unbind val") + holdout_removed = { + "classify_train": n_classify_train, + "classify_val": n_classify_val, + "unbind_train": unbind_recipe.get("n_holdout_removed_train", 0), + "unbind_val": unbind_recipe.get("n_holdout_removed_val", 0), + } + log_holdout(holdout_spec, holdout_removed) + if len(classify_tr) < 8: + raise RuntimeError("not enough classify train rows") + + select_metric = resolve_select_metric() + select_on_classify = select_metric == SELECT_METRIC_CLASSIFY + if select_on_classify: + selection_rows = list(classify_va or classify_tr[:8]) + else: + selection_rows = list(unbind_va or unbind_tr[:8]) + classify_va, unbind_va, force_overlap = enforce_force_train_disjoint( + classify_va, + unbind_va, + selection_rows, + ) + + _enter_training_execution() + import torch + from torch import nn + from torch.optim import AdamW + + seed_receipt = apply_training_seed(torch) + maps = label_maps_for_splits(unbind_tr, unbind_va) + if filter_mode() == "strict": + assert_publishable_vocab(maps["filler_vocab"]) + tok, encoder = _require_local_model(trunk) + hidden = int(getattr(encoder.config, "hidden_size", HIDDEN)) + if hidden != HIDDEN: + raise RuntimeError(f"hidden {hidden} != {HIDDEN}") + n_unfrozen, last_trainable_used = freeze_encoder(encoder) + classify = nn.Linear(hidden, len(FAMILIES)) + role_head = nn.Linear(hidden, len(maps["role_vocab"])) + filler_head = nn.Linear(hidden, len(maps["filler_vocab"])) + init_from = resolve_init_from() + init_receipt: dict = {"init_from": None, "warm_start": False} + if init_from is not None: + init_receipt = { + "warm_start": True, + **warm_load_checkpoint( + encoder, classify, role_head, filler_head, maps, init_from + ), + } + trainable = [p for p in encoder.parameters() if p.requires_grad] + list(classify.parameters()) + list(role_head.parameters()) + list(filler_head.parameters()) + epochs = int(os.environ.get("HYPERLEX_TRAIN_EPOCHS", "2")) + early_stop = resolve_early_stop_config(max_epochs=epochs, select_metric=select_metric) + opt = AdamW(trainable, lr=float(os.environ.get("HYPERLEX_TRAIN_LR", "2e-5"))) + batch = int(os.environ.get("HYPERLEX_TRAIN_BATCH", "8")) + unbind_loss_weight = resolve_unbind_loss_weight() + unbind_every_n = resolve_unbind_every_n() + morph_margin = resolve_unbind_morph_margin() + inferred_weight = resolve_unbind_inferred_weight() + slot_ce_mode = resolve_unbind_primary_mode() + unbind_primary = slot_ce_mode["unbind_primary"] + head_slot_weight = resolve_unbind_head_slot_weight() + second_slot_weight = resolve_unbind_second_slot_weight() + curriculum = resolve_curriculum_schedule() + curriculum_plan = plan_unbind_curriculum(unbind_tr, epochs, curriculum) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + for mod in (encoder, classify, role_head, filler_head): + mod.to(device) + encoder.train() + losses = [] + epoch_metrics = [] + + def encode_texts(texts): + enc = tok(texts, padding=True, truncation=True, max_length=MAX_LEN, return_tensors="pt") + return {k: v.to(device) for k, v in enc.items()} + + def unbind_loss(row): + fillers = list(row.get("fillers") or []) + roles = list(row.get("roles") or []) + if not fillers: + return None + out = encoder(**encode_texts([row["text"]])) + states = out.last_hidden_state[0] + offs = _offsets(tok, row["text"]) + slot_ces = [] + aux_terms = [] + for k, fill in enumerate(fillers): + idxs = pool_indices(states.size(0), atom_token_index(row["text"], fill, offs)) + h = states[idxs].mean(0) + gold_f = maps["filler_of"].get(fill, maps["filler_of"][UNK]) + logits = filler_head(h) + slot_ces.append( + nn.functional.cross_entropy( + logits.unsqueeze(0), torch.tensor([gold_f], device=device) + ) + ) + hard = [nf for nf in (row.get("hard_neg_fillers") or []) if nf] + if hard: + gold_logit = logits[gold_f] + neg_vals = [] + for nf in hard: + ni = maps["filler_of"].get(nf) + if ni is None or ni == gold_f: + continue + neg_vals.append(logits[ni]) + if neg_vals: + stacked = torch.stack(neg_vals) + aux_terms.append(torch.relu(stacked + morph_margin - gold_logit).sum()) + if k < len(roles): + gold_r = maps["role_of"].get(roles[k], maps["role_of"][UNK]) + aux_terms.append( + nn.functional.cross_entropy( + role_head(h).unsqueeze(0), torch.tensor([gold_r], device=device) + ) + ) + weighted_slots = apply_head_slot_weight(slot_ces, head_slot_weight) + weighted_slots = apply_second_slot_weight(weighted_slots, second_slot_weight) + return combine_unbind_train_terms( + weighted_slots, + aux_terms, + primary=unbind_primary, + aux_lambda=UNBIND_SLOT_CE_AUX_LAMBDA, + ) + + # Keep last train loss on-device; avoid per-step .cpu() sync (morph68 hang: + # post-SAVE_BEST host spin at ~98% CPU / GPU util 0 with mem held). + last_train_loss = None + + def step_unbind(row) -> None: + nonlocal last_train_loss + if unbind_loss_weight == 0: + return + uloss = unbind_loss(row) + if uloss is None: + return + row_w = unbind_row_sample_weight(row, inferred_weight) + scaled = uloss * unbind_loss_weight * row_w + opt.zero_grad() + scaled.backward() + opt.step() + last_train_loss = scaled.detach() + + residual_dump_path = resolve_unbind_residual_dump_path() + last_residual_records: list[dict] = [] + save_best_unbind = resolve_save_best_unbind() + best_exact = float("-inf") + best_macro = float("-inf") + best_epoch_index: int | None = None + best_metrics: dict | None = None + best_state: dict | None = None + best_residual_records: list[dict] = [] + progress_path = out_dir / "epoch-progress.jsonl" + + @torch.no_grad() + def score(): + nonlocal last_residual_records + encoder.eval() + classify.eval() + filler_head.eval() + hit = tot = 0 + classify_golds: list[str] = [] + classify_preds: list[str] = [] + for row in classify_va or classify_tr[:8]: + out = encoder(**encode_texts([row["text"]])) + pred = int(classify(out.last_hidden_state[:, 0]).argmax(-1)[0]) + gold = maps["family_of"].get(row["lineage"], maps["family_of"]["none"]) + hit += int(pred == gold) + tot += 1 + if select_on_classify: + classify_golds.append(FAMILIES[gold] if 0 <= gold < len(FAMILIES) else NONE_LABEL) + classify_preds.append(FAMILIES[pred] if 0 <= pred < len(FAMILIES) else NONE_LABEL) + pairs: list[tuple[list[str], list[str]]] = [] + strict_pairs: list[tuple[list[str], list[str]]] = [] + residual_records: list[dict] = [] + for row in unbind_va or unbind_tr[:8]: + fillers = list(row.get("fillers") or []) + if not fillers: + continue + out = encoder(**encode_texts([row["text"]])) + states = out.last_hidden_state[0] + offs = _offsets(tok, row["text"]) + gold_strs: list[str] = [] + raw_strs: list[str] = [] + pred_strs: list[str] = [] + for fill in fillers: + idxs = pool_indices(states.size(0), atom_token_index(row["text"], fill, offs)) + pred = int(filler_head(states[idxs].mean(0)).argmax()) + gold_strs.append(mapped_filler(maps, fill)) + raw_strs.append(str(fill).lower()) + pred_strs.append(mapped_pred(maps, pred)) + pairs.append((gold_strs, pred_strs)) + strict_pairs.append((raw_strs, pred_strs)) + if residual_dump_path: + rec = residual_row_record( + text=str(row.get("text") or ""), + gold=gold_strs, + pred=pred_strs, + row=row, + ) + if rec is not None: + residual_records.append(rec) + last_residual_records = residual_records + encoder.train() + classify.train() + filler_head.train() + metrics = summarize_unbind_pairs(pairs, strict_pairs=strict_pairs) + metrics["classify_acc"] = hit / max(1, tot) + metrics["n_classify_eval"] = tot + if select_on_classify: + metrics["classify_macro_f1_nonnone"] = macro_f1_nonnone(classify_golds, classify_preds) + metrics["none_fpr"] = none_false_positive_rate(classify_golds, classify_preds) + return metrics + + out_dir.mkdir(parents=True, exist_ok=True) + layout = describe(maps) + layout["last_trainable"] = last_trainable_used + layout["aligner"] = "char_span + offset_mapping" + write_skeleton(out_dir, maps=maps) + + training_started = monotonic_seconds() + stop_reason = STOP_REASON_MAX_EPOCHS + for ep in range(epochs): + epoch_started = monotonic_seconds() + phase_rows, phase_meta = select_unbind_for_epoch(unbind_tr, ep, curriculum) + unbind_cycle = 0 + classify_batch_i = 0 + for i in range(0, len(classify_tr), batch): + chunk = classify_tr[i : i + batch] + y = torch.tensor([maps["family_of"].get(c["lineage"], maps["family_of"]["none"]) for c in chunk], device=device) + out = encoder(**encode_texts([c["text"] for c in chunk])) + loss = nn.functional.cross_entropy(classify(out.last_hidden_state[:, 0]), y) + opt.zero_grad() + loss.backward() + opt.step() + last_train_loss = loss.detach() + if should_interleave_unbind(classify_batch_i, unbind_every_n) and phase_rows: + step_unbind(phase_rows[unbind_cycle % len(phase_rows)]) + unbind_cycle += 1 + classify_batch_i += 1 + for row in phase_rows: + step_unbind(row) + metrics = score() + metrics["epoch"] = ep + metrics["unbind_phase"] = phase_meta["phase"] + metrics["n_unbind_phase"] = phase_meta["n_rows"] + metrics["unbind_phase_fallback_full_mix"] = phase_meta["fallback_full_mix"] + epoch_metrics.append(metrics) + exact = float(metrics.get("unbind_exact") or 0.0) + if last_train_loss is not None: + # One host sync per epoch (not per step). + losses.append(float(last_train_loss.item())) + saved_best = False + if select_on_classify: + if metrics.get("classify_macro_f1_nonnone") is None: + raise RuntimeError( + "HLX_SELECT_METRIC=classify_macro_f1_nonnone but val has no " + "non-none gold; macro-F1 is NOT_COMPUTABLE" + ) + score_now = float(metrics["classify_macro_f1_nonnone"]) + best_macro, best_epoch_index, improved = note_strict_improvement( + best_macro, + best_epoch_index, + score_now, + ep, + ) + if improved: + if device.type == "cuda": + torch.cuda.synchronize() + best_metrics = dict(metrics) + best_state = _build_weight_state( + encoder, classify, role_head, filler_head, maps, layout + ) + best_residual_records = list(last_residual_records) + best_dir = out_dir / "best" + write_skeleton(best_dir, maps=maps) + save_heads(best_dir, best_state) + (best_dir / "best-checkpoint.json").write_text( + json.dumps( + { + "metric": SELECT_METRIC_CLASSIFY, + "epoch": ep, + "classify_macro_f1_nonnone": score_now, + "none_fpr": metrics.get("none_fpr"), + "unbind_exact": exact, + "unbind_token_f1": best_metrics.get("unbind_token_f1"), + "unbind_slot_f1": best_metrics.get("unbind_slot_f1"), + "unbind_exact_strict": best_metrics.get("unbind_exact_strict"), + "unbind_token_f1_strict": best_metrics.get("unbind_token_f1_strict"), + "unbind_slot_f1_strict": best_metrics.get("unbind_slot_f1_strict"), + "classify_acc": best_metrics.get("classify_acc"), + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + saved_best = True + if device.type == "cuda": + torch.cuda.synchronize() + torch.cuda.empty_cache() + none_fpr = metrics.get("none_fpr") + none_fpr_text = "null" if none_fpr is None else f"{float(none_fpr):.6f}" + best_text = None if best_macro == float("-inf") else best_macro + print( + f"[epoch {ep}] classify_macro_f1_nonnone={score_now:.6f} " + f"none_fpr={none_fpr_text} best={best_text} saved_best={saved_best}", + flush=True, + ) + elif save_best_unbind and exact > best_exact: + if device.type == "cuda": + torch.cuda.synchronize() + best_exact = exact + best_metrics = dict(metrics) + best_state = _build_weight_state( + encoder, classify, role_head, filler_head, maps, layout + ) + best_residual_records = list(last_residual_records) + best_dir = out_dir / "best" + write_skeleton(best_dir, maps=maps) + save_heads(best_dir, best_state) + (best_dir / "best-checkpoint.json").write_text( + json.dumps( + { + "metric": "unbind_exact", + "epoch": ep, + "unbind_exact": exact, + "unbind_token_f1": best_metrics.get("unbind_token_f1"), + "unbind_slot_f1": best_metrics.get("unbind_slot_f1"), + "unbind_exact_strict": best_metrics.get("unbind_exact_strict"), + "unbind_token_f1_strict": best_metrics.get("unbind_token_f1_strict"), + "unbind_slot_f1_strict": best_metrics.get("unbind_slot_f1_strict"), + "classify_acc": best_metrics.get("classify_acc"), + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + saved_best = True + if device.type == "cuda": + torch.cuda.synchronize() + torch.cuda.empty_cache() + # Durable heartbeat — morph68 hung silently after ep4 with frozen docker logs. + progress = { + "epoch": ep, + "unbind_exact": exact, + "best_unbind_exact": None if best_exact == float("-inf") else best_exact, + "saved_best": saved_best, + "n_unbind_phase": phase_meta["n_rows"], + "unbind_phase": phase_meta["phase"], + "classify_acc": metrics.get("classify_acc"), + } + if select_on_classify: + progress["classify_macro_f1_nonnone"] = metrics.get("classify_macro_f1_nonnone") + progress["none_fpr"] = metrics.get("none_fpr") + progress["best_classify_macro_f1_nonnone"] = ( + None if best_macro == float("-inf") else best_macro + ) + if seed_receipt is not None: + progress["seed"] = seed_receipt["seed"] + progress.update( + epoch_timing_fields( + epoch_started=epoch_started, + epoch_ended=monotonic_seconds(), + training_started=training_started, + ) + ) + with progress_path.open("a", encoding="utf-8") as pf: + pf.write(json.dumps(progress, sort_keys=True) + "\n") + pf.flush() + if not select_on_classify: + print( + f"[epoch {ep}] unbind_exact={exact:.6f} best={best_exact if best_exact != float('-inf') else None} saved_best={saved_best}", + flush=True, + ) + if early_stop_break( + enabled=early_stop.enabled, + epoch_index=ep, + best_epoch=best_epoch_index, + epochs_scored=ep + 1, + minimum_epochs=early_stop.minimum_epochs, + patience=early_stop.patience, + max_epochs=epochs, + ): + stop_reason = STOP_REASON_EARLY_STOPPING + break + + training_elapsed_seconds = round_seconds(monotonic_seconds() - training_started) + final_state = _build_weight_state( + encoder, classify, role_head, filler_head, maps, layout + ) + # Always write final epoch weights under a distinct name when best-save is on, + # then promote best → primary model.safetensors (fixes morph35 peak-not-saved). + if (save_best_unbind or select_on_classify) and best_state is not None: + final_file = save_heads(out_dir, final_state) + # rename primary final dump aside, then write best as primary + final_path = out_dir / final_file + aside = out_dir / ( + "model.final.safetensors" if final_file == "model.safetensors" else "heads.final.pt" + ) + if final_path.exists(): + final_path.replace(aside) + weight_file = save_heads(out_dir, best_state) + primary_val = best_metrics or {} + gated_from = ( + "best_classify_macro_f1_nonnone" if select_on_classify else "best_unbind_exact" + ) + residual_for_dump = best_residual_records + else: + weight_file = save_heads(out_dir, final_state) + primary_val = epoch_metrics[-1] if epoch_metrics else {} + gated_from = "final_epoch" + residual_for_dump = last_residual_records + aside = None + + last = epoch_metrics[-1] if epoch_metrics else {} + residual_receipt: dict = { + "unbind_residual_dump": "", + "n_unbind_residual": 0, + "unbind_residual_themes": {}, + } + if residual_dump_path: + residual_receipt = write_residual_dump(residual_dump_path, residual_for_dump) + receipt = { + "schema": "hyperlex.hyperlexical.train_receipt.v0.1", + "model_id": MODEL_ID_SEED, + "trunk": TRUNK, + "trunk_dir": str(trunk), + "device": str(device), + "cuda": bool(torch.cuda.is_available()), + "epochs": epochs, + "stop_reason": stop_reason, + "training_elapsed_seconds": training_elapsed_seconds, + **provenance(root), + "n_train_classify": len(classify_tr), + "task_accounting": task_accounting, + "n_train_unbind": len(unbind_tr), + "n_unfrozen_encoder": n_unfrozen, + "n_encoder_tensors": len(final_state.get("encoder") or {}), + "last_trainable": last_trainable_used, + "unbind_loss_weight": unbind_loss_weight, + "unbind_every_n": unbind_every_n, + "unbind_primary": slot_ce_mode["unbind_primary"], + "unbind_slot_ce_armed": slot_ce_mode["unbind_slot_ce_armed"], + "unbind_slot_ce_aux_lambda": slot_ce_mode["unbind_slot_ce_aux_lambda"], + "unbind_head_slot_weight": head_slot_weight, + "unbind_second_slot_weight": second_slot_weight, + "n_unbind_observed": unbind_recipe["n_unbind_observed"], + "n_unbind_inferred": unbind_recipe["n_unbind_inferred"], + "unbind_observed_upsample": unbind_recipe["unbind_observed_upsample"], + "unbind_inferred_cap": unbind_recipe["unbind_inferred_cap"], + "unbind_inferred_weight": inferred_weight, + "n_unbind_morph_negatives": unbind_recipe["n_unbind_morph_negatives"], + "unbind_morph_margin": morph_margin, + "unbind_curriculum": curriculum_plan["enabled"], + "unbind_curriculum_pos_epochs": curriculum_plan["pos_epochs"], + "unbind_curriculum_type_epochs": curriculum_plan["type_epochs"], + "unbind_curriculum_phases": curriculum_plan["phases"], + "n_unbind_curriculum_positional": curriculum_plan["n_unbind_positional"], + "n_unbind_curriculum_type_slot": curriculum_plan["n_unbind_type_slot"], + "n_unbind_curriculum_joint": curriculum_plan["n_unbind_joint"], + "unbind_filler_denylist_lineages": unbind_recipe.get("unbind_filler_denylist_lineages", 0), + "unbind_hard_atoms_path": unbind_recipe.get("unbind_hard_atoms_path", ""), + "unbind_hard_upsample": unbind_recipe.get("unbind_hard_upsample", 1), + "n_unbind_hard_atoms_matched": unbind_recipe.get("n_unbind_hard_atoms_matched", 0), + "unbind_force_train_path": unbind_recipe.get("unbind_force_train_path", ""), + "release_set": release_stats, + "filler_filter": unbind_recipe.get("filler_filter"), + "n_filler_rows_dropped_train": unbind_recipe.get("n_filler_rows_dropped_train", 0), + "n_filler_rows_dropped_val": unbind_recipe.get("n_filler_rows_dropped_val", 0), + "n_unbind_force_train": unbind_recipe.get("n_unbind_force_train", 0), + "n_unbind_force_train_keys": unbind_recipe.get("n_unbind_force_train_keys", 0), + "n_unbind_val_after_force_train": unbind_recipe.get( + "n_unbind_val_after_force_train", 0 + ), + "n_unbind_hard_extra_copies": unbind_recipe.get("n_unbind_hard_extra_copies", 0), + "unbind_residual_dump": residual_receipt.get("unbind_residual_dump", ""), + "n_unbind_residual": residual_receipt.get("n_unbind_residual", 0), + "unbind_residual_themes": residual_receipt.get("unbind_residual_themes", {}), + "unbind_residual_by_scheme": residual_receipt.get("unbind_residual_by_scheme", {}), + "unbind_residual_by_class": residual_receipt.get("unbind_residual_by_class", {}), + "unbind_residual_summary": residual_receipt.get("unbind_residual_summary", ""), + "last_loss": losses[-1] if losses else None, + "val": primary_val, + "val_final": last, + "val_best": best_metrics, + "save_best_unbind": save_best_unbind, + "primary_weights_from": gated_from, + "best_unbind_exact": None if best_metrics is None else best_metrics.get("unbind_exact"), + "best_epoch": None if best_metrics is None else best_metrics.get("epoch"), + "final_weight_file": None if aside is None else aside.name, + "init_from": init_receipt.get("init_from"), + "warm_start": bool(init_receipt.get("warm_start")), + "init_weight_file": init_receipt.get("weight_file"), + "init_encoder_trainable_loaded": init_receipt.get("encoder_trainable_loaded"), + "init_expand_vocab": bool(init_receipt.get("init_expand_vocab")), + "init_expand_vocab_applied": bool(init_receipt.get("expand_vocab")), + "init_expand_role": init_receipt.get("role"), + "init_expand_filler": init_receipt.get("filler"), + "epoch_metrics": epoch_metrics, + "weight_file": weight_file, + "aligner": "char_span + offset_mapping", + "data_sha256": input_receipt["data_sha256"], + "training_input_mode": input_receipt["training_input_mode"], + "training_export_path": input_receipt["training_export_path"], + "training_export_sha256_expected": input_receipt["training_export_sha256_expected"], + "training_export_sha256_actual": input_receipt["training_export_sha256_actual"], + "training_export_rows": input_receipt["training_export_rows"], + "live_export_generation_enabled": input_receipt["live_export_generation_enabled"], + "holdout": holdout_receipt(holdout_spec, holdout_removed), + "holdout_training_disjoint": disjoint_receipt, + "classify_admission": classify_admission_receipt, + "include_live": include_live, + "live_included": bundle["counts"].get("live_included", 0), + "name_gate": False, + "e2_pass": False, + "brier": None, + "forecast_eligible": False, + "note": "HF-shaped dump. Not Hyperlexical until E2.", + } + if reserve_receipt is not None: + receipt["eval_reserve_disjoint"] = reserve_receipt + if classify_split_receipt is not None: + receipt["classify_split"] = classify_split_receipt + if seed_receipt is not None: + receipt["seed"] = seed_receipt + if select_on_classify: + receipt["select_metric"] = SELECT_METRIC_CLASSIFY + if force_overlap.get("disjoint"): + receipt["force_train_disjoint"] = { + "n_overlap": force_overlap["n_overlap"], + "n_in_val": force_overlap["n_in_val"], + "n_in_selection": force_overlap["n_in_selection"], + "n_force_rows": force_overlap["n_force_rows"], + "n_dropped_classify_val": force_overlap["n_dropped_classify_val"], + "n_dropped_unbind_val": force_overlap["n_dropped_unbind_val"], + "n_classify_val": force_overlap["n_classify_val"], + "n_unbind_val": force_overlap["n_unbind_val"], + } + if resolve_vocab_train_only(): + receipt["vocab_train_only"] = True + (out_dir / "layout.json").write_text(json.dumps(layout, indent=2, sort_keys=True) + "\n", encoding="utf-8") + (out_dir / "train-receipt.json").write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8") + train_config = { + "lr": os.environ.get("HYPERLEX_TRAIN_LR", "2e-5"), + "epochs": epochs, + "batch": batch, + "max_len": MAX_LEN, + "last_trainable": last_trainable_used, + "unbind_loss_weight": unbind_loss_weight, + "unbind_every_n": unbind_every_n, + "unbind_primary": slot_ce_mode["unbind_primary"], + "unbind_slot_ce_armed": slot_ce_mode["unbind_slot_ce_armed"], + "unbind_slot_ce_aux_lambda": slot_ce_mode["unbind_slot_ce_aux_lambda"], + "unbind_head_slot_weight": head_slot_weight, + "unbind_second_slot_weight": second_slot_weight, + "unbind_observed_upsample": unbind_recipe["unbind_observed_upsample"], + "unbind_inferred_cap": unbind_recipe["unbind_inferred_cap"], + "unbind_inferred_weight": inferred_weight, + "n_unbind_morph_negatives": unbind_recipe["n_unbind_morph_negatives"], + "unbind_morph_margin": morph_margin, + "unbind_curriculum": curriculum_plan["enabled"], + "unbind_curriculum_pos_epochs": curriculum_plan["pos_epochs"], + "unbind_curriculum_type_epochs": curriculum_plan["type_epochs"], + "unbind_filler_denylist_lineages": unbind_recipe.get( + "unbind_filler_denylist_lineages", 0 + ), + "unbind_hard_atoms_path": unbind_recipe.get("unbind_hard_atoms_path", ""), + "unbind_hard_upsample": unbind_recipe.get("unbind_hard_upsample", 1), + "n_unbind_hard_atoms_matched": unbind_recipe.get( + "n_unbind_hard_atoms_matched", 0 + ), + "n_unbind_hard_extra_copies": unbind_recipe.get( + "n_unbind_hard_extra_copies", 0 + ), + "unbind_force_train_path": unbind_recipe.get("unbind_force_train_path", ""), + "n_unbind_force_train": unbind_recipe.get("n_unbind_force_train", 0), + "n_unbind_force_train_keys": unbind_recipe.get( + "n_unbind_force_train_keys", 0 + ), + "n_unbind_val_after_force_train": unbind_recipe.get( + "n_unbind_val_after_force_train", 0 + ), + "save_best_unbind": save_best_unbind, + "holdout": holdout_receipt(holdout_spec, holdout_removed), + "init_from": init_receipt.get("init_from"), + "warm_start": bool(init_receipt.get("warm_start")), + } + if seed_receipt is not None: + train_config["seed"] = seed_receipt["seed"] + (out_dir / "config-train.json").write_text( + json.dumps(train_config, indent=2) + "\n", + encoding="utf-8", + ) + return receipt