"""Package the per-L bucket models into ONE multifunction .mlpackage (functions L8, L16, ...; default L16),
coremltools save_multifunction dedups the shared weights. usage: make_multifunction.py 25m fp16_emb [8,16,32,64,256]"""
import sys, subprocess
import coremltools as ct
size, var = sys.argv[1], sys.argv[2]
Ls = [int(x) for x in (sys.argv[3] if len(sys.argv) > 3 else "8,16,32,64,256").split(",")]
d = ct.utils.MultiFunctionDescriptor()
for L in Ls:
    d.add_function(f"models/enc{size}_L{L}_{var}.mlpackage", src_function_name="main", target_function_name=f"L{L}")
d.default_function_name = "L16"
out = f"models/enc{size}_multi_{var}.mlpackage"
ct.utils.save_multifunction(d, out)
du = lambda p: subprocess.run(["du", "-sk", p], capture_output=True, text=True).stdout.split()[0]
sep = sum(int(du(f"models/enc{size}_L{L}_{var}.mlpackage")) for L in Ls)
print(f"{out}: {int(du(out))/1024:.1f} MB  vs separate packages {sep/1024:.1f} MB")