File size: 7,623 Bytes
f9d0103 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """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()
|