|
4 | 4 | # (scripts/convert-nemo-to-gguf.py). |
5 | 5 | # qvac-parakeet.cpp owns the GGUF tensor/metadata layout this script |
6 | 6 | # produces; if the layout changes upstream, resync the file rather than |
7 | | -# diverging it locally. Copy is verbatim apart from this header. |
| 7 | +# diverging it locally. Copy is verbatim apart from this header, PLUS the |
| 8 | +# rnnt head support, which is ahead of the pinned commit: it mirrors the |
| 9 | +# engine-side change hunk-for-hunk (see the companion parakeet-cpp PR). |
| 10 | +# Re-pin and resync once that merges. |
8 | 11 | # Original copyright/license: MIT, qvac-parakeet.cpp authors -- see |
9 | 12 | # the top-level NOTICE for attribution. |
10 | 13 | """Convert an NVIDIA NeMo .nemo archive to a single GGUF for the |
|
23 | 26 | cache-aware streaming, end-of- |
24 | 27 | utterance token detection; |
25 | 28 | parakeet_realtime_eou_120m-v1) |
| 29 | + - ``EncDecHybridRNNTCTCBPEModel`` / ``EncDecRNNTBPEModel`` |
| 30 | + (no TDT durations, no ``<EOU>`` token) |
| 31 | + -> plain RNN-T (Transducer) head; |
| 32 | + for hybrids only the Transducer |
| 33 | + branch is exported (the CTC aux |
| 34 | + head ``ctc_decoder.*`` is ignored) |
26 | 35 | - ``EncDecDiarLabelModel`` -> Sortformer (diar_sortformer_4spk-v1, |
27 | 36 | diar_streaming_sortformer_4spk-v2) |
28 | 37 |
|
29 | | -The FastConformer encoder topology is shared across all four flavours; only |
| 38 | +The FastConformer encoder topology is shared across all five flavours; only |
30 | 39 | the decoder / head tensors + metadata differ. EOU additionally swaps the |
31 | 40 | conv module's BatchNorm for a LayerNorm and carries cache-aware streaming |
32 | 41 | hyperparameters (att_context_size, subsampling-output cache lookback, and the |
|
43 | 52 | Metadata: |
44 | 53 | general.architecture = "parakeet-ctc" (kept for GGUF compat) |
45 | 54 | general.name = "<derived from cfg>" |
46 | | - parakeet.model.type = "ctc", "tdt", "eou", or "sortformer" |
| 55 | + parakeet.model.type = "ctc", "tdt", "rnnt", "eou", or "sortformer" |
47 | 56 | parakeet.encoder.* (hyperparameters, incl. use_bias, xscaling, |
48 | 57 | conv_norm_type, att_context_size, |
49 | 58 | causal_downsampling, conv_context_size) |
50 | 59 | parakeet.preproc.* (mel/stft hyperparameters) |
51 | 60 | parakeet.ctc.* (vocab_size, blank_id) [CTC only] |
52 | 61 | parakeet.tdt.* (predictor + joint hyperparameters |
53 | 62 | + durations) [TDT only] |
| 63 | + parakeet.rnnt.* (vocab_size, blank_id, pred_hidden, |
| 64 | + pred_rnn_layers, joint_hidden, |
| 65 | + max_symbols_per_step) [RNNT only] |
54 | 66 | parakeet.eou.* (vocab_size, blank_id, eou_id, eob_id, |
55 | 67 | pred_hidden, pred_rnn_layers, joint_hidden, |
56 | 68 | encoder_chunk_mel_frames, |
|
73 | 85 | tdt.predict.lstm.{l}.{w_ih,w_hh,b_ih,b_hh} [TDT only] |
74 | 86 | tdt.joint.{enc,pred}.{weight,bias} [TDT only] |
75 | 87 | tdt.joint.out.{weight,bias} [TDT only] |
| 88 | + rnnt.predict.embed.weight [RNNT only] |
| 89 | + rnnt.predict.lstm.{l}.{w_ih,w_hh,b_ih,b_hh} [RNNT only] |
| 90 | + rnnt.joint.{enc,pred}.{weight,bias} [RNNT only] |
| 91 | + rnnt.joint.out.{weight,bias} [RNNT only] |
76 | 92 | eou.predict.embed.weight [EOU only] |
77 | 93 | eou.predict.lstm.0.{w_ih,w_hh,b_ih,b_hh} [EOU only] |
78 | 94 | eou.joint.{enc,pred}.{weight,bias} [EOU only] |
@@ -393,15 +409,31 @@ def write_gguf(out: Path, ckpt: Path, cfg: dict, sd: dict, tok_bytes: bytes, qua |
393 | 409 | # identical to TDT/EOU; the hybrid's CTC aux head (ctc_decoder.*) is |
394 | 410 | # ignored. Cache-aware streaming params (att_context_size, conv_norm_type) |
395 | 411 | # ride in the shared parakeet.encoder.* metadata already emitted above. |
| 412 | + if "prednet" not in dec or "joint" not in cfg: |
| 413 | + raise RuntimeError( |
| 414 | + "rnnt head requires a Transducer checkpoint " |
| 415 | + "(decoder.prednet / joint config missing -- is this a CTC model?)") |
396 | 416 | pred_hidden = int(dec["prednet"]["pred_hidden"]) |
397 | 417 | pred_rnn_layers = int(dec["prednet"]["pred_rnn_layers"]) |
398 | 418 | joint_hidden = int(cfg["joint"]["jointnet"]["joint_hidden"]) |
399 | 419 | pred_vocab_size = int(dec["vocab_size"]) # label vocab (no blank) |
400 | | - joint_num_classes = int(cfg["joint"]["num_classes"]) # label vocab + blank |
401 | | - blank_id = joint_num_classes # blank_as_pad at vocab_size |
| 420 | + joint_num_classes = int(cfg["joint"]["num_classes"]) # label vocab, blank excluded (RNNTJoint adds +1) |
| 421 | + blank_id = joint_num_classes # blank_as_pad at index vocab_size |
| 422 | + if joint_num_classes != pred_vocab_size: |
| 423 | + raise RuntimeError( |
| 424 | + f"rnnt: joint.num_classes ({joint_num_classes}) != decoder.vocab_size " |
| 425 | + f"({pred_vocab_size}); blank_id placement would be wrong") |
402 | 426 | greedy_cfg = (cfg.get("decoding") or {}).get("greedy") or {} |
403 | | - max_symbols = int(greedy_cfg.get("max_symbols") |
404 | | - or greedy_cfg.get("max_symbols_per_step") or 10) |
| 427 | + max_symbols_cfg = greedy_cfg.get("max_symbols", |
| 428 | + greedy_cfg.get("max_symbols_per_step")) |
| 429 | + if not max_symbols_cfg: |
| 430 | + # NeMo treats an unset/None max_symbols as "no cap"; the C++ greedy |
| 431 | + # loop needs a finite cap, so bake in NeMo's usual default of 10. |
| 432 | + print("[convert] rnnt: decoding.greedy.max_symbols unset (uncapped in " |
| 433 | + "NeMo); capping at 10 in the GGUF", file=sys.stderr) |
| 434 | + max_symbols = 10 |
| 435 | + else: |
| 436 | + max_symbols = int(max_symbols_cfg) |
405 | 437 |
|
406 | 438 | writer.add_uint32("parakeet.rnnt.vocab_size", pred_vocab_size) |
407 | 439 | writer.add_uint32("parakeet.rnnt.blank_id", blank_id) |
@@ -656,6 +688,16 @@ def try_bias(name: str, key: str): |
656 | 688 | add_f32("rnnt.joint.enc.bias", sd["joint.enc.bias"]) |
657 | 689 | add_2d ("rnnt.joint.pred.weight", sd["joint.pred.weight"]) |
658 | 690 | add_f32("rnnt.joint.pred.bias", sd["joint.pred.bias"]) |
| 691 | + # A plain RNN-T joint emits exactly vocab+1 logits (labels + blank). |
| 692 | + # A TDT-shaped joint (vocab+1+num_durations rows) reaching this branch |
| 693 | + # (e.g. via --type rnnt) must fail here, not decode garbage later. |
| 694 | + out_rows = int(sd["joint.joint_net.2.weight"].shape[0]) |
| 695 | + vocab_p1 = int(cfg["decoder"]["vocab_size"]) + 1 |
| 696 | + if out_rows != vocab_p1: |
| 697 | + raise RuntimeError( |
| 698 | + f"rnnt: joint output has {out_rows} rows, expected vocab+1 = " |
| 699 | + f"{vocab_p1}; duration logits present? (TDT checkpoint -- use " |
| 700 | + f"the auto-detected head instead of forcing rnnt)") |
659 | 701 | add_2d ("rnnt.joint.out.weight", sd["joint.joint_net.2.weight"]) |
660 | 702 | add_f32("rnnt.joint.out.bias", sd["joint.joint_net.2.bias"]) |
661 | 703 | else: |
|
0 commit comments