| """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") |
|
|
