evidence/scripts/train_workflow_backbone.py

"""Shared-backbone adaptation with frozen-baseline replay and development-only checkpoints."""

import argparse
import hashlib
import json
import math
import random
import time
from collections import Counter
from pathlib import Path

import torch
from safetensors.torch import load_file
from train_foundation_head import write

from veyra.data import read_records
from veyra.option_model import OptionModel
from veyra.proper_learning import distribution_losses, retention_losses
from veyra.workflow_learning import SHARES, annotate, selection_key, summarize_logits


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def forward(model, record, root):
    states, mappings, _, conditions = model.encode(record.request, root, capture_condition=True)
    if states.shape[0] != 1 or len(mappings) != 1:
        raise ValueError("require one independent decision")
    logits = model.readout(states) + model.binding_head(states, model.condition_readout(conditions))
    return logits, states


@torch.inference_mode()
def assess(model, records, root, rows, tensors, indices, temperatures):
    model.eval()
    logits = torch.stack(
        [forward(model, records[rows[i]["id"]], root)[0][0].cpu() for i in indices.tolist()]
    )
    return summarize_logits(logits, tensors, rows, indices, temperatures)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--learning-rate", type=float, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--records", type=Path, default=Path("data/workflow-v12/records.jsonl"))
    parser.add_argument("--features", type=Path, default=Path("data/features-workflow-v12"))
    parser.add_argument("--head-experiment", type=Path, default=Path("runs/workflow-v12-head"))
    args = parser.parse_args()
    if args.output.exists():
        raise FileExistsError("backbone experiment output is immutable")
    configuration = Path("configs/workflow-release-v12.json")
    config = json.loads(configuration.read_text(encoding="utf-8"))
    if (
        args.learning_rate not in config["backbone_learning_rates"]
        or config["backbone_epochs"] != 1
    ):
        raise ValueError("run outside the frozen experiment plan")
    selected = json.loads((args.head_experiment / "selection.json").read_text(encoding="utf-8"))[
        "selected"
    ]
    parent = Path(selected["checkpoint"])
    parent_manifest = json.loads((parent / "manifest.json").read_text(encoding="utf-8"))
    temperatures = parent_manifest["calibration"]["temperatures"]
    baseline = json.loads((args.head_experiment / "baseline.json").read_text(encoding="utf-8"))[
        "dev"
    ]
    spec = json.loads((args.features / "specification.json").read_text(encoding="utf-8"))
    if digest(args.records) != spec["records_sha256"]:
        raise ValueError("data changed after feature extraction")
    tensors = load_file(args.features / "features.safetensors")
    rows = json.loads((args.features / "records.json").read_text(encoding="utf-8"))
    records = {r.id: r for r in read_records(args.records) if r.split in {"train", "dev"}}
    annotate(rows, records.values())
    train = [i for i, r in enumerate(rows) if r["split"] == "train"]
    dev = torch.tensor([i for i, r in enumerate(rows) if r["split"] == "dev"])
    groups = {}
    for i in train:
        groups.setdefault(rows[i]["group"], []).append(i)
    units = list(groups.values())
    random.Random(config["seed"]).shuffle(units)
    order = [i for unit in units for i in unit]
    segments = [order[: len(order) // 2], order[len(order) // 2 :]]
    counts = Counter(rows[i]["domain"] for i in train)
    weights = {d: SHARES[d] * len(train) / n for d, n in counts.items()}
    protocol = {
        "arguments": vars(args),
        "release_protocol_sha256": digest(configuration),
        "parent": str(parent),
        "parent_weights_sha256": digest(parent / "head.safetensors"),
        "dataset_sha256": digest(args.records),
        "method": selected["method"],
        "seed": config["seed"],
        "domain_shares": SHARES,
        "accumulation": 8,
        "head_lr_multiplier": 0.1,
        "retention": (
            "0.25 supervised + 4*temperature-2 KL + 0.1 cosine feature anchoring on replay"
        ),
        "source_sha256": {
            str(p): digest(p)
            for p in (
                Path(__file__),
                Path("src/veyra/workflow_learning.py"),
                Path("src/veyra/proper_learning.py"),
                Path("src/veyra/option_model.py"),
            )
        },
        "condition_readout_frozen": True,
        "vision_weights_frozen": True,
        "final_or_calibration_used": False,
    }
    write(args.output / "protocol.json", protocol)
    torch.manual_seed(config["seed"])
    torch.set_num_threads(4)
    model = OptionModel.load(parent, local_files_only=True, merge=False)
    model.condition_readout.requires_grad_(False)
    adapters = [p for name, p in model.named_parameters() if ".lora_" in name]
    optimizer = torch.optim.AdamW(
        [
            {"params": adapters, "lr": args.learning_rate},
            {
                "params": list(model.readout.parameters()) + list(model.binding_head.parameters()),
                "lr": args.learning_rate * 0.1,
            },
        ],
        weight_decay=0.01,
    )
    total_steps = sum(math.ceil(len(segment) / 8) for segment in segments)

    def rate(step):
        if step < 30:
            return (step + 1) / 30
        return 0.1 + 0.45 * (1 + math.cos(math.pi * min(1, (step - 30) / max(1, total_steps - 30))))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, rate)
    started, processed, steps, total_loss, history = time.perf_counter(), 0, 0, 0.0, []
    torch.cuda.reset_peak_memory_stats()
    for segment_index, segment in enumerate(segments):
        model.train()
        model.encoder.model.eval()
        for start in range(0, len(segment), 8):
            batch = segment[start : start + 8]
            optimizer.zero_grad(set_to_none=True)
            for i in batch:
                logits, hidden = forward(model, records[rows[i]["id"]], args.records.parent)
                valid = tensors["valid"][i : i + 1].to(logits.device)
                types = tensors["types"][i : i + 1].to(logits.device)
                scale = torch.tensor(temperatures, device=logits.device)[types, None]
                scaled = logits / scale
                loss = distribution_losses(
                    scaled,
                    tensors["targets"][i : i + 1].to(logits.device),
                    valid,
                    types,
                    selected["method"],
                ).mean()
                if rows[i]["retention"]:
                    reference = tensors["parent_logits"][i : i + 1].to(logits.device) / scale
                    anchoring = (
                        1
                        - torch.nn.functional.cosine_similarity(
                            hidden, tensors["states"][i : i + 1].to(hidden.device)
                        ).mean()
                    )
                    loss = (
                        0.25 * loss
                        + 4 * retention_losses(scaled, reference, valid).mean()
                        + 0.1 * anchoring
                    )
                loss = loss * weights[rows[i]["domain"]]
                if not torch.isfinite(loss):
                    raise ValueError("nonfinite training objective")
                (loss / len(batch)).backward()
                total_loss += float(loss.detach())
                processed += 1
            torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], 1.0, error_if_nonfinite=True
            )
            optimizer.step()
            scheduler.step()
            steps += 1
            if steps % 25 == 0:
                progress = {
                    "processed": processed,
                    "total": len(train),
                    "steps": steps,
                    "mean_loss": total_loss / processed,
                    "elapsed_seconds": time.perf_counter() - started,
                }
                write(args.output / "progress.json", progress)
                print(json.dumps(progress), flush=True)
        metrics = assess(model, records, args.records.parent, rows, tensors, dev, temperatures)
        path = args.output / ("half" if segment_index == 0 else "full")
        metadata = {
            "completed": True,
            "intermediate": True,
            "optimizer_steps": steps,
            "method": selected["method"],
            "seed": config["seed"],
            "epochs": 0.5 * (segment_index + 1),
            "backbone_updated": True,
            "parent_training": parent_manifest["training"],
            "dataset_sha256": digest(args.records),
            "protocol_sha256": digest(args.output / "protocol.json"),
            "selected_dev": metrics,
            "selection": selection_key(metrics, baseline),
            "elapsed_seconds": time.perf_counter() - started,
            "peak_memory_mib": torch.cuda.max_memory_allocated() / 2**20,
            "final_evaluated_at_selection": False,
            "smoke": False,
        }
        model.save(path, metadata)
        # OptionModel.save preserves this model's frozen parent calibration until final fitting.
        history.append({"checkpoint": str(path), **metadata})
        write(args.output / "history.json", history)
        print(
            json.dumps(
                {
                    "checkpoint": str(path),
                    "selection": metadata["selection"],
                    "domains": metrics["by_domain"],
                }
            ),
            flush=True,
        )
    write(
        args.output / "complete.json",
        {
            "completed": True,
            "candidates": [h["checkpoint"] for h in history],
            "requires_merged_dev_selection": True,
        },
    )


if __name__ == "__main__":
    main()