File size: 3,193 Bytes
deaf348 | 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 | """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 # noqa: E402
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 # noqa: F401
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")
|