| """Task 2: convert QDec to Core ML. usage: convert.py B T [variant...]; variants fp16 | pal8 | lin8. |
| Writes models/dec_b{B}_t{T}_{variant}.mlpackage and appends one line per variant to logs/convert.log.""" |
| import sys, time, traceback, torch, numpy as np, coremltools as ct |
| import coremltools.optimize.coreml as cto |
| from pathlib import Path |
| from qdec import QDec, ane_kwargs |
|
|
| B, T = int(sys.argv[1]), int(sys.argv[2]) |
| variants = sys.argv[3:] or ["fp16"] |
| Path("models").mkdir(exist_ok=True); Path("logs").mkdir(exist_ok=True) |
| log = open("logs/convert.log", "a") |
| def L(msg): |
| print(msg); log.write(msg + "\n"); log.flush() |
|
|
| import os |
| CH = int(os.environ.get('CHUNK', '0')) or None |
| m = QDec(T, B=B, chunk=CH, **ane_kwargs()).eval() |
| SUF = f'c{CH}' if CH else '' |
| ex = torch.zeros((B, T), dtype=torch.int32) |
| t0 = time.time() |
| with torch.no_grad(): |
| tr = torch.jit.trace(m, ex.long()) |
| base = None |
| for v in variants: |
| name = f"models/dec_b{B}_t{T}_{v}{SUF}.mlpackage" |
| try: |
| t0 = time.time() |
| if base is None: |
| base = ct.convert(tr, inputs=[ct.TensorType("ids", shape=(B, T), dtype=np.int32)], |
| outputs=[ct.TensorType("lp", dtype=np.float32)], |
| compute_precision=ct.precision.FLOAT16, |
| minimum_deployment_target=ct.target.macOS15, convert_to="mlprogram") |
| mm = base |
| if v == "pal8": |
| mm = cto.palettize_weights(base, cto.OptimizationConfig(global_config=cto.OpPalettizerConfig(nbits=8, mode="kmeans"))) |
| elif v == "lin8": |
| mm = cto.linear_quantize_weights(base, cto.OptimizationConfig(global_config=cto.OpLinearQuantizerConfig(mode="linear_symmetric"))) |
| mm.save(name) |
| sz = sum(f.stat().st_size for f in Path(name).rglob("*") if f.is_file()) / 1e6 |
| L(f"OK b{B} t{T} {v}{SUF}: {time.time()-t0:.1f}s {sz:.1f} MB -> {name}") |
| except Exception as e: |
| L(f"FAIL b{B} t{T} {v}{SUF}: {type(e).__name__}: {str(e).splitlines()[0][:300]}") |
| traceback.print_exc() |
|
|
