Document standalone ONNX PnC quickstart

README.md

@@ -42,38 +42,112 @@ is a separate speech recognizer; PnC v2 can process its output.
 Save this as `pnc_demo.py` (Python 3.11+):
 
 ```python
+from bisect import bisect_right
+import json
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
 from huggingface_hub import snapshot_download
-from ekko import Punctuator
+from tokenizers import Tokenizer
 
-model_dir = snapshot_download(
+model_dir = Path(snapshot_download(
     repo_id="RyeAI/ekko-pnc",
     revision="d3129544e4e0155093e192cac58a9e69a4dea1ff",
     allow_patterns=["pnc.int8.onnx", "tokenizer.json", "config.json"],
-)
-
-pnc = Punctuator(model_dir, max_length=128, overlap_words=12)
-print(pnc("hej mit navn er emil hvordan går det i dag"))
+))
+tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
+labels = json.loads((model_dir / "config.json").read_text())["id2label"]
+session = ort.InferenceSession(str(model_dir / "pnc.int8.onnx"), providers=["CPUExecutionProvider"])
+
+
+def capitalize_initial(word):
+    for index, character in enumerate(word):
+        if character.isalpha():
+            return word[:index] + character.upper() + word[index + 1:]
+    return word
+
+
+def punctuate(text):
+    words = text.split()
+    if not words:
+        return ""
+
+    encoding = tokenizer.encode(words, is_pretokenized=True)
+    counts = [0] * len(words)
+    special_tokens = 0
+    for word_index in encoding.word_ids:
+        if word_index is None:
+            special_tokens += 1
+        else:
+            counts[word_index] += 1
+    prefix = [0]
+    for count in counts:
+        prefix.append(prefix[-1] + count)
+
+    def window_end(start):
+        budget = prefix[start] + 128 - special_tokens
+        return min(len(words), start + 128,
+                   max(start + 1, bisect_right(prefix, budget) - 1))
+
+    output = []
+    next_word = 0
+    while next_word < len(words):
+        start = max(0, next_word - 12)
+        end = window_end(start)
+        if end <= next_word:
+            start, end = next_word, window_end(next_word)
+        window_words = words[start:end]
+        encoded = tokenizer.encode(window_words, is_pretokenized=True)
+        inputs = {
+            "input_ids": np.array([encoded.ids], dtype=np.int64),
+            "attention_mask": np.ones((1, len(encoded.ids)), dtype=np.int64),
+        }
+        predictions = session.run(None, inputs)[0][0].argmax(-1)
+        rendered = window_words.copy()
+        seen = set()
+        for token_index, word_index in enumerate(encoded.word_ids):
+            if word_index is None or word_index in seen:
+                continue
+            seen.add(word_index)
+            label = labels[str(int(predictions[token_index]))]
+            upper = label.endswith("|U")
+            mark = label[:-2] if upper else label
+            word = window_words[word_index]
+            rendered[word_index] = (word.capitalize() if upper else word) + ("" if mark == "O" else mark)
+        emit_end = end if end == len(words) else max(next_word + 1, end - 12)
+        output.extend(rendered[next_word - start:emit_end - start])
+        next_word = emit_end
+
+    if len(output) != len(words):
+        raise RuntimeError("PnC changed the number of words")
+    output[0] = capitalize_initial(output[0])
+    for index, word in enumerate(output[:-1]):
+        abbreviation = any(word[i - 1].isalpha() and word[i + 1].isalpha()
+                           for i in range(1, len(word) - 1) if word[i] == ".")
+        if word.endswith((".", "!", "?")) and not abbreviation \
+                and not any(char.isupper() for char in output[index + 1]):
+            output[index + 1] = capitalize_initial(output[index + 1])
+    return " ".join(output)
+
+print(punctuate("hej mit navn er emil hvordan går det i dag"))
 # Hej, mit navn er Emil. Hvordan går det i dag?
 ```
 
 Run it with [uv](https://docs.astral.sh/uv/):
 
 ```bash
-uv run --with "ekko-stt[pnc] @ git+https://github.com/Rye-A1/ekko-danish-stt.git" python pnc_demo.py
+uv run --no-project --with huggingface-hub --with onnxruntime==1.29.0 \
+  --with tokenizers==0.22.2 --with numpy python pnc_demo.py
 ```
 
 The first run downloads the pinned int8 model; later runs use the Hugging Face
-cache. `Punctuator` also accepts longer text and keeps the input words in order.
-For longer transcripts, use **`max_length=128, overlap_words=12`**, as shown
-above. These are the current `Punctuator` defaults. `max_length` counts
-tokenizer subwords (including special tokens); `overlap_words` is a context
-margin measured in original words. Pass a continuous transcript when possible
-so the model can use context across ASR chunks. If you run the ONNX model
-directly, reproduce Ekko's overlapping windows and combine the word
-predictions; a single long model call will not give the same result. These
-settings change inference only; they do not change the v2 model weights. The
-[reference implementation](https://github.com/Rye-A1/ekko-danish-stt/blob/main/src/ekko/pnc.py)
-shows the window assembly.
+cache. The example uses only published packages and model files; it does not
+install code from GitHub. It mirrors the Ekko runtime's **128-token windows
+with 12 words of overlap**, so it also handles longer transcripts while keeping
+every input word in order. Pass continuous text when possible so the model can
+use context across ASR chunks. These settings affect inference, not the v2
+model weights.
 
 Version 2 is tuned for **continuous speech** such as dictation and meetings.
 The earlier version often missed sentence endings when given a long transcript;