From 2c247727698389361fd65366dfcfb5fab1955ac7 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 28 Jun 2026 18:29:19 +0800 Subject: [PATCH] gguf metadata correctness: streaming flags, file_type ordering, language fixes --- .../models/nemotron-3.5-asr-streaming-0.6b.md | 12 +- scripts/audit_gguf_metadata.py | 242 ++++++++++++++++++ scripts/convert-granite_nar.py | 27 +- scripts/convert-medasr.py | 24 ++ scripts/convert-moonshine.py | 63 ++++- scripts/convert-parakeet.py | 43 +++- scripts/convert-sensevoice.py | 1 + .../nemotron-3.5-asr-streaming-0.6b.yaml | 10 +- scripts/lib/test_quantize_bulk_trailer.py | 180 +++++++++++++ tools/transcribe-quantize/main.cpp | 91 +++++++ 10 files changed, 664 insertions(+), 29 deletions(-) create mode 100644 scripts/audit_gguf_metadata.py create mode 100644 scripts/lib/test_quantize_bulk_trailer.py diff --git a/docs/models/nemotron-3.5-asr-streaming-0.6b.md b/docs/models/nemotron-3.5-asr-streaming-0.6b.md index b7614a74..2d49c178 100644 --- a/docs/models/nemotron-3.5-asr-streaming-0.6b.md +++ b/docs/models/nemotron-3.5-asr-streaming-0.6b.md @@ -8,9 +8,10 @@ successor to ## What it's for -Multilingual speech-to-text across **40 language-locales** (19 -transcription-ready, plus broad-coverage and adaptation tiers — see the -upstream model card for the full list) with greedy RNN-T decoding. +Multilingual speech-to-text across **32 supported language-locales** (19 +transcription-ready + 13 broad-coverage; the tokenizer also recognizes 8 +adaptation-ready locales that require fine-tuning — see the upstream model +card for the full list) with greedy RNN-T decoding. Outputs cased, punctuated transcripts (native PnC). Token- and word-level timestamps are available. @@ -171,9 +172,10 @@ on WER (Stage 7), not tensor tolerances. ## Capabilities -- **Languages:** 40 language-locales (e.g. `en-US`, `en-GB`, `es-ES`, +- **Languages:** 32 supported language-locales (e.g. `en-US`, `en-GB`, `es-ES`, `fr-FR`, `de-DE`, `it-IT`, `pt-BR`, `nl-NL`, `ru-RU`, `zh-CN`, `ja-JP`, - `ko-KR`, `hi-IN`, `ar-AR`, …), selected via `--language `. + `ko-KR`, `hi-IN`, `ar-AR`, …), selected via `--language `. (The + tokenizer recognizes 40; the 8 adaptation-ready locales need fine-tuning.) - **Language detection:** `auto` mode emits `` locale tags (e.g. ``, ``) in the raw token stream. A tag can appear anywhere in the sequence, not only at the end. They are stripped from the returned diff --git a/scripts/audit_gguf_metadata.py b/scripts/audit_gguf_metadata.py new file mode 100644 index 00000000..ea23fd6f --- /dev/null +++ b/scripts/audit_gguf_metadata.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +"""Audit GGUF metadata for the header-query contract. + +For each GGUF it reports the four things that decide whether a published file +is "online-queryable" and self-describing: + + identity — general.* identity block complete? (name, size_label, author, + organization, license{,.name,.link}, languages) + trailer — are the bulk tokenizer arrays (tokens/scores/token_type/merges/ + chat_template) the LAST KVs, so a header range-read reaches the + small metadata without pulling them? (the #56 gguf_writer layout) + file_type — does general.file_type sit BEFORE the tokenizer trailer? (the + quantizer used to strand it dead-last) + streaming — is stt.capability.streaming present, and does it agree with what + the slug implies (streaming / unified / nemotron / realtime)? + +Usage + # local files / dirs / globs (default: models/**/*.gguf under the repo) + uv run --project scripts/envs/moonshine scripts/audit_gguf_metadata.py [PATH ...] + + # a published HF repo — range-fetches only the header of each .gguf + uv run --project scripts/envs/moonshine scripts/audit_gguf_metadata.py \ + --hf-repo handy-computer/parakeet-unified-en-0.6b-gguf + +Exit code is non-zero when any file has an issue, so it can gate a re-export. +""" + +from __future__ import annotations + +import argparse +import glob +import os +import sys +import tempfile +from pathlib import Path + +from gguf import GGUFReader, GGUFValueType + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from lib.gguf_common import BULK_KV_KEYS # noqa: E402 + +REPO_ROOT = Path(__file__).resolve().parent.parent + +# general.* keys we expect every shipped GGUF to carry. CORE is the subset +# whose absence is a hard failure; the rest are reported but advisory. +IDENTITY_KEYS = [ + "general.name", "general.size_label", "general.author", + "general.organization", "general.license", "general.license.name", + "general.license.link", "general.languages", +] +IDENTITY_CORE = {"general.name", "general.size_label", "general.author", + "general.license", "general.languages"} + +# Slug tokens that imply the model is a streaming variant. +STREAMING_HINTS = ("streaming", "unified", "nemotron", "realtime") + + +def field_value(field): + if field is None or not field.types: + return None + t = field.types[0] + if t == GGUFValueType.STRING: + return bytes(field.parts[field.data[0]]).decode("utf-8", "replace") + if t == GGUFValueType.ARRAY: + return "" + return field.parts[field.data[0]][0] + + +def audit_reader(reader: GGUFReader, label: str) -> dict: + keys = list(reader.fields.keys()) + kset = set(keys) + + missing = [k for k in IDENTITY_KEYS if k not in kset] + identity_fail = bool(IDENTITY_CORE - kset) + + present_bulk = [k for k in BULK_KV_KEYS if k in kset] + if present_bulk: + first_bulk = min(keys.index(k) for k in present_bulk) + tail = keys[-len(present_bulk):] + trailer_ok = tail == present_bulk + else: + first_bulk = len(keys) + trailer_ok = True # no tokenizer tables → nothing to trailer + + if "general.file_type" not in kset: + ft_status = "absent" + elif not present_bulk: + ft_status = "no-trailer" + elif keys.index("general.file_type") < first_bulk: + ft_status = "ok" + else: + ft_status = "AFTER-TRAILER" + + ident = (field_value(reader.fields.get("stt.variant")) + or field_value(reader.fields.get("general.basename")) + or label) + looks_streaming = any(h in str(ident).lower() for h in STREAMING_HINTS) + sval = field_value(reader.fields.get("stt.capability.streaming")) + streaming = "ABSENT" if sval is None else bool(sval) + if looks_streaming and streaming is not True: + streaming_suspect = f"slug implies streaming but flag={streaming}" + elif (not looks_streaming) and streaming is True: + streaming_suspect = "flag=true but slug implies offline" + else: + streaming_suspect = None + + issues = [] + if identity_fail: + issues.append("identity") + if not trailer_ok: + issues.append("trailer") + if ft_status == "AFTER-TRAILER": + issues.append("file_type") + if streaming_suspect: + issues.append("streaming") + + return { + "label": label, "missing": missing, "identity_fail": identity_fail, + "trailer_ok": trailer_ok, "file_type": ft_status, + "streaming": streaming, "streaming_suspect": streaming_suspect, + "issues": issues, + } + + +def audit_local(path: Path) -> dict: + return audit_reader(GGUFReader(str(path)), path.name) + + +def audit_hf(repo: str, filename: str, prefix_mb: int) -> dict: + """Range-fetch only the header of an HF-hosted GGUF, then audit it. + + GGUFReader memmaps the whole declared file, so we fetch the first + `prefix_mb` MB (enough for any of our tokenizer tables) and sparse-pad the + temp file to the real size; tensor-data views land in the zero region and + are never read.""" + import requests + from huggingface_hub import get_hf_file_metadata, hf_hub_url + from huggingface_hub.utils import build_hf_headers + + url = hf_hub_url(repo, filename) + total = get_hf_file_metadata(url).size + prefix = min(prefix_mb * 1024 * 1024, total) + headers = build_hf_headers() + headers["Range"] = f"bytes=0-{prefix - 1}" + resp = requests.get(url, headers=headers, timeout=120) + resp.raise_for_status() + + with tempfile.NamedTemporaryFile(suffix=".gguf", delete=False) as tf: + tmp = Path(tf.name) + tf.write(resp.content) + tf.truncate(total) # sparse-pad so memmap views stay in bounds + try: + return audit_reader(GGUFReader(str(tmp)), f"{repo}/{filename}") + finally: + tmp.unlink(missing_ok=True) + + +def collect_local(paths: list[str]) -> list[Path]: + out: list[Path] = [] + for p in paths: + pp = Path(p) + if pp.is_dir(): + out += [Path(x) for x in glob.glob(str(pp / "**" / "*.gguf"), recursive=True)] + elif any(c in p for c in "*?["): + out += [Path(x) for x in glob.glob(p, recursive=True)] + elif pp.is_file(): + out.append(pp) + else: + print(f"warning: no match for {p!r}", file=sys.stderr) + return sorted(set(out)) + + +def print_report(rows: list[dict]) -> None: + w = max((len(r["label"]) for r in rows), default=10) + print(f"\n{'file':<{w}} ident trailer file_type streaming") + print("-" * (w + 42)) + for r in rows: + ident = "MISS" if r["identity_fail"] else ("warn" if r["missing"] else "ok") + trailer = "ok" if r["trailer_ok"] else "BAD" + sflag = r["streaming"] + smark = " <-- SUSPECT" if r["streaming_suspect"] else "" + print(f"{r['label']:<{w}} {ident:<5} {trailer:<7} " + f"{r['file_type']:<13} {str(sflag):<7}{smark}") + # Detail on anything with issues. + flagged = [r for r in rows if r["issues"]] + if flagged: + print("\nNeeds attention:") + for r in flagged: + bits = [] + if r["identity_fail"] or r["missing"]: + bits.append("missing " + ", ".join(k.replace("general.", "") for k in r["missing"])) + if not r["trailer_ok"]: + bits.append("tokenizer arrays not trailered") + if r["file_type"] == "AFTER-TRAILER": + bits.append("file_type stranded after tokenizer trailer") + if r["streaming_suspect"]: + bits.append(r["streaming_suspect"]) + print(f" {r['label']}: {'; '.join(bits)}") + n_issue = len(flagged) + print(f"\n{len(rows)} file(s) audited; {n_issue} with issue(s), " + f"{len(rows) - n_issue} clean.") + + +def main(argv: list[str]) -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("paths", nargs="*", help="GGUF files, dirs, or globs") + p.add_argument("--hf-repo", action="append", default=[], + help="audit a published HF repo's .gguf files (range-fetch headers)") + p.add_argument("--hf-file", action="append", default=[], + help="limit --hf-repo to these filenames (default: all .gguf)") + p.add_argument("--prefix-mb", type=int, default=16, + help="header bytes to range-fetch per HF file (default: 16)") + args = p.parse_args(argv) + + rows: list[dict] = [] + + for repo in args.hf_repo: + from huggingface_hub import HfApi + files = args.hf_file or [f for f in HfApi().list_repo_files(repo) if f.endswith(".gguf")] + for fn in files: + try: + rows.append(audit_hf(repo, fn, args.prefix_mb)) + except Exception as e: # noqa: BLE001 - report and continue the sweep + print(f"ERROR {repo}/{fn}: {e}", file=sys.stderr) + + local_paths = args.paths or ([] if args.hf_repo else [str(REPO_ROOT / "models")]) + for path in collect_local(local_paths): + try: + rows.append(audit_local(path)) + except Exception as e: # noqa: BLE001 + print(f"ERROR {path}: {e}", file=sys.stderr) + + if not rows: + print("no GGUFs found to audit", file=sys.stderr) + return 2 + + print_report(rows) + return 1 if any(r["issues"] for r in rows) else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/convert-granite_nar.py b/scripts/convert-granite_nar.py index 3b355bc0..24b738dc 100644 --- a/scripts/convert-granite_nar.py +++ b/scripts/convert-granite_nar.py @@ -47,7 +47,7 @@ import numpy as np import torch -from gguf import GGUFWriter, GGUFValueType +from gguf import GGUFWriter, GGUFValueType, LlamaFileType from safetensors.torch import safe_open sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -452,6 +452,26 @@ def add_tensor(writer: GGUFWriter, name: str, t: torch.Tensor) -> None: add_bf16(writer, name, t) +def compute_size_label(total_params: int) -> str: + if total_params >= 1_000_000_000: + return f"{total_params / 1_000_000_000:.1f}B" + if total_params >= 1_000_000: + return f"{total_params / 1_000_000:.0f}M" + return f"{total_params / 1_000:.0f}K" + + +def total_safetensors_params(model_dir: Path) -> int: + """Element count summed across all safetensors shards — drives the coarse + general.size_label bucket. get_slice().get_shape() reads only the header, + so no tensor data is materialized.""" + total = 0 + for sf in sorted(model_dir.glob("*.safetensors")): + with safe_open(sf, framework="pt") as h: + for k in h.keys(): + total += int(np.prod(h.get_slice(k).get_shape())) + return total + + # ----- Main ----------------------------------------------------------------- @@ -497,6 +517,9 @@ def main(argv: list[str]) -> int: out_path = outdir / f"{variant}-{REF_DTYPE}.gguf" print(f"Writing GGUF: {out_path}") + size_label = compute_size_label(total_safetensors_params(model_dir)) + print(f" params: ~{size_label}") + writer = gguf_writer(str(out_path), "granite_speech_nar") # ---- general.* ---- @@ -505,6 +528,8 @@ def main(argv: list[str]) -> int: writer, name="Granite Speech 4.1 2B NAR", basename="granite-speech-nar", + size_label=size_label, + file_type=int(LlamaFileType.MOSTLY_BF16), languages=languages, author="IBM", organization="ibm-granite", diff --git a/scripts/convert-medasr.py b/scripts/convert-medasr.py index 91771058..b6565621 100644 --- a/scripts/convert-medasr.py +++ b/scripts/convert-medasr.py @@ -328,6 +328,26 @@ def add_f32_array(writer: GGUFWriter, name: str, arr: np.ndarray) -> None: writer.add_tensor(name, arr) +def compute_size_label(total_params: int) -> str: + if total_params >= 1_000_000_000: + return f"{total_params / 1_000_000_000:.1f}B" + if total_params >= 1_000_000: + return f"{total_params / 1_000_000:.0f}M" + return f"{total_params / 1_000:.0f}K" + + +def total_safetensors_params(model_dir: Path) -> int: + """Element count summed across all safetensors shards — drives the coarse + general.size_label bucket. get_slice().get_shape() reads only the header, + so no tensor data is materialized.""" + total = 0 + for sf in sorted(model_dir.glob("*.safetensors")): + with safe_open(sf, framework="pt") as h: + for k in h.keys(): + total += int(np.prod(h.get_slice(k).get_shape())) + return total + + # ---- Main ----------------------------------------------------------------- @@ -373,6 +393,9 @@ def main(argv: list[str]) -> int: out_path = outdir / f"{slug}-{REF_DTYPE}.gguf" print(f"Writing GGUF: {out_path}") + size_label = compute_size_label(total_safetensors_params(model_dir)) + print(f" params: ~{size_label}") + writer = gguf_writer(str(out_path), ARCH_KEY) # ---- general.* ---- @@ -380,6 +403,7 @@ def main(argv: list[str]) -> int: writer, name="MedASR", basename="medasr", + size_label=size_label, file_type=REFERENCE_FILE_TYPE, languages=["en"], author="Google", diff --git a/scripts/convert-moonshine.py b/scripts/convert-moonshine.py index 164db028..a3051673 100644 --- a/scripts/convert-moonshine.py +++ b/scripts/convert-moonshine.py @@ -477,6 +477,61 @@ def _infer_languages(variant: str) -> list[str]: return [tail] +# Human-friendly pieces for composing general.name from the variant slug. +# UsefulSensors ships moonshine-[-]; this offline converter handles +# the base/tiny sizes plus the 12 language variants +# (moonshine-{tiny,base}-{ar,ja,ko,uk,vi,zh}). Streaming variants are a separate +# converter (convert-moonshine_streaming.py). +_SIZE_DISPLAY = { + "tiny": "Tiny", + "base": "Base", + "small": "Small", + "medium": "Medium", + "large": "Large", +} +_LANGUAGE_DISPLAY = { + "ar": "Arabic", + "ja": "Japanese", + "ko": "Korean", + "uk": "Ukrainian", + "vi": "Vietnamese", + "zh": "Chinese", + "en": "English", +} + + +def _display_name(variant: str) -> str: + """Compose general.name from the variant slug. + + moonshine-tiny → "Moonshine Tiny" + moonshine-base → "Moonshine Base" + moonshine-tiny-vi → "Moonshine Tiny (Vietnamese)" + moonshine-base-zh → "Moonshine Base (Chinese)" + + Raises on an unrecognized size or language suffix so a typo fails the + build loudly rather than shipping a wrong title. + """ + parts = variant.lower().split("-") + if parts[0] != "moonshine" or len(parts) not in (2, 3): + raise ValueError(f"unrecognized moonshine variant slug: {variant!r}") + size = parts[1] + if size not in _SIZE_DISPLAY: + raise ValueError( + f"unknown moonshine size {size!r} in variant {variant!r}; " + f"add it to _SIZE_DISPLAY" + ) + name = f"Moonshine {_SIZE_DISPLAY[size]}" + if len(parts) == 3: + lang = parts[2] + if lang not in _LANGUAGE_DISPLAY: + raise ValueError( + f"unknown moonshine language suffix {lang!r} in variant " + f"{variant!r}; add it to _LANGUAGE_DISPLAY" + ) + name += f" ({_LANGUAGE_DISPLAY[lang]})" + return name + + def convert(model_dir: Path, out_path: Path, variant: str, languages: list[str] | None = None, repo_id: str | None = None) -> None: @@ -540,15 +595,9 @@ def convert(model_dir: Path, out_path: Path, variant: str, writer = gguf_writer(str(out_path), "moonshine") # ---- general.* ---- - _DISPLAY_NAMES = { - "moonshine-tiny": "Moonshine Tiny", - "moonshine-base": "Moonshine Base", - } - if variant not in _DISPLAY_NAMES: - raise ValueError(f"unknown moonshine variant slug: {variant!r}") add_general_identity( writer, - name=_DISPLAY_NAMES[variant], + name=_display_name(variant), basename="moonshine", size_label=size_label, file_type=int(REFERENCE_FILE_TYPE), diff --git a/scripts/convert-parakeet.py b/scripts/convert-parakeet.py index 48831cc0..b2f52cdd 100644 --- a/scripts/convert-parakeet.py +++ b/scripts/convert-parakeet.py @@ -325,19 +325,22 @@ "basename": "parakeet-rnnt", "head_kind": "rnnt", "expected_vocab_size": 13087, - # 40 BCP-47 locales from the model card's transcription-ready - # tier (19) + broad-coverage tier (13) + adaptation-ready tier - # (8). The full prompt_dictionary read from the .nemo carries - # extra aliases (en, en-US, enGB ...) — those are emitted in - # the prompt KVs, not here. + # 32 BCP-47 locales the model card supports for production + # transcription: transcription-ready tier (19) + broad-coverage + # tier (13). The 8 "adaptation-ready" locales (el-GR, lt-LT, + # lv-LV, mt-MT, sl-SI, he-IL, th-TH, nn-NO) are recognized by the + # tokenizer but require fine-tuning for real transcription, so + # they are NOT advertised in general.languages. They remain + # selectable through the prompt dictionary + # (stt.parakeet.prompt.dictionary.*), which still carries all 40 + # locales + aliases (en, en-US, enGB ...) + the auto slot — that + # is the conditioning table, not the transcription-support claim. "languages": [ "en-US", "en-GB", "es-US", "es-ES", "fr-FR", "fr-CA", "it-IT", "pt-BR", "pt-PT", "nl-NL", "de-DE", "tr-TR", "ru-RU", "ar-AR", "hi-IN", "ja-JP", "ko-KR", "vi-VN", "uk-UA", "pl-PL", "sv-SE", "cs-CZ", "nb-NO", "da-DK", "bg-BG", "fi-FI", "hr-HR", "sk-SK", "zh-CN", "hu-HU", "ro-RO", "et-EE", - "el-GR", "lt-LT", "lv-LV", "mt-MT", "sl-SI", "he-IL", "th-TH", - "nn-NO", ], "lang_detect": True, "has_prompt": True, @@ -1389,6 +1392,32 @@ def convert(model_spec: str, out_path: Path, repo_id: str | None = None) -> None if profile["lang_detect"]: writer.add_bool("stt.capability.lang_detect", True) + # Streaming capability. Derived from the encoder attention geometry so + # the header bool can never disagree with what the C++ loader computes: + # this mirrors `Derive supports_streaming from hparams` in + # src/arch/parakeet/model.cpp exactly — + # chunked_limited + (L, R) >= 0 -> cache-aware streaming + # (nemotron-*-streaming) + # chunked_limited_with_rc + non-empty (L, C, R) -> buffered streaming + # (parakeet-unified-en-0.6b) + # regular / anything else -> offline only + # Emitted unconditionally (true AND false) so a header-only reader gets the + # correct answer without replaying this derivation. The loader still treats + # the architecture as the floor and ignores a contradictory KV. + cache_aware_streaming = ( + hp["enc_att_context_style"] == "chunked_limited" + and hp["enc_att_context_left"] >= 0 + and hp["enc_att_context_right"] >= 0 + ) + buffered_streaming = ( + hp["enc_att_context_style"] == "chunked_limited_with_rc" + and bool(hp.get("enc_att_chunk_left_choices")) + and bool(hp.get("enc_att_chunk_chunk_choices")) + and bool(hp.get("enc_att_chunk_right_choices")) + ) + writer.add_bool("stt.capability.streaming", + cache_aware_streaming or buffered_streaming) + # Head discriminator. The C++ loader currently always reads TDT # KV (durations etc.); Stage 4 will gate predictor/joint/tdt # reads on this. Emit unconditionally so the GGUF carries the diff --git a/scripts/convert-sensevoice.py b/scripts/convert-sensevoice.py index 58fbca1f..0ce2d2d3 100644 --- a/scripts/convert-sensevoice.py +++ b/scripts/convert-sensevoice.py @@ -449,6 +449,7 @@ def convert(model_dir: Path, out_path: Path, variant: str, repo_id: str | None = # consumers see the real license, and retain the MODEL_LICENSE link # the agreement's attribution clause (2.2) requires. license="FunASR-Model-License-1.1", + license_name="FunASR Model Open Source License Agreement v1.1", license_link="https://github.com/modelscope/FunASR/blob/main/MODEL_LICENSE", repo_url=(f"https://huggingface.co/{repo_id}" if repo_id else None), url="https://huggingface.co/FunAudioLLM/SenseVoiceSmall", diff --git a/scripts/hf_cards/nemotron-3.5-asr-streaming-0.6b.yaml b/scripts/hf_cards/nemotron-3.5-asr-streaming-0.6b.yaml index 2fc1c7b8..a98ae4a1 100644 --- a/scripts/hf_cards/nemotron-3.5-asr-streaming-0.6b.yaml +++ b/scripts/hf_cards/nemotron-3.5-asr-streaming-0.6b.yaml @@ -52,14 +52,6 @@ languages: - hu - ro - et - - el - - lt - - lv - - mt - - sl - - he - - th - - nn tags: - gguf - transcribe.cpp @@ -73,7 +65,7 @@ tags: - multilingual summary: | - Multilingual speech-to-text across 40 language-locales with punctuation and capitalization. A 0.6B-parameter cache-aware streaming FastConformer encoder with a prompt-conditioned RNN-T transducer decoder; the target language is selected per call (--language en-US, fr-FR, de-DE, ...) and an auto mode emits a tag. Ships both the offline path (att_context_size=[56, 13], 1.12s, headline accuracy) and runtime-selectable chunked streaming (--stream-chunk-ms 1120 --stream-att-right {0,3,6,13}). + Multilingual speech-to-text across 32 supported language-locales (the model's tokenizer recognizes 40, but 8 are adaptation-ready and need fine-tuning) with punctuation and capitalization. A 0.6B-parameter cache-aware streaming FastConformer encoder with a prompt-conditioned RNN-T transducer decoder; the target language is selected per call (--language en-US, fr-FR, de-DE, ...) and an auto mode emits a tag. Ships both the offline path (att_context_size=[56, 13], 1.12s, headline accuracy) and runtime-selectable chunked streaming (--stream-chunk-ms 1120 --stream-att-right {0,3,6,13}). default_quant_index: 2 # Q8_0 diff --git a/scripts/lib/test_quantize_bulk_trailer.py b/scripts/lib/test_quantize_bulk_trailer.py new file mode 100644 index 00000000..841747cc --- /dev/null +++ b/scripts/lib/test_quantize_bulk_trailer.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Regression test: tools/transcribe-quantize keeps general.file_type (and all +small scalar KVs) ahead of the multi-MB tokenizer trailer. + +The C++ quantizer copies every input KV (gguf_set_kv) then overrides +general.file_type. Because every ggml gguf setter is remove-then-append, that +override lands file_type at the very end — behind the tokenizer tables — unless +the quantizer re-asserts the bulk-KV trailer (move_bulk_kv_last). This builds a +reference GGUF with the converter's real trailer layout (via gguf_writer()), +runs the built binary, and asserts the fix holds and the bulk values survive. + +Exit-code driven: uv run --project scripts/envs/moonshine \ + scripts/lib/test_quantize_bulk_trailer.py [path-to-binary] +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +from gguf import GGUFReader, GGMLQuantizationType, LlamaFileType +from gguf.quants import dequantize + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from lib.gguf_common import ( # noqa: E402 + BULK_KV_KEYS, + TOKEN_TYPE_CONTROL, + TOKEN_TYPE_NORMAL, + add_general_identity, + gguf_writer, +) + +# Small fake vocab — exercises every bulk key type: tokens/merges (string +# arrays), token_type (int32 array), scores (float32 array). +TOKENS = ["", "", ""] + [f"tok{i}" for i in range(29)] +TYPES = [TOKEN_TYPE_CONTROL, TOKEN_TYPE_CONTROL, TOKEN_TYPE_CONTROL] + \ + [TOKEN_TYPE_NORMAL] * 29 +SCORES = [float(-i) for i in range(len(TOKENS))] +MERGES = ["t o", "to k", "tok 0"] + +# One quantizable linear (n_per_row=64, a clean multiple of the Q8_0 block) plus +# a norm + bias that stay F32 — mirrors the real per-tensor policy buckets. +FC1 = np.arange(8 * 64, dtype=np.float32).reshape(8, 64) * 0.01 - 2.0 +NORM = np.linspace(0.5, 1.5, 64, dtype=np.float32) +BIAS = np.linspace(-0.1, 0.1, 64, dtype=np.float32) + + +def build_reference(path: Path) -> None: + w = gguf_writer(str(path), "moonshine") + add_general_identity( + w, + name="Synthetic Moonshine", + basename="moonshine", + size_label="0M", + file_type=int(LlamaFileType.ALL_F32), + languages=["en"], + ) + w.add_string("stt.variant", "synthetic") + w.add_bool("stt.capability.timestamps", False) + + w.add_string("tokenizer.ggml.model", "bpe") + w.add_string("tokenizer.ggml.pre", "default") + w.add_array("tokenizer.ggml.tokens", TOKENS) + w.add_array("tokenizer.ggml.token_type", TYPES) + w.add_array("tokenizer.ggml.scores", SCORES) + w.add_array("tokenizer.ggml.merges", MERGES) + w.add_uint32("tokenizer.ggml.bos_token_id", 1) + + w.add_tensor("dec.blocks.0.ffn.fc1.weight", FC1, raw_dtype=GGMLQuantizationType.F32) + w.add_tensor("dec.blocks.0.norm_self.weight", NORM, raw_dtype=GGMLQuantizationType.F32) + w.add_tensor("dec.blocks.0.ffn.fc1.bias", BIAS, raw_dtype=GGMLQuantizationType.F32) + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + + +def kv_keys(path: Path) -> list[str]: + return list(GGUFReader(str(path)).fields.keys()) + + +def check(cond: bool, msg: str) -> None: + if not cond: + print(f"FAIL {msg}") + raise SystemExit(1) + print(f"ok {msg}") + + +def main(argv: list[str]) -> int: + binary = Path(argv[1]) if len(argv) > 1 else \ + Path(__file__).resolve().parents[2] / "bin" / "transcribe-quantize" + if not binary.is_file(): + print(f"FAIL quantizer binary not found: {binary}") + return 1 + + # Honour TRANSCRIBE_TMPDIR; otherwise reuse the sandbox tmp/ next to the + # repo when present, else the system default. Keeps local runs inside the + # sandbox without pinning a machine-specific path into the test. + tmp_root = os.environ.get("TRANSCRIBE_TMPDIR") + if not tmp_root: + sandbox_tmp = Path(__file__).resolve().parents[3] / "tmp" + tmp_root = str(sandbox_tmp) if sandbox_tmp.is_dir() else None + + with tempfile.TemporaryDirectory(dir=tmp_root) as td: + ref = Path(td) / "ref-F32.gguf" + out = Path(td) / "out-Q8_0.gguf" + build_reference(ref) + + # Precondition: the reference really does carry the trailer layout. + rk = kv_keys(ref) + bulk_in_ref = [k for k in BULK_KV_KEYS if k in rk] + first_bulk = min(rk.index(k) for k in bulk_in_ref) + check(rk.index("general.file_type") < first_bulk, + "reference: file_type precedes the tokenizer trailer") + check(rk.index("tokenizer.ggml.tokens") >= len(rk) - len(bulk_in_ref), + "reference: bulk keys are the trailer") + + r = subprocess.run( + [str(binary), str(ref), str(out), "--quant", "Q8_0"], + capture_output=True, text=True, + ) + print(r.stdout.strip()) + check(r.returncode == 0, f"quantizer exit 0 (got {r.returncode}); stderr={r.stderr.strip()}") + check(out.is_file(), "quantizer produced an output file") + + # --- The fix: file_type stays ahead of the tokenizer trailer --- + ok = kv_keys(out) + check("general.file_type" in ok, "output carries general.file_type") + present_bulk = [k for k in BULK_KV_KEYS if k in ok] + first_bulk_out = min(ok.index(k) for k in present_bulk) + check(ok.index("general.file_type") < first_bulk_out, + "output: file_type precedes the tokenizer trailer (THE FIX)") + + # Bulk keys are the final KVs, in canonical BULK_KV_KEYS order. + tail = ok[-len(present_bulk):] + check(tail == present_bulk, + f"output: bulk keys are last, in canonical order (got {tail})") + + # --- Values survived the move-then-reappend round trip --- + rd = GGUFReader(str(out)) + + def arr_str(key): + f = rd.fields[key] + return [bytes(f.parts[i]).decode("utf-8") for i in f.data] + + def arr_num(key): + f = rd.fields[key] + return [f.parts[i][0] for i in f.data] + + check(arr_str("tokenizer.ggml.tokens") == TOKENS, "tokens intact") + check(arr_str("tokenizer.ggml.merges") == MERGES, "merges intact") + check([int(x) for x in arr_num("tokenizer.ggml.token_type")] == TYPES, + "token_type intact") + got_scores = [float(x) for x in arr_num("tokenizer.ggml.scores")] + check(np.allclose(got_scores, SCORES), "scores intact") + + ft = int(rd.fields["general.file_type"].parts[rd.fields["general.file_type"].data[0]][0]) + check(ft == int(LlamaFileType.MOSTLY_Q8_0), + f"file_type updated to MOSTLY_Q8_0 ({int(LlamaFileType.MOSTLY_Q8_0)}), got {ft}") + + # --- Tensor actually requantized and round-trips within Q8_0 error --- + fc1 = next(t for t in rd.tensors if t.name == "dec.blocks.0.ffn.fc1.weight") + check(fc1.tensor_type == GGMLQuantizationType.Q8_0, + f"fc1 is Q8_0 (got {fc1.tensor_type})") + # GGUFReader hands back the raw Q8_0 block bytes; dequantize to compare. + deq = dequantize(fc1.data, GGMLQuantizationType.Q8_0).reshape(FC1.shape) + check(np.max(np.abs(deq - FC1)) < 0.05, + "fc1 dequantizes back within Q8_0 tolerance") + + print("\nALL PASS") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/tools/transcribe-quantize/main.cpp b/tools/transcribe-quantize/main.cpp index 169f9476..f8a63da3 100644 --- a/tools/transcribe-quantize/main.cpp +++ b/tools/transcribe-quantize/main.cpp @@ -115,6 +115,89 @@ void print_stats(const std::map & stats, label); } +// --------------------------------------------------------------------------- +// Bulk-KV trailer +// --------------------------------------------------------------------------- + +// Large tokenizer KVs, in canonical trailer order. Mirrors BULK_KV_KEYS in +// scripts/lib/gguf_common.py: the Python converter relocates these to the end +// of the KV section so range-read consumers can fetch the small scalar +// metadata (general.*, stt.*, tokenizer identity) without pulling the multi-MB +// tokenizer tables. We re-assert that layout here because the quantizer's +// general.file_type override appends via ggml's remove-then-append setter, +// which would otherwise strand file_type — a tiny scalar a remote consumer +// wants — *after* the tokenizer trailer. +const char * const kBulkKvKeys[] = { + "tokenizer.ggml.tokens", + "tokenizer.ggml.scores", + "tokenizer.ggml.token_type", + "tokenizer.ggml.merges", + "tokenizer.chat_template", +}; + +// Byte size of one element for the scalar element types that can appear inside +// a bulk array (strings are handled separately). gguf_type_size is not part of +// the public gguf.h, so we size the handful of types we actually move; 0 means +// "unexpected type, leave the key where it is." +size_t bulk_kv_elem_size(gguf_type t) { + switch (t) { + case GGUF_TYPE_UINT8: case GGUF_TYPE_INT8: case GGUF_TYPE_BOOL: return 1; + case GGUF_TYPE_UINT16: case GGUF_TYPE_INT16: return 2; + case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: case GGUF_TYPE_FLOAT32: return 4; + case GGUF_TYPE_UINT64: case GGUF_TYPE_INT64: case GGUF_TYPE_FLOAT64: return 8; + default: return 0; + } +} + +// Move the bulk tokenizer KVs to the very end of ctx's KV section, preserving +// each value. Every ggml gguf setter is remove-then-append, so reading a key's +// value into a local copy and re-setting it relocates it to the tail. Keys are +// processed in kBulkKvKeys order, so the trailer ends up canonical regardless +// of how they were interleaved on the way in. Absent keys are skipped. Mirrors +// move_bulk_metadata_last() in scripts/lib/gguf_common.py. +void move_bulk_kv_last(gguf_context * ctx) { + for (const char * key : kBulkKvKeys) { + const int64_t kid = gguf_find_key(ctx, key); + if (kid < 0) { + continue; // not present — nothing to move + } + const gguf_type kv_type = gguf_get_kv_type(ctx, kid); + if (kv_type == GGUF_TYPE_ARRAY) { + const gguf_type elem = gguf_get_arr_type(ctx, kid); + const size_t n = gguf_get_arr_n(ctx, kid); + if (elem == GGUF_TYPE_STRING) { + // Copy the strings out before the remove-then-append re-set + // frees the originals. + std::vector store(n); + std::vector ptrs(n); + for (size_t i = 0; i < n; ++i) { + store[i] = gguf_get_arr_str(ctx, kid, i); + } + for (size_t i = 0; i < n; ++i) { + ptrs[i] = store[i].c_str(); + } + gguf_set_arr_str(ctx, key, ptrs.data(), n); + } else { + const size_t esz = bulk_kv_elem_size(elem); + if (esz == 0) { + continue; // unexpected element type — leave it in place + } + // Copy the raw element bytes out before the re-set frees them. + const void * data = gguf_get_arr_data(ctx, kid); + std::vector tmp(n * esz); + if (!tmp.empty()) { + std::memcpy(tmp.data(), data, n * esz); + } + gguf_set_arr_data(ctx, key, elem, tmp.data(), n); + } + } else if (kv_type == GGUF_TYPE_STRING) { + const std::string val = gguf_get_val_str(ctx, kid); + gguf_set_val_str(ctx, key, val.c_str()); + } + // Other scalar types are not expected among the bulk keys; ignore. + } +} + // --------------------------------------------------------------------------- // Main // --------------------------------------------------------------------------- @@ -298,6 +381,14 @@ int main(int argc, char ** argv) { // gguf-dump display it. gguf_set_val_u32(gguf_out, "general.file_type", preset->file_type); + // Re-assert the bulk-KV trailer so general.file_type (and every other + // small scalar) precedes the multi-MB tokenizer tables. The override + // above appends file_type last via ggml's remove-then-append setter; + // without this, a range-read consumer couldn't reach file_type without + // parsing the whole tokenizer blob. Mirrors the converter's gguf_writer() + // layout (scripts/lib/gguf_common.py). + move_bulk_kv_last(gguf_out); + // ---- Per-tensor: dequant → fp32 → requant → add to ctx_out ---- std::vector fp32_scratch; int64_t requantized = 0;