File size: 4,583 Bytes
42b4a6a | 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 | """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")
# Fixed shapes: with EnumeratedShapes Core ML ran every shape on the CPU. With a
# fixed batch the planner puts all ops on the ANE from B=16 up; B=1 stays on CPU.
# b2048 / b8192 are for the GPU, which wants bigger batches.
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])
# lean builds for bulk scanning: no concepts output (only p_fake + evidence leave the accelerator)
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()})
# ---- parity: torch fp32 vs Core ML fp16 on every train+eval citation
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]}")
# ---- placement
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) # compile_model will not overwrite
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)
|