Add custom vocabulary (phrase boosting)

transcribe(..., boost=[...]) adds phrase boosting: our implementation of
NeMo's phrase boosting (phrase_boost.py), using the model's tokenizer.

Claude-Session: https://claude.ai/code/session_01JeoiSkEdCzw16SqBenVs8w

README.md

@@ -39,7 +39,7 @@ For other computers, use the [ONNX versions](https://huggingface.co/KlangAI/pian
 ## Usage
 
 ```bash
-pip install parakeet-mlx==0.5.2 huggingface_hub
+pip install parakeet-mlx==0.5.2 huggingface_hub sentencepiece
 brew install ffmpeg   # to read audio files
 ```
 
@@ -60,6 +60,16 @@ Audio longer than 2 minutes is transcribed in 2-minute chunks with 15 s overlap.
 
 Use `pianissimo_mlx.py` from this repository rather than parakeet-mlx's own `transcribe`. It can load the 8-bit and 4-bit versions, which parakeet-mlx cannot do on its own, and it prepares the audio exactly as it was prepared during training. parakeet-mlx prepares it slightly differently, which adds about 0.2 percentage points of WER.
 
+### Vocabulary
+
+Supports custom vocabulary using our custom implementation of NeMo's phrase boosting. The implementation matches NeMo's results on the original model. Listing names and terms increases the chance that the model spells them correctly.
+
+```python
+result = pianissimo_mlx.transcribe(model, "audio.wav", boost=["Klang AI", "Pianissimo", "RixVox", "Common Voice"])
+```
+
+Test results are on the [main model card](https://huggingface.co/KlangAI/pianissimo-sv#vocabulary).
+
 ## How they were made
 
 All versions are converted from the released checkpoint and keep its local attention (256 frames on each side). In the 8-bit and 4-bit versions, the encoder's linear layers are quantized with MLX's built-in quantization, in groups of 64 weights for 8-bit and 32 for 4-bit. Everything else is stored in bfloat16.

phrase_boost.py

@@ -0,0 +1,84 @@
+"""Phrase boosting for Klang Pianissimo. A numpy port of NeMo's GPU-PB boosting tree
+(ContextGraph and GPUBoostingTreeModel, NeMo 2.7, Apache 2.0)."""
+import math
+from collections import deque
+
+import numpy as np
+import sentencepiece as spm
+
+
+class _Node:
+    __slots__ = ("id", "children", "token_score", "node_score", "is_end", "fail")
+
+    def __init__(self, node_id, token_score, node_score, is_end):
+        self.id = node_id
+        self.children = {}
+        self.token_score = token_score
+        self.node_score = node_score
+        self.is_end = is_end
+        self.fail = None
+
+
+class PhraseBoost:
+    def __init__(self, phrases, tokenizer_path, vocab_size, context_score=1.0, depth_scaling=2.0):
+        sp = spm.SentencePieceProcessor(model_file=str(tokenizer_path))
+        self.vocab_size = vocab_size
+        self.root = _Node(0, 0.0, 0.0, False)
+        self.root.fail = self.root
+        self.nodes = [self.root]
+        for phrase in dict.fromkeys(phrases):
+            self._add(sp.encode(phrase), context_score, depth_scaling)
+        self._fill_fail()
+        self._cache = {}
+
+    def _add(self, tokens, context_score, depth_scaling):
+        node = self.root
+        for i, token in enumerate(tokens):
+            last = i == len(tokens) - 1
+            child = node.children.get(token)
+            if child is None:
+                token_score = context_score * depth_scaling + math.log(i + 1) if i > 0 else context_score
+                child = _Node(len(self.nodes), token_score, node.node_score + token_score, last)
+                node.children[token] = child
+                self.nodes.append(child)
+            else:
+                child.token_score = max(context_score, child.token_score)
+                child.node_score = node.node_score + child.token_score
+                child.is_end = last or child.is_end
+            node = child
+
+    def _fill_fail(self):
+        queue = deque()
+        for child in self.root.children.values():
+            child.fail = self.root
+            queue.append(child)
+        while queue:
+            node = queue.popleft()
+            for token, child in node.children.items():
+                fail = node.fail
+                while token not in fail.children and fail is not self.root:
+                    fail = fail.fail
+                child.fail = fail.children.get(token, self.root)
+                queue.append(child)
+
+    def scores(self, state):
+        """Bonus and next state for every token from `state`."""
+        if state not in self._cache:
+            bonus = np.zeros(self.vocab_size, dtype=np.float32)
+            nxt = np.full(self.vocab_size, -1, dtype=np.int64)
+            node, taken_back = self.nodes[state], 0.0
+            while True:
+                for token, child in node.children.items():
+                    if nxt[token] < 0:
+                        bonus[token] = taken_back + child.token_score
+                        nxt[token] = child.id
+                if node is self.root:
+                    unset = nxt < 0
+                    bonus[unset] = taken_back
+                    nxt[unset] = 0
+                    break
+                if not node.is_end:
+                    taken_back += node.fail.node_score - node.node_score
+                node = node.fail
+            self._cache[state] = (bonus, nxt)
+        return self._cache[state]

pianissimo_mlx.py

@@ -1,13 +1,8 @@
-"""Load Klang Pianissimo for MLX and transcribe Swedish speech on a Mac with Apple silicon.
+"""Klang Pianissimo on Apple silicon with parakeet-mlx: loads the quantized versions, computes NeMo's
+audio features, splits long audio into chunks and supports phrase boosting.
 
-The model itself runs in parakeet-mlx. This file loads the 8-bit and 4-bit versions, prepares the
-audio exactly as it was prepared during training (parakeet-mlx does it slightly differently, which
-adds about 0.2 percentage points of WER), and splits long recordings into overlapping pieces.
-
-    import pianissimo_mlx
     model = pianissimo_mlx.load("KlangAI/pianissimo-sv-mlx-8bit")
-    result = pianissimo_mlx.transcribe(model, "audio.wav")
-    print(result.text)
+    print(pianissimo_mlx.transcribe(model, "audio.wav", boost=["Klang AI", "Pianissimo", "RixVox", "Common Voice"]).text)
 """
 import json
 import os
@@ -17,10 +12,13 @@ import subprocess
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
-from parakeet_mlx.alignment import (merge_longest_common_subsequence, merge_longest_contiguous,
+from parakeet_mlx import tokenizer
+from parakeet_mlx.alignment import (AlignedToken, merge_longest_common_subsequence, merge_longest_contiguous,
                                     sentences_to_result, tokens_to_sentences)
 from parakeet_mlx.utils import from_config
 
+from phrase_boost import PhraseBoost
+
 SAMPLE_RATE = 16000
 
 
@@ -29,7 +27,7 @@ def load(repo_or_path):
     path = repo_or_path
     if not os.path.isdir(path):
         from huggingface_hub import snapshot_download
-        path = snapshot_download(repo_or_path, allow_patterns=["config.json", "model.safetensors"])
+        path = snapshot_download(repo_or_path, allow_patterns=["config.json", "model.safetensors", "tokenizer.model"])
     cfg = json.load(open(os.path.join(path, "config.json")))
     model = from_config(cfg)
     q = cfg.get("quantization")
@@ -37,6 +35,7 @@ def load(repo_or_path):
         nn.quantize(model, group_size=q["group_size"], bits=q["bits"],
                     class_predicate=lambda p, m: isinstance(m, nn.Linear) and p.startswith("encoder.layers."))
     model.load_weights(os.path.join(path, "model.safetensors"))
+    model.tokenizer_path = os.path.join(path, "tokenizer.model")
     return model
 
 
@@ -50,7 +49,7 @@ def read_audio(path):
 
 
 def features(audio, args):
-    """The 128-band log spectrogram the model was trained on, computed the same way as in NeMo."""
+    """NeMo's log-mel featurizer, ported to MLX."""
     n_fft, win = args.n_fft, int(args.window_size * args.sample_rate)
     hop = int(args.window_stride * args.sample_rate)
     x = mx.concatenate([audio[:1], audio[1:] - 0.97 * audio[:-1]])
@@ -66,24 +65,21 @@ def features(audio, args):
     return mx.expand_dims((logmel - mean) / (std + 1e-5), 0)
 
 
-def transcribe(model, audio, chunk_duration=120.0, overlap_duration=15.0):
-    """audio is a file path or a 16 kHz mono numpy array. The result has .text, and .sentences
-    with start and end times per sentence and per token.
-
-    Audio longer than chunk_duration seconds is transcribed in overlapping pieces and joined where
-    they overlap, which keeps memory use flat however long the recording is. Pass
-    chunk_duration=None to transcribe the whole recording in one pass."""
+def transcribe(model, audio, chunk_duration=120.0, overlap_duration=15.0, boost=None, boost_alpha=1.0):
+    """audio is a path or a 16 kHz mono numpy array. Audio longer than chunk_duration seconds is
+    transcribed in overlapping chunks (None for one pass). boost lists phrases to spell as written."""
     if isinstance(audio, (str, os.PathLike)):
         audio = read_audio(audio)
     audio = np.asarray(audio, dtype=np.float32)
+    tree = PhraseBoost(boost, model.tokenizer_path, vocab_size=len(model.vocabulary)) if boost else None
     if chunk_duration is None or len(audio) <= chunk_duration * SAMPLE_RATE:
-        return _transcribe_piece(model, audio)
+        return _transcribe_piece(model, audio, tree, boost_alpha)
 
     step = int((chunk_duration - overlap_duration) * SAMPLE_RATE)
     size = int(chunk_duration * SAMPLE_RATE)
     tokens = []
     for start in range(0, len(audio), step):
-        result = _transcribe_piece(model, audio[start:start + size])
+        result = _transcribe_piece(model, audio[start:start + size], tree, boost_alpha)
         offset = start / SAMPLE_RATE
         for sentence in result.sentences:
             for token in sentence.tokens:
@@ -92,7 +88,6 @@ def transcribe(model, audio, chunk_duration=120.0, overlap_duration=15.0):
         if not tokens:
             tokens = result.tokens
         else:
-            # the same joining parakeet-mlx uses for its own chunked transcription
             try:
                 tokens = merge_longest_contiguous(tokens, result.tokens, overlap_duration=overlap_duration)
             except RuntimeError:
@@ -103,5 +98,50 @@ def transcribe(model, audio, chunk_duration=120.0, overlap_duration=15.0):
     return sentences_to_result(tokens_to_sentences(tokens))
 
 
-def _transcribe_piece(model, audio):
-    return model.generate(features(mx.array(audio), model.preprocessor_config))[0]
+def _transcribe_piece(model, audio, tree, alpha):
+    mel = features(mx.array(audio), model.preprocessor_config)
+    if tree is None:
+        return model.generate(mel)[0]
+    encoded, lengths = model.encoder(mel)
+    mx.eval(encoded, lengths)
+    return sentences_to_result(tokens_to_sentences(_boosted_greedy(model, encoded, int(lengths[0]), tree, alpha)))
+
+
+def _boosted_greedy(model, encoded, length, tree, alpha):
+    """parakeet-mlx's greedy TDT decoding with the phrase bonus added."""
+    blank = len(model.vocabulary)
+    hypothesis, last_token, hidden_state = [], None, None
+    boost_state, step, new_symbols = 0, 0, 0
+    while step < length:
+        decoder_out, (hidden, cell) = model.decoder(
+            mx.array([[last_token]]) if last_token is not None else None, hidden_state)
+        decoder_out = decoder_out.astype(encoded.dtype)
+        decoder_hidden = (hidden.astype(encoded.dtype), cell.astype(encoded.dtype))
+        joint_out = model.joint(encoded[:, step:step + 1], decoder_out)
+
+        token_logits = joint_out[0, 0, :, :blank + 1]
+        pred_token = int(mx.argmax(token_logits))
+        token_probs = mx.softmax(token_logits, axis=-1)
+        entropy = -mx.sum(token_probs * mx.log(token_probs + 1e-10), axis=-1)
+        confidence = float(1.0 - entropy / mx.log(mx.array(blank + 1, dtype=token_probs.dtype)))
+        decision = int(mx.argmax(joint_out[0, 0, :, blank + 1:]))
+
+        if pred_token != blank:
+            bonus, next_states = tree.scores(boost_state)
+            logits = np.array(token_logits[0, :blank].astype(mx.float32))
+            pred_token = int(np.argmax(logits + alpha * bonus))
+            boost_state = int(next_states[pred_token])
+            hypothesis.append(AlignedToken(
+                pred_token, start=step * model.time_ratio, duration=model.durations[decision] * model.time_ratio,
+                confidence=confidence, text=tokenizer.decode([pred_token], model.vocabulary)))
+            last_token = pred_token
+            hidden_state = decoder_hidden
+
+        step += model.durations[decision]
+        new_symbols += 1
+        if model.durations[decision] != 0:
+            new_symbols = 0
+        elif model.max_symbols is not None and model.max_symbols <= new_symbols:
+            step += 1
+            new_symbols = 0
+    return hypothesis

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eacec2b0a77f336d4a2ca4a25a7047575d3c2b74de47e997f4c205126ed3135e
+size 360916