File size: 2,816 Bytes
616b1e2 | 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 | """Compare the no-cache Core ML boundary with RLCD's cached native engine."""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import torch
from decision import CandidateScorer, load_native, source_path
from fixtures import CONTEXT, SCHEMA
from preprocessing import Shape, batch_arrays, prepare_candidates, select_values
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--length", type=int, default=256)
parser.add_argument("--batch", type=int, default=8)
parser.add_argument("--device", choices=["cpu", "mps"], default="cpu")
args = parser.parse_args()
torch.set_num_threads(2)
snapshot = source_path()
sys.path.insert(0, str(snapshot))
from rlcd.engine import Engine # noqa: E402
tokenizer, model = load_native(args.device)
upstream = Engine.__new__(Engine)
upstream.model = model
upstream.tokenizer = tokenizer
upstream.device = args.device
upstream.dtype = "float32"
shape = Shape(length=args.length, candidates=args.batch)
candidates = prepare_candidates(tokenizer, CONTEXT, SCHEMA, shape)
scorer = CandidateScorer(model, shape.length).eval()
started = time.perf_counter()
with torch.inference_mode():
native = upstream.constrained(CONTEXT, SCHEMA)
ours = []
for offset in range(0, len(candidates), shape.candidates):
group = candidates[offset : offset + shape.candidates]
arrays = batch_arrays(tokenizer, group, shape)
tensors = [torch.from_numpy(value).to(args.device) for value in arrays.values()]
ours.extend(scorer(*tensors).float().cpu().tolist()[: len(group)])
expected = [entry["log_likelihood"] for entries in native["scores"].values() for entry in entries]
errors = [abs(left - right) for left, right in zip(ours, expected)]
same_answer = select_values(candidates, ours) == json.loads(native["text"])
report = {
"source_repo": "notnotsamuel/LFM2.5-350M-RLCD",
"source_revision": "deb589d803d141cabd158ef55f6617b128529f36",
"device": args.device,
"shape": vars(shape),
"candidates": len(candidates),
"max_log_likelihood_error": max(errors),
"same_selected_values": same_answer,
"seconds": round(time.perf_counter() - started, 2),
"expected": expected,
"actual": ours,
}
output = Path("build/native-parity.json")
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2), flush=True)
if not same_answer or max(errors) > 1e-2:
raise SystemExit("native cached/no-cache parity failed")
if __name__ == "__main__":
main()
|