Parakeet TDT 0.6B v3 for Core AI: 8-bit palettized encoder (0.65 GB), Mac and iPhone validated
56ae5ae verified | #!/usr/bin/env python3 | |
| """Export pinned original Parakeet TDT weights; save references and provenance. | |
| Writes a fresh portable bundle; no publication and no claimed parity until the | |
| native reference suite has run. Requires the existing apple/.venv toolchain. | |
| `--encoder-weights pal8` (or pal6) palettizes the encoder's weights with 8-bit (6-bit) | |
| per-tensor k-means, spec 159's technique, keeping its 16-bit inputs and compute. | |
| `--small-weights fp16` stores the predictor's and joint's weights in 16 bits behind | |
| their 32-bit interface. Either makes a smaller bundle with its own adapter ID | |
| (spec 172); the defaults reproduce the original FP16 bundle. | |
| """ | |
| import argparse | |
| import hashlib | |
| import importlib.metadata | |
| import json | |
| from pathlib import Path | |
| import shutil | |
| import time | |
| import traceback | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from parakeet_adapters import build_components, load_source | |
| class StoredHalf(nn.Module): | |
| """16-bit weights behind a 32-bit interface: inputs cast down, outputs cast up.""" | |
| def __init__(self, inner): | |
| super().__init__() | |
| self.inner = inner.half() | |
| def forward(self, *args): | |
| out = self.inner(*(x.half() if x.is_floating_point() else x for x in args)) | |
| return tuple(x.float() for x in out) if isinstance(out, tuple) else out.float() | |
| class HalfOutput(nn.Module): | |
| """The palettized encoder's output as 16-bit, as the FP16 encoder's was: the runner | |
| reads `encoded` as Float16 (Parakeet.swift's contract).""" | |
| def __init__(self, inner): | |
| super().__init__() | |
| self.inner = inner | |
| def forward(self, *args): | |
| return self.inner(*args).half() | |
| def palettize(module, bits, example): | |
| from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig | |
| preset = {8: KMeansPalettizerConfig.presets.w8, 6: KMeansPalettizerConfig.presets.w6}[bits]() | |
| palettizer = KMeansPalettizer(module, preset) | |
| palettizer.prepare(example, num_workers=8) | |
| return palettizer.finalize().eval() | |
| def sha(path): | |
| with Path(path).open("rb") as f: | |
| return hashlib.file_digest(f, "sha256").hexdigest() | |
| def main(): | |
| p = argparse.ArgumentParser(description=__doc__) | |
| p.add_argument("--source", type=Path, required=True) | |
| p.add_argument("--output", type=Path, required=True) | |
| p.add_argument("--resume", action="store_true", help="Keep already-exported components after an interrupted export") | |
| p.add_argument("--encoder-weights", choices=["fp16", "pal8", "pal6"], default="fp16") | |
| p.add_argument("--small-weights", choices=["fp32", "fp16"], default="fp32", help="Predictor and joint weights") | |
| a = p.parse_args() | |
| slim = a.encoder_weights != "fp16" or a.small_weights != "fp32" | |
| adapter_id = "parakeet-tdt-v3-coreai-batch15-v1" if not slim else f"parakeet-tdt-v3-coreai-{a.encoder_weights}-{a.small_weights}-v1" | |
| if a.output.exists() and not a.resume: | |
| p.error("Use a fresh output directory, or explicitly --resume the same source export") | |
| a.output.mkdir(parents=True, exist_ok=True) | |
| torch.set_num_threads(4) | |
| torch.manual_seed(1729) | |
| provenance = json.loads((a.source / "source-provenance.json").read_text()) | |
| provenance["weights_sha256"] = sha(a.source / "model.safetensors") | |
| report_path = a.output / "conversion.json" | |
| previous = json.loads(report_path.read_text()) if report_path.exists() else None | |
| if previous and previous["source"] != provenance: | |
| raise ValueError("Cannot resume an export from different source weights") | |
| report = previous or {"source": provenance, "status": "started", "components": {}, "weights": {"encoder": a.encoder_weights, "predictor_joint": a.small_weights}, | |
| "versions": {n: importlib.metadata.version(n) for n in ["torch", "transformers", "coreai-torch", "coreai-core"]}} | |
| def save(): | |
| report_path.write_text(json.dumps(report, indent=2, allow_nan=False)+"\n") | |
| save() | |
| try: | |
| model = load_source(a.source) | |
| from coreai_torch import TorchConverter, get_decomp_table | |
| for name, (graph, inputs, names, outputs, precision) in build_components(model).items(): | |
| if report["components"].get(name, {}).get("status") == "exported": | |
| continue | |
| start = time.perf_counter() | |
| entry = {"status": "started", "precision": precision} | |
| report["components"][name] = entry | |
| save() | |
| # Encoder/projector precision does not modify predictor or joint modules. | |
| graph = graph.eval() | |
| if name == "encoder" and a.encoder_weights != "fp16": | |
| graph = HalfOutput(palettize(graph, int(a.encoder_weights[3:]), inputs)) | |
| precision = a.encoder_weights + "-fp16" | |
| entry["precision"] = precision | |
| if name in ("predictor", "joint") and a.small_weights == "fp16": | |
| graph = StoredHalf(graph).eval() | |
| precision = "fp16-weights" | |
| entry["precision"] = precision | |
| if precision.endswith("fp16") and name == "encoder": | |
| graph.half() | |
| inputs = tuple(x.half() if x.is_floating_point() else x for x in inputs) | |
| with torch.no_grad(): | |
| result = graph(*inputs) | |
| result = result if isinstance(result, tuple) else (result,) | |
| if not all(torch.isfinite(x).all() for x in result): | |
| raise ValueError("Nonfinite PyTorch reference: " + name) | |
| reference = a.output / "references" | |
| reference.mkdir(exist_ok=True) | |
| np.savez(reference/(name+".npz"), **{n: x.numpy() for n,x in zip(names,inputs)}, **{"expected_"+n:x.numpy() for n,x in zip(outputs,result)}) | |
| # Portable raw fixtures consumed by SwiftParakeetComponentBench. | |
| for tensor_name, tensor in [*zip(names, inputs), *(("expected_"+n, x) for n,x in zip(outputs,result))]: | |
| tensor.float().numpy().astype("<f4").tofile(reference/(name+"-"+tensor_name+".f32")) | |
| entry.update(status="reference_ready", inputs=[{"name":n,"shape":list(x.shape),"dtype":str(x.dtype).removeprefix("torch.")} for n,x in zip(names,inputs)], outputs=outputs) | |
| save() | |
| ep = torch.export.export(graph, args=inputs).run_decompositions(get_decomp_table()) | |
| entry["status"] = "torch_exported" | |
| save() | |
| program = TorchConverter(mode=TorchConverter.Mode.RELEASE).add_exported_program(ep, input_names=names, output_names=outputs).to_coreai() | |
| entry["status"] = "coreai_lowered" | |
| save() | |
| program.optimize() | |
| program.save_asset(a.output/(name+".aimodel")) | |
| entry.update(status="exported", seconds=time.perf_counter()-start) | |
| save() | |
| print(json.dumps({"component":name, **entry}),flush=True) | |
| shutil.copy2(a.source/"tokenizer.json", a.output/"tokenizer.json") | |
| shutil.copy2(a.source/"README.md", a.output/"SOURCE_MODEL_CARD.md") | |
| files=[] | |
| for asset in sorted(a.output.glob("*.aimodel")): | |
| for f in sorted(asset.rglob("*")): | |
| if f.is_file(): files.append({"path":str(f.relative_to(a.output)),"bytes":f.stat().st_size,"sha256":sha(f)}) | |
| for f in [a.output/"tokenizer.json",a.output/"SOURCE_MODEL_CARD.md"]: | |
| files.append({"path":f.name,"bytes":f.stat().st_size,"sha256":sha(f)}) | |
| manifest={"adapter_id":adapter_id,"source":provenance,"license":"CC-BY-4.0","sample_rate":16000,"max_window_samples":240000,"max_utterance_samples":560000,"blank_id":8192,"durations":[0,1,2,3,4],"encoder_frames":188,"encoder_width":640,"frontend_frames":1500,"components":{name:{"asset":name+".aimodel","inputs":v["inputs"],"outputs":v["outputs"],"precision":v["precision"]} for name,v in report["components"].items()},"files":files,"validation":"pending_native_parity"} | |
| (a.output/"manifest.json").write_text(json.dumps(manifest,indent=2)+"\n") | |
| report["status"]="exported" | |
| save() | |
| except BaseException as error: | |
| report.update(status="failed",error=str(error),traceback=traceback.format_exc()) | |
| save() | |
| raise | |
| if __name__ == "__main__": | |
| main() | |