"""One immutable merged-model evaluation; optional calibration uses separate groups.""" import argparse import hashlib import json import shutil import time from pathlib import Path import torch from train_foundation_head import write from veyra.calibrate import fit_abstention, fit_temperature from veyra.candidates import candidates_for from veyra.constants import QUESTION_TYPES from veyra.data import read_records from veyra.decision_metrics import observation, report from veyra.features import FeatureSample from veyra.option_model import OptionModel from veyra.probability import Calibration, typed_answer def digest(path): return hashlib.sha256(path.read_bytes()).hexdigest() @torch.inference_mode() def main(): parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--records", type=Path, required=True) parser.add_argument("--split", choices=["dev", "calibration", "test"], required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--calibrated-output", type=Path) parser.add_argument("--legacy-regression", action="store_true") args = parser.parse_args() if args.output.exists() and any(args.output.iterdir()): raise FileExistsError("evaluation output must be empty") if args.calibrated_output and args.split != "calibration": raise ValueError("only the calibration split can fit calibration") records = [r for r in read_records(args.records) if r.split == args.split] if not records: raise ValueError("empty evaluation") protocol = { "arguments": vars(args), "weights_sha256": digest(args.checkpoint / "head.safetensors"), "manifest_sha256": digest(args.checkpoint / "manifest.json"), "records_sha256": digest(args.records), "source_sha256": digest(Path(__file__)), "dependency_source_sha256": { str(path): digest(path) for path in ( Path("src/veyra/decision_metrics.py"), Path("src/veyra/probability.py"), Path("src/veyra/option_model.py"), Path("src/veyra/calibrate.py"), ) }, "subset_sha256": hashlib.sha256( "".join(r.model_dump_json() + "\n" for r in records).encode() ).hexdigest(), "metric": ( "All probability-argmax decisions, including abstained answers, " "against explicit hard gold; soft targets scored separately" ), "merged_bf16_deployment": True, "probability_serialization": "typed_answer float32 softmax on the model device", } write(args.output / "protocol.json", protocol) model = OptionModel.load(args.checkpoint, local_files_only=True, merge=True) torch.set_num_threads(4) samples, logits, descriptors = [], [], [] started = time.perf_counter() for index, record in enumerate(records, start=1): output, _ = model(record.request, args.records.parent) for name, value in output.items(): question = record.request.questions[name] candidates = candidates_for(question) samples.append( FeatureSample( features=torch.zeros(len(candidates), 1), context=torch.zeros(1), levels=torch.tensor([c.level for c in candidates]), targets=torch.tensor([record.targets[name][c.key] for c in candidates]), type_id=QUESTION_TYPES.index(question.type), record_id=record.id, group_id=record.group_id, split=record.split, family=record.family, language=record.language, question_id=name, tags=record.tags, ) ) logits.append(value.cpu().float()) descriptors.append((record, name, question)) if index % 200 == 0: print( json.dumps({"evaluated": index, "total": len(records), "split": args.split}), flush=True, ) calibration = model.calibration fitting_report = None if args.calibrated_output: temperatures, thresholds, always, fitting_report = [], [], [], {} for type_id, kind in enumerate(QUESTION_TYPES): fitting, policy = [], [] for i, sample in enumerate(samples): if sample.type_id == type_id: bucket = int(hashlib.sha256(sample.group_id.encode()).hexdigest()[:8], 16) % 2 (fitting if bucket == 0 else policy).append(i) temperature = max( 1.0, fit_temperature([samples[i] for i in fitting], [logits[i] for i in fitting]) ) selection = fit_abstention( [samples[i] for i in policy], [logits[i] for i in policy], temperature ) temperatures.append(temperature) thresholds.append(selection["threshold"]) if selection["always_abstain"]: always.append(kind) fitting_report[kind] = { **selection, "temperature": temperature, "temperature_questions": len(fitting), "policy_questions": len(policy), } calibration = Calibration( tuple(temperatures), tuple(thresholds), tuple(QUESTION_TYPES), tuple(QUESTION_TYPES), tuple(always), ) if args.calibrated_output.exists() and any(args.calibrated_output.iterdir()): raise FileExistsError("calibrated checkpoint output must be empty") args.calibrated_output.mkdir(parents=True, exist_ok=True) shutil.copyfile( args.checkpoint / "head.safetensors", args.calibrated_output / "head.safetensors" ) manifest = json.loads((args.checkpoint / "manifest.json").read_text()) manifest["calibration"] = calibration.to_dict() manifest["training"]["intermediate"] = False manifest["training"]["calibration"] = { "split": "calibration", "temperature_floor": 1.0, "group_disjoint": True, "protocol": protocol, "fits": fitting_report, } write(args.calibrated_output / "manifest.json", manifest) rows = [] for (record, name, question), value in zip(descriptors, logits, strict=True): answer = typed_answer(question, value.to(model.encoder.device), calibration) row = observation(record, name, answer["probabilities"], abstained=answer["abstained"]) row["confidence"] = answer["confidence"] if args.legacy_regression: row["scope"] = "legacy_regression" rows.append(row) prediction_path = args.output / "predictions.jsonl" with prediction_path.open("x", encoding="utf-8", newline="\n") as stream: for row in rows: stream.write(json.dumps(row, ensure_ascii=False) + "\n") result = { "protocol": protocol, **report(rows, bootstrap=args.split == "test"), "calibration": calibration.to_dict(), "fitting": fitting_report, "predictions_sha256": digest(prediction_path), "elapsed_seconds": time.perf_counter() - started, } write(args.output / "evaluation.json", result) print(json.dumps({"overall": result["overall"], "by_domain": result["by_domain"]}), flush=True) if __name__ == "__main__": main()