"""Build the FluidInference/vela-3.0-coreml release folder (release/):
  Vela3Decision.mlpackage   multifunction fp16 encoder (decision backbone), functions L128 / L256 / L512 / L1024
  Vela3SpanSets.mlpackage   multifunction fp16 encoder (backbone + span_sets adapter, LoRA merged), functions L512 / L1024
  heads_decision.safetensors, heads_span_sets.safetensors   fp32 heads (run on the host)
  tokenizer.json (base tokenizer, the 8 marker tokens removed: markers are inserted by id), config.json, vela3_config.json
  upstream/ the release's own Python (Apache-2.0, unmodified) that renders inputs, decodes spans and picks thresholds"""
import json, shutil, sys
from pathlib import Path
import coremltools as ct
import torch
from safetensors.torch import load_file, save_file

OUT = Path("release"); OUT.mkdir(exist_ok=True)
for name, src, Ls in (("Vela3Decision", "vela3_decision", (128, 256, 512, 1024)), ("Vela3SpanSets", "vela3_span_sets", (512, 1024))):
    d = ct.utils.MultiFunctionDescriptor()
    for L in Ls:
        d.add_function(f"coreml/{src}_L{L}.mlpackage", src_function_name="main", target_function_name=f"L{L}")
    d.default_function_name = f"L{Ls[0]}"
    ct.utils.save_multifunction(d, str(OUT / f"{name}.mlpackage"))
    print("packed", name)

# heads: decision (from model.safetensors) and span_sets (the adapter's own heads; everything outside the encoder)
sd = load_file("src/model.safetensors")
save_file({k: v.float().contiguous() for k, v in sd.items() if not k.startswith("encoder.")}, str(OUT / "heads_decision.safetensors"))
ad = load_file("src/span_sets_adapter.safetensors")
heads = {k: v.float().contiguous() for k, v in ad.items() if not k.startswith("encoder.")}
for k, v in sd.items():  # heads the adapter did not change come from the backbone file
    if not k.startswith("encoder.") and k not in heads:
        heads[k] = v.float().contiguous()
save_file(heads, str(OUT / "heads_span_sets.safetensors"))
print("heads", sorted(heads))

cfg = json.load(open("src/config.json")); mk = cfg["decision"]["marker_ids"]
tj = json.load(open("src/tokenizer.json", encoding="utf-8"))
json.dump(tj, open(OUT / "tokenizer.json", "w", encoding="utf-8"), ensure_ascii=False)  # the release's own file (its loaders strip the markers)
for f in ("config.json", "vela3_config.json", "tokenizer_config.json", "LICENSE"):
    shutil.copy(f"src/{f}", OUT / f)
up = OUT / "upstream"; up.mkdir(exist_ok=True)
for f in ("vela3_encoder.py", "vela3_schema.py", "ub_d3_common.py"):
    shutil.copy(f"src/{f}", up / f)
print("done", sorted(p.name for p in OUT.iterdir()))