| """Finnish punctuation and casing for speech recognition output (Mateli-punct, ONNX Runtime). |
| |
| python punctuate.py "mä luulen et se on ihan hyvä juttu mut en mä tiiä mitä mieltä sä oot" |
| |
| from punctuate import Punctuator |
| p = Punctuator(".") # directory with config.json, tokenizer.json, *.onnx |
| p.punctuate("mä luulen et se on ihan hyvä juttu") # -> "Mä luulen, et se on ihan hyvä juttu." |
| p.punctuate(words, pauses=[0.05, 0.4, ...]) # pauses[i] = silence after words[i] in seconds |
| |
| stream = p.stream() # live use, one word at a time |
| for word, pause_before in asr_words: |
| print(stream.push(word, pause_before), end=" ") # words come out once 3 words follow them |
| print(stream.flush()) |
| |
| Words are never changed, added or removed: each word gets one label, the punctuation after it |
| (none , . ? !) and its casing (lower, Capitalised, ALL CAPS). Pauses are optional; they help |
| mostly with sentence boundaries. Needs numpy, onnxruntime and tokenizers. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import bisect |
| import json |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from tokenizers import Tokenizer |
|
|
| SENTENCE_END = {".", "?", "!"} |
|
|
|
|
| class Model: |
| """input ids -> logits [seq, labels].""" |
|
|
| def __init__(self, model_dir: Path, precision: str = "int8", threads: int = 1): |
| cfg = json.loads((model_dir / "config.json").read_text(encoding="utf-8")) |
| so = ort.SessionOptions() |
| so.intra_op_num_threads = threads |
| so.log_severity_level = 3 |
| self.sess = ort.InferenceSession(str(model_dir / cfg["files"][precision]), so, |
| providers=["CPUExecutionProvider"]) |
| self.cfg = cfg |
|
|
| def __call__(self, ids: list[int]) -> np.ndarray: |
| x = np.array([ids], np.int64) |
| return self.sess.run(["logits"], {"input_ids": x, "attention_mask": np.ones_like(x)})[0][0] |
|
|
|
|
| class Punctuator: |
| def __init__(self, model_dir=".", precision: str = "int8", threads: int = 1, |
| window: int = 150, overlap: int = 50): |
| d = Path(model_dir) |
| self.model = Model(d, precision, threads) |
| cfg = self.model.cfg |
| self.labels = [cfg["id2label"][str(i)] for i in range(len(cfg["id2label"]))] |
| self.tok = Tokenizer.from_file(str(d / "tokenizer.json")) |
| pm = cfg["pause_markers"] |
| self.edges = pm["edges"] |
| self.flag_id = self.tok.token_to_id(pm["flag"]) |
| self.bucket_ids = [self.tok.token_to_id(b) for b in pm["buckets"]] |
| self.cls_id, self.sep_id = self.tok.token_to_id("[CLS]"), self.tok.token_to_id("[SEP]") |
| self.unk_id = self.tok.token_to_id("[UNK]") |
| if window < 1 or not 0 <= overlap < window: |
| raise ValueError("need window >= 1 and 0 <= overlap < window") |
| self.window, self.overlap = window, overlap |
|
|
| |
| def _encode(self, words, pauses): |
| """[CLS] (flag) w1 pieces (pause marker) w2 pieces ... [SEP]; returns ids and the word |
| index of each position (-1 for special tokens). The flag says pauses are given; a marker |
| follows a word whose pause is at least the first bucket edge (unknown pause: no marker).""" |
| ids, widx = [self.cls_id], [-1] |
| if pauses is not None: |
| ids.append(self.flag_id) |
| widx.append(-1) |
| for w, word in enumerate(words): |
| pieces = self.tok.encode(word, add_special_tokens=False).ids or [self.unk_id] |
| ids += pieces |
| widx += [w] * len(pieces) |
| p = None if pauses is None else pauses[w] |
| if p is not None and not math.isnan(p): |
| b = bisect.bisect_right(self.edges, p) |
| if b > 0: |
| ids.append(self.bucket_ids[b - 1]) |
| widx.append(-1) |
| return ids[:511] + [self.sep_id], widx[:511] + [-1] |
|
|
| def _call(self, words, pauses) -> list[int]: |
| """Label id of each word (the prediction on its first sub-token); -1 if truncated away.""" |
| ids, widx = self._encode(words, pauses) |
| logits = self.model(ids) |
| out = [-1] * len(words) |
| for i, w in enumerate(widx): |
| if w >= 0 and out[w] < 0: |
| out[w] = int(np.argmax(logits[i])) |
| return out |
|
|
| |
| def predict(self, words: list[str], pauses: list[float] | None = None) -> list[str]: |
| """One label per word ("O", ",|U", ".", ...). Long inputs run in overlapping windows of |
| `window` words; each word keeps the prediction from the window where it has the most |
| context on both sides.""" |
| n = len(words) |
| if pauses is not None and len(pauses) != n: |
| raise ValueError("give one pause per word (the silence after it)") |
| best = [(-1, -1)] * n |
|
|
| def run(s): |
| chunk = words[s:s + self.window] |
| labels = self._call(chunk, None if pauses is None else pauses[s:s + self.window]) |
| for i, lab in enumerate(labels): |
| if lab < 0: |
| continue |
| left = i if s > 0 else n |
| right = len(chunk) - 1 - i if s + len(chunk) < n else n |
| if min(left, right) > best[s + i][0]: |
| best[s + i] = (min(left, right), lab) |
|
|
| stride = self.window - self.overlap |
| for s in (range(0, max(1, n - self.overlap), stride) if n > self.window else [0]): |
| run(s) |
| while (i := next((i for i, (_, lab) in enumerate(best) if lab < 0), None)) is not None: |
| run(i) |
| if best[i][1] < 0: |
| raise RuntimeError(f"word {i} does not fit into the model input") |
| return [self.labels[lab] for _, lab in best] |
|
|
| def punctuate(self, text_or_words, pauses: list[float] | None = None) -> str: |
| words = text_or_words.split() if isinstance(text_or_words, str) else list(text_or_words) |
| return render(words, self.predict(words, pauses)) |
|
|
| def stream(self, **kw) -> "StreamingPunctuator": |
| return StreamingPunctuator(self, **kw) |
|
|
|
|
| def render(words: list[str], labels: list[str], capitalise_first: bool = True) -> str: |
| """Apply labels; also capitalise the first word and every word after . ? !""" |
| if len(words) != len(labels): |
| raise ValueError(f"{len(words)} words but {len(labels)} labels") |
| out, prev = [], "." if capitalise_first else "O" |
| for w, lab in zip(words, labels): |
| punct, _, case = lab.partition("|") |
| if case == "A": |
| w = w.upper() |
| elif case == "U" or prev in SENTENCE_END: |
| w = w[:1].upper() + w[1:] |
| out.append(w if punct == "O" else w + punct) |
| prev = punct |
| return " ".join(out) |
|
|
|
|
| class StreamingPunctuator: |
| """Live punctuation, one word at a time. A word is committed once `lookahead` later words are |
| known; the model sees up to `context` committed words before the pending ones and runs at |
| most every `every` words. With the defaults a word appears about 2.4 s after it is spoken |
| (at a typical speaking rate) and boundary F1 is within 2 points of whole-text punctuation.""" |
|
|
| def __init__(self, p: Punctuator, lookahead: int = 3, context: int = 64, every: int = 5): |
| if lookahead < 0 or context < 0 or every < 1: |
| raise ValueError("need lookahead >= 0, context >= 0 and every >= 1") |
| self.p, self.k, self.L, self.every = p, lookahead, context, every |
| self.context, self.pending, self.pauses = [], [], [] |
| self.new, self.prev = 0, "." |
|
|
| def push(self, word: str, pause_before: float | None = None) -> str: |
| """Add a word; pause_before is the silence since the previous word (None = no pauses). |
| Returns the newly committed text ("" if nothing is committed yet).""" |
| if self.pauses: |
| self.pauses[-1] = math.nan if pause_before is None else pause_before |
| self.pauses.append(math.nan) |
| self.pending.append(word) |
| self.new += 1 |
| if len(self.pending) <= self.k or self.new < self.every: |
| return "" |
| return self._run(final=False) |
|
|
| def flush(self, pause_after: float | None = None) -> str: |
| """Commit everything still pending; pause_after is the silence after the last word, if known.""" |
| if self.pending and pause_after is not None: |
| self.pauses[-1] = pause_after |
| parts = [] |
| while self.pending: |
| text = self._run(final=True) |
| if not text: |
| if not self.context: |
| raise RuntimeError("cannot fit the next word into the model input") |
| self.pauses = self.pauses[len(self.context):] |
| self.context = [] |
| continue |
| parts.append(text) |
| return " ".join(parts) |
|
|
| def _run(self, final: bool) -> str: |
| words = self.context + self.pending |
| use_pauses = any(not math.isnan(x) for x in self.pauses) |
| labels = self.p._call(words, self.pauses if use_pauses else None)[len(self.context):] |
| covered = next((i for i, lab in enumerate(labels) if lab < 0), len(labels)) |
| n = covered if final else max(0, covered - self.k) |
| committed = self.pending[:n] |
| text = render(committed, [self.p.labels[lab] for lab in labels[:n]], capitalise_first=False) |
| if n: |
| if self.prev in SENTENCE_END: |
| text = text[:1].upper() + text[1:] |
| self.prev = self.p.labels[labels[n - 1]].partition("|")[0] |
| self.context = (self.context + committed)[-self.L:] if self.L else [] |
| self.pending = self.pending[n:] |
| self.pauses = self.pauses[len(self.pauses) - len(self.context) - len(self.pending):] |
| self.new = 0 |
| return text |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("text", help="lowercase words without punctuation") |
| ap.add_argument("--model-dir", default=str(Path(__file__).resolve().parent)) |
| ap.add_argument("--fp32", action="store_true", help="use model.onnx instead of model.int8.onnx") |
| a = ap.parse_args() |
| print(Punctuator(a.model_dir, "fp32" if a.fp32 else "int8").punctuate(a.text)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
