| """vela3-fast span adapters for Core ML: build the release SchemaNet (frozen backbone + adapter), merge the LoRA deltas into plain |
| Linear weights (W + alpha / r * B A), and export the merged encoder with the same fixed-length graph as the decision backbone. |
| |
| python span_export.py # parity: merged export vs the release SchemaNet hidden states |
| python span_export.py convert L [adapter] # -> coreml/vela3_<adapter>_L<L>.mlpackage (default adapter: span_sets)""" |
| import sys |
| import torch |
| from torch import nn |
| from vela3_export import Vela3EncoderExport, load_decision |
|
|
| sys.path.insert(0, "src") |
| import vela3_schema as VS |
|
|
|
|
| def build(adapter="span_sets"): |
| dec = load_decision() |
| return dec, VS.build_adapter(dec, f"src/{adapter}_adapter.safetensors", "cpu") |
|
|
|
|
| @torch.no_grad() |
| def merge_lora(net): |
| """Replace every LoRALinear by an nn.Linear with the merged weight (a copy: the shared backbone tensors are untouched).""" |
| n = 0 |
| for mod in list(net.modules()): |
| for name, child in list(mod.named_children()): |
| if isinstance(child, VS.LoRALinear): |
| b = child.base |
| lin = nn.Linear(b.in_features, b.out_features, bias=b.bias is not None) |
| lin.weight.copy_(b.weight + child.scale * (child.lora_B @ child.lora_A)) |
| if b.bias is not None: |
| lin.bias.copy_(b.bias) |
| setattr(mod, name, lin); n += 1 |
| return n |
|
|
|
|
| if __name__ == "__main__": |
| adapter = sys.argv[3] if len(sys.argv) > 3 else "span_sets" |
| dec, net = build(adapter) |
| if sys.argv[1:2] == ["convert"]: |
| import numpy as np, coremltools as ct |
| import ct_patch |
| ids = torch.randint(5, 49152, (1, 300)); ids[0, 0] = 0 |
| with torch.no_grad(): |
| ref = net.hidden(ids, torch.ones_like(ids))[0] |
| print("merged LoRA layers:", merge_lora(net)) |
| if sys.argv[1:2] != ["convert"]: |
| for L in (300, 512): |
| g = Vela3EncoderExport(net.encoder, L).eval() |
| x = torch.zeros(1, L, dtype=torch.int32); x[0, :300] = ids[0]; mk = torch.zeros(1, L, dtype=torch.int32); mk[0, :300] = 1 |
| h = g(x, mk)[0, :300] |
| print(f"L={L} max|dh|={(h - ref).abs().max():.2e} min cos={nn.functional.cosine_similarity(h, ref, dim=-1).min():.6f}") |
| sys.exit() |
| L = int(sys.argv[2]) |
| g = Vela3EncoderExport(net.encoder, L).eval() |
| x = torch.zeros(1, L, dtype=torch.int32); mk = torch.ones(1, L, dtype=torch.int32) |
| with torch.no_grad(): |
| tr = torch.jit.trace(g, (x, mk)) |
| ml = ct.convert(tr, inputs=[ct.TensorType("input_ids", shape=(1, L), dtype=np.int32), ct.TensorType("attention_mask", shape=(1, L), dtype=np.int32)], |
| outputs=[ct.TensorType("hidden", dtype=np.float16)], minimum_deployment_target=ct.target.macOS15, |
| compute_precision=ct.precision.FLOAT16, convert_to="mlprogram") |
| ml.short_description = f"vela3-fast UltraBERT-380M encoder + {adapter} adapter (LoRA merged), fixed length {L}" |
| ml.save(f"coreml/vela3_{adapter}_L{L}.mlpackage"); print("saved", f"coreml/vela3_{adapter}_L{L}.mlpackage") |
|
|
