File size: 8,284 Bytes
56ae5ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | #!/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()
|