File size: 3,502 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 | """Measure backbone calls for real text and photograph requests on the frozen model."""
import argparse
import hashlib
import json
import math
from pathlib import Path
import torch
from train_foundation_head import write
from veyra.data import read_records
from veyra.option_model import OptionModel
from veyra.schema import DecisionRequest
def digest(path):
return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--records", type=Path, default=Path("data/foundation-v11/records.jsonl"))
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
if args.output.exists():
raise FileExistsError("contract evidence is immutable")
records = sorted(
(r for r in read_records(args.records) if r.split == "test"), key=lambda r: r.id
)
photo = next(r for r in records if "domain:image_waste" in r.tags)
text = next(r for r in records if not r.request.state.images)
questions = {}
for row in records:
if row.group_id == photo.group_id:
for question in row.request.questions.values():
questions.setdefault(question.type, question)
if set(questions) != {"choice", "score", "noul"}:
raise ValueError("photo fixture must cover all three decision types")
requests = {
"text": text.request,
"photograph_three_questions": DecisionRequest(
state=photo.request.state, questions=questions
),
}
torch.set_num_threads(4)
model = OptionModel.load(args.checkpoint, local_files_only=True, merge=True)
forwards, results = [], {}
hook = model.encoder.model.register_forward_hook(lambda *_: forwards.append(1))
try:
for name, request in requests.items():
start = len(forwards)
response = model.predict(request, args.records.parent.resolve())
count = len(forwards) - start
valid = all(
all(math.isfinite(p) and 0 <= p <= 1 for p in answer["probabilities"].values())
and abs(sum(answer["probabilities"].values()) - 1) <= 1e-5
for answer in response["answers"].values()
)
results[name] = {
"backbone_forwards": count,
"generated_answer_tokens": response["usage"]["output_tokens"],
"valid_probabilities": valid,
"questions": len(request.questions),
"types": sorted({q.type for q in request.questions.values()}),
"input_tokens": response["usage"]["input_tokens"],
}
finally:
hook.remove()
result = {
"weights_sha256": digest(args.checkpoint / "head.safetensors"),
"manifest_sha256": digest(args.checkpoint / "manifest.json"),
"records_sha256": digest(args.records),
"source_sha256": digest(Path(__file__)),
"gpu": torch.cuda.get_device_name(),
"record_ids": [text.id, photo.id],
"results": results,
"passed": all(
r["backbone_forwards"] == 1
and r["generated_answer_tokens"] == 0
and r["valid_probabilities"]
for r in results.values()
),
}
write(args.output, result)
print(json.dumps(result), flush=True)
if not result["passed"]:
raise SystemExit(1)
if __name__ == "__main__":
main()
|