File size: 2,521 Bytes
628e2fd | 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 | """Agreement of shipped Core ML packages with native Fast Decisions predictions on bucket-fitting heads."""
import argparse
import json
import re
from pathlib import Path
import coremltools as ct
import numpy as np
from huggingface_hub import snapshot_download
from convert_names import MODEL_ID, MODEL_REVISION
from preprocessing import load_processor, prepare_decision
from runtime import decode
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--native-rows", default="build/fast-decisions-native.rows.jsonl")
parser.add_argument("--per-domain", type=int, default=20)
parser.add_argument("--out", default="build/verify-packages.json")
parser.add_argument("packages", nargs="+")
args = parser.parse_args()
data = Path(snapshot_download("fastino/fast-decisions", repo_type="dataset",
revision="1a33070cabf94ce2e29105482dd2ef6c157ad7f2"))
processor = load_processor(snapshot_download(MODEL_ID, revision=MODEL_REVISION))
native = {(r["domain"], r["index"], r["task"]): r["predicted"]
for r in map(json.loads, Path(args.native_rows).open())}
domains = sorted({key[0] for key in native})
report = {}
for package in args.packages:
length, heads, options = map(int, re.search(r"_L(\d+)_H(\d+)_K(\d+)", package).groups())
model = ct.models.MLModel(package, compute_units=ct.ComputeUnit.ALL)
checked = agree = 0
for domain in domains:
used = 0
for index, line in enumerate((data / f"{domain}.jsonl").open()):
if used >= args.per_domain:
break
row = json.loads(line)
for head in row["output"]["classifications"]:
tasks = {head["task"]: head["labels"]}
try:
arrays = prepare_decision(processor, row["input"], tasks, length, heads, options)
except ValueError:
continue
got = decode(tasks, np.asarray(model.predict(arrays)["logits"])[0])[head["task"]]
checked += 1
agree += [got["label"]] == native[(domain, index, head["task"])]
used += 1
report[Path(package).name] = {"checked_heads": checked, "agree_with_native": agree}
print(Path(package).name, checked, agree, flush=True)
Path(args.out).write_text(json.dumps(report, indent=2) + "\n")
if __name__ == "__main__":
main()
|