Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert moonshine-ai/moonshine streaming safetensors checkpoints to audio.cpp GGUF packages."""
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import struct
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
SPEC = REPO_ROOT / "model_specs" / "moonshine_asr.json"
|
|
SCALAR_TENSORS = {"model.encoder.embedder.comp.log_k"}
|
|
|
|
|
|
def require_v1_spec(path: Path) -> None:
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
spec = json.load(handle)
|
|
if spec.get("schema_version") != 1:
|
|
raise SystemExit(f"{path} must be a schema_version 1 model spec")
|
|
|
|
|
|
def rank1_scalar_safetensors(source: Path, destination: Path) -> None:
|
|
"""Copy safetensors while exposing known rank-0 scalars as length-1 tensors."""
|
|
with source.open("rb") as f:
|
|
header_len = struct.unpack("<Q", f.read(8))[0]
|
|
header = json.loads(f.read(header_len))
|
|
payload = f.read()
|
|
|
|
changed = False
|
|
for name in SCALAR_TENSORS:
|
|
entry = header.get(name)
|
|
if entry is not None and entry.get("shape") == []:
|
|
entry["shape"] = [1]
|
|
changed = True
|
|
|
|
if not changed:
|
|
shutil.copyfile(source, destination)
|
|
return
|
|
|
|
encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
|
padding = (8 - (len(encoded) % 8)) % 8
|
|
destination.write_bytes(struct.pack("<Q", len(encoded) + padding) + encoded + (b" " * padding) + payload)
|
|
|
|
|
|
def convert(converter: Path, checkpoint: Path, output: Path, quant_type: str, overwrite: bool) -> None:
|
|
require_v1_spec(SPEC)
|
|
ckpt = checkpoint / "model.safetensors"
|
|
if not ckpt.exists():
|
|
candidates = sorted(checkpoint.glob("*.safetensors"))
|
|
if not candidates:
|
|
raise SystemExit(f"No .safetensors checkpoint found in {checkpoint}")
|
|
ckpt = candidates[-1]
|
|
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="moonshine-asr-convert-") as tmp_dir:
|
|
converted_input = Path(tmp_dir) / ckpt.name
|
|
rank1_scalar_safetensors(ckpt, converted_input)
|
|
command = [
|
|
str(converter),
|
|
"--input",
|
|
str(converted_input),
|
|
"--root",
|
|
str(checkpoint),
|
|
"--family",
|
|
"moonshine_asr",
|
|
"--model-spec",
|
|
str(SPEC),
|
|
"--type",
|
|
quant_type,
|
|
"--output",
|
|
str(output),
|
|
]
|
|
if overwrite:
|
|
command.append("--overwrite")
|
|
print("+", " ".join(command))
|
|
subprocess.run(command, check=True)
|
|
|
|
for name in ["config.json", "tokenizer.json"]:
|
|
source = checkpoint / name
|
|
if source.exists():
|
|
shutil.copyfile(source, output.parent / name)
|
|
print(f"copied {name} -> {output.parent}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
parser.add_argument("--checkpoint", type=Path, required=True, help="HF checkpoint directory")
|
|
parser.add_argument("--converter", type=Path, required=True, help="path to audiocpp_gguf")
|
|
parser.add_argument("--output", type=Path, required=True, help="output .gguf path")
|
|
parser.add_argument(
|
|
"--type",
|
|
default="q8_0",
|
|
choices=["orig", "f32", "f16", "bf16", "q8_0", "q4_k", "q5_k", "q6_k"],
|
|
help="GGUF tensor storage type",
|
|
)
|
|
parser.add_argument("--overwrite", action="store_true", help="overwrite output")
|
|
args = parser.parse_args()
|
|
convert(args.converter, args.checkpoint, args.output, args.type, args.overwrite)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|