| """Hand-written CiteNet -> Core ML (.mlpackage, fp16), plus parity and placement checks.""" |
| import json, os, shutil, sys, time |
| import numpy as np, torch, coremltools as ct |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from citenet import net as cn, text, spec |
| from lepc import load |
|
|
| OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build") |
| |
| |
| |
| BATCHES = [1, 32, 512, 2048, 8192] |
|
|
| model, meta = cn.make() |
| tot, nz = cn.param_counts(model) |
| paths = {} |
| for B in BATCHES: |
| traced = torch.jit.trace(model, torch.zeros(B, 1, 1, spec.SEQ)) |
| ml = ct.convert( |
| traced, |
| inputs=[ct.TensorType(name="codes", shape=(B, 1, 1, spec.SEQ), dtype=np.float16)], |
| outputs=[ct.TensorType(name=n, dtype=np.float16) for n in ("p_fake", "evidence", "concepts")], |
| compute_precision=ct.precision.FLOAT16, |
| minimum_deployment_target=ct.target.macOS15, |
| ) |
| ml.short_description = "CiteNet: phantom legal-citation detector with hand-written weights (no training)" |
| ml.user_defined_metadata.update({"params": str(tot), "nonzero_params": str(nz), "trained": "false", |
| "batch": str(B), "evidence": json.dumps(meta["evidence"])}) |
| paths[B] = os.path.join(OUT, f"CiteNet_b{B}.mlpackage") |
| ml.save(paths[B]) |
| print("saved", paths[B]) |
| |
| class Lean(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__(); self.m = m |
| def forward(self, x): |
| p, e, _ = self.m(x) |
| return p, e |
| for B in (512, 2048, 8192): |
| traced = torch.jit.trace(Lean(model).eval(), torch.zeros(B, 1, 1, spec.SEQ)) |
| ml = ct.convert(traced, inputs=[ct.TensorType(name="codes", shape=(B, 1, 1, spec.SEQ), dtype=np.float16)], |
| outputs=[ct.TensorType(name=n, dtype=np.float16) for n in ("p_fake", "evidence")], |
| compute_precision=ct.precision.FLOAT16, minimum_deployment_target=ct.target.macOS15) |
| ml.short_description = "CiteNet lean (p_fake + evidence only), hand-written weights" |
| lp = os.path.join(OUT, f"CiteNet_lean_b{B}.mlpackage"); ml.save(lp) |
| dst = os.path.join(OUT, f"CiteNet_lean_b{B}.mlmodelc"); shutil.rmtree(dst, ignore_errors=True) |
| ct.utils.compile_model(lp, dst) |
| print("saved", lp) |
| json.dump({k: meta[k] for k in ("features", "concepts", "rules", "evidence", "groups", "anchored", "floating")}, |
| open(os.path.join(OUT, "meta.json"), "w")) |
| np.savez_compressed(os.path.join(OUT, "weights.npz"), **{k: v for k, v in cn.build()[0].items()}) |
|
|
| |
| rows = load("aux_train") + load("eval") |
| X = text.encode([r["cite"] for r in rows]) |
| with torch.no_grad(): |
| pt = model(torch.from_numpy(X))[0][:, 0].numpy() |
| for B, cu in ((512, "CPU_AND_NE"), (2048, "CPU_AND_GPU")): |
| m = ct.models.MLModel(paths[B], compute_units=getattr(ct.ComputeUnit, cu)) |
| outs = [] |
| for i in range(0, len(X), B): |
| xb = X[i:i + B]; n = len(xb) |
| pad = np.zeros((B, 1, 1, spec.SEQ), np.float16); pad[:n] = xb |
| outs.append(m.predict({"codes": pad})["p_fake"][:n, 0]) |
| pc = np.concatenate(outs).astype(np.float32) |
| flip = np.where((pc >= 0.5) != (pt >= 0.5))[0] |
| print(f"{cu} b{B} vs torch fp32: {len(X)-len(flip)}/{len(X)} labels agree; flips: {[rows[i]['cite'] for i in flip]}") |
|
|
| |
| from coremltools.models.compute_plan import MLComputePlan |
| for B, p in paths.items(): |
| dst = os.path.join(OUT, f"CiteNet_b{B}.mlmodelc") |
| shutil.rmtree(dst, ignore_errors=True) |
| c = ct.utils.compile_model(p, dst) |
| for cu in ("CPU_AND_NE", "CPU_AND_GPU"): |
| plan = MLComputePlan.load_from_path(path=c, compute_units=getattr(ct.ComputeUnit, cu)) |
| cnt = {} |
| for fn in plan.model_structure.program.functions.values(): |
| for op in fn.block.operations: |
| u = plan.get_compute_device_usage_for_mlprogram_operation(op) |
| if u: |
| k = type(u.preferred_compute_device).__name__.replace("ML", "").replace("ComputeDevice", "") |
| cnt[k] = cnt.get(k, 0) + 1 |
| print(f"b{B} {cu} placement:", cnt) |
|
|
