| """Dewpoint: multilingual punctuation restoration and truecasing.
|
|
|
| Self-contained inference for the released checkpoints, in one file, with two
|
| interchangeable backends:
|
|
|
| from dewpoint import Punctuator
|
| p = Punctuator.from_pretrained("valkayuh/dewpoint")
|
| p.restore("so i said meet at three thirty tuesday what do you think", lang="en")
|
| # -> 'So I said meet at three thirty Tuesday. What do you think?'
|
|
|
| backend="torch" needs torch + transformers + safetensors, and uses a GPU if present.
|
| backend="onnx" needs only onnxruntime + tokenizers + numpy (no torch at all),
|
| and downloads the ONNX graphs instead of the safetensors.
|
| backend="auto" (the default) uses torch when it is installed, otherwise ONNX.
|
|
|
| The model is an ensemble of two dual-head token taggers (mmBERT-base and
|
| XLM-RoBERTa-large). Each member reads the same word list and returns one
|
| posterior per word; the posteriors are averaged in probability space, a
|
| per-class decision bias fitted on IWSLT dev2010 is applied, and the argmax is
|
| taken. Pass ``members=["mmbert-base"]`` for the single 307M-parameter model,
|
| which uses its own calibration.
|
|
|
| It is a tagger, not a generator: it never adds, drops, reorders or rewrites a
|
| word. It only decides, per word, which mark follows it (none / , / . / ? / !)
|
| and how it is cased (lower / Capitalised / UPPER).
|
| """
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import re
|
| import unicodedata
|
| from typing import Iterable, Sequence
|
|
|
| import numpy as np
|
|
|
| try:
|
| import torch
|
| import torch.nn as nn
|
| except ImportError:
|
| torch = None
|
| nn = None
|
|
|
| __all__ = ["Punctuator", "StreamingPunctuator"]
|
|
|
| PUNCT_LABELS = ["O", "COMMA", "PERIOD", "QUESTION", "EXCLAM"]
|
| CASE_LABELS = ["LOWER", "CAP", "UPPER"]
|
| _SENT_END = {"PERIOD", "QUESTION", "EXCLAM"}
|
|
|
|
|
|
|
| CASELESS_LANGS = {
|
| "zh", "zh-cn", "zh-tw", "ja", "ko", "ar", "he", "fa", "ur", "th", "hi", "bn",
|
| "ta", "te", "ml", "kn", "mr", "gu", "pa", "ne", "si", "my", "km", "lo", "am",
|
| "ti", "dv", "bo", "dz", "ps", "sd", "ug", "yi", "yue", "wuu", "as", "or",
|
| "arq", "arz", "ary", "acm", "apc", "ckb", "ka", "he-il", "sa", "ks", "bho",
|
| }
|
|
|
| NO_SPACE_LANGS = {"zh", "zh-cn", "zh-tw", "ja", "th", "my", "km", "lo", "bo",
|
| "dz", "yue", "wuu"}
|
|
|
| _SURFACE = {
|
| "ascii": {"O": "", "COMMA": ",", "PERIOD": ".", "QUESTION": "?", "EXCLAM": "!"},
|
| "cjk": {"O": "", "COMMA": "，", "PERIOD": "。", "QUESTION": "？", "EXCLAM": "！"},
|
| "ja": {"O": "", "COMMA": "、", "PERIOD": "。", "QUESTION": "？", "EXCLAM": "！"},
|
| "arab": {"O": "", "COMMA": "،", "PERIOD": ".", "QUESTION": "؟", "EXCLAM": "!"},
|
| "deva": {"O": "", "COMMA": ",", "PERIOD": "।", "QUESTION": "?", "EXCLAM": "!"},
|
|
|
| "el": {"O": "", "COMMA": ",", "PERIOD": ".", "QUESTION": ";", "EXCLAM": "!"},
|
| "hy": {"O": "", "COMMA": ",", "PERIOD": "։", "QUESTION": "՞", "EXCLAM": "՜"},
|
| }
|
| _SURFACE_BY_LANG = {}
|
| for _l in ("zh", "zh-cn", "zh-tw", "yue", "wuu"):
|
| _SURFACE_BY_LANG[_l] = _SURFACE["cjk"]
|
| for _l in ("ar", "fa", "ur", "ps", "arq", "arz", "ary", "ckb", "sd", "ug"):
|
| _SURFACE_BY_LANG[_l] = _SURFACE["arab"]
|
| for _l in ("hi", "bn", "mr", "ne", "pa", "as", "or", "sa", "bho"):
|
| _SURFACE_BY_LANG[_l] = _SURFACE["deva"]
|
| _SURFACE_BY_LANG.update(el=_SURFACE["el"], hy=_SURFACE["hy"], ja=_SURFACE["ja"])
|
|
|
|
|
|
|
| _MARKS = set(
|
| ".。．۔।॥။։។…⋯።,，،٫၊、᠂᠈፣?？؟;՞⸮፧!！՜:：；؛·-–—―−፡"
|
| "\"'‘’“”„‚«»()[]{}〈〉《》「」『』（）‹›‟″′¡¿⸘‛*_`•⁃・＂＇´‐‑"
|
| )
|
| _LATIN_RUN = re.compile(r"[A-Za-z0-9À-ɏ]+(?:[.'’-][A-Za-z0-9]+)*")
|
|
|
|
|
| def base_lang(lang: str) -> str:
|
| return lang.split("-")[0].lower()
|
|
|
|
|
| def is_caseless(lang: str) -> bool:
|
| return lang.lower() in CASELESS_LANGS or base_lang(lang) in CASELESS_LANGS
|
|
|
|
|
| def is_no_space(lang: str) -> bool:
|
| return lang.lower() in NO_SPACE_LANGS or base_lang(lang) in NO_SPACE_LANGS
|
|
|
|
|
| def surface_for(lang: str) -> dict:
|
| return _SURFACE_BY_LANG.get(base_lang(lang), _SURFACE["ascii"])
|
|
|
|
|
| def apply_case(word: str, case: str) -> str:
|
| if case == "UPPER":
|
| return word.upper()
|
| if case == "CAP":
|
| return word[:1].upper() + word[1:]
|
| return word
|
|
|
|
|
| def split_words(text: str, lang: str) -> list[str]:
|
| """Raw text -> the lowercased, mark-free word list the model was trained on."""
|
| text = unicodedata.normalize("NFC", text)
|
| if is_no_space(lang):
|
| return _split_charwise(text)
|
| out = []
|
| for raw in text.split():
|
| chars = [i for i, ch in enumerate(raw) if ch not in _MARKS]
|
| if not chars:
|
| continue
|
| core = raw[chars[0]:chars[-1] + 1]
|
| if any(ch.isalnum() for ch in core):
|
| out.append(core.lower())
|
| return out
|
|
|
|
|
| def _split_charwise(text: str) -> list[str]:
|
| """CJK / Thai / Khmer: grapheme clusters, with Latin runs kept whole."""
|
| words, i, n = [], 0, len(text)
|
| while i < n:
|
| ch = text[i]
|
| if ch.isspace() or ch in _MARKS:
|
| i += 1
|
| continue
|
| m = _LATIN_RUN.match(text, i)
|
| if m and m.group():
|
| words.append(m.group().lower())
|
| i = m.end()
|
| continue
|
| if ch.isalnum() or unicodedata.category(ch) in ("Mn", "Mc", "Me"):
|
| j = i + 1
|
| while j < n and unicodedata.category(text[j]) in ("Mn", "Mc", "Me"):
|
| j += 1
|
| words.append(text[i:j].lower())
|
| i = j
|
| continue
|
| i += 1
|
| return words
|
|
|
|
|
|
|
| if torch is not None:
|
| class PunctCaseModel(nn.Module):
|
| """Encoder -> residual shared trunk -> punctuation head and case head."""
|
|
|
| def __init__(self, config, n_punct=5, n_case=3):
|
| super().__init__()
|
| from transformers import AutoModel
|
| self.encoder = AutoModel.from_config(config)
|
| h = config.hidden_size
|
| self.drop = nn.Dropout(0.1)
|
| self.trunk = nn.Sequential(nn.Linear(h, h), nn.GELU(),
|
| nn.LayerNorm(h, eps=1e-5))
|
| self.head_punct = nn.Linear(h, n_punct)
|
| self.head_case = nn.Linear(h, n_case)
|
|
|
| def forward(self, input_ids, attention_mask=None):
|
| x = self.encoder(input_ids=input_ids,
|
| attention_mask=attention_mask).last_hidden_state
|
| x = self.drop(x)
|
| x = x + self.trunk(x)
|
| return self.head_punct(x), self.head_case(x)
|
|
|
| __all__.append("PunctCaseModel")
|
|
|
|
|
| def plan_windows(sub_counts, max_sub=508, overlap_sub=256):
|
| """Overlapping windows under a subword budget; only each window's centre
|
| commits. Returns (win_start, win_end, commit_start, commit_end) in words."""
|
| n = len(sub_counts)
|
| if n == 0:
|
| return []
|
| overlap_sub = min(overlap_sub, max_sub // 2)
|
| prefix = [0] * (n + 1)
|
| for i, c in enumerate(sub_counts):
|
| prefix[i + 1] = prefix[i] + max(1, c)
|
|
|
| def end_for(start):
|
| lo, hi = start + 1, n
|
| budget = prefix[start] + max_sub
|
| while lo < hi:
|
| mid = (lo + hi + 1) // 2
|
| if prefix[mid] <= budget:
|
| lo = mid
|
| else:
|
| hi = mid - 1
|
| return max(lo, start + 1)
|
|
|
| wins, start = [], 0
|
| while True:
|
| end = end_for(start)
|
| wins.append([start, end])
|
| if end >= n:
|
| break
|
| target = prefix[end] - overlap_sub
|
| nxt = end
|
| while nxt > start + 1 and prefix[nxt] > target:
|
| nxt -= 1
|
| start = max(nxt, start + 1)
|
| out = []
|
| for i, (s, e) in enumerate(wins):
|
| cs = s if i == 0 else (wins[i - 1][1] + s) // 2
|
| ce = e if i == len(wins) - 1 else (e + wins[i + 1][0] + 1) // 2
|
| cs = max(cs, s)
|
| out.append((s, e, cs, min(max(ce, cs), e)))
|
| return out
|
|
|
|
|
| class _MemberBase:
|
| """Windowing, first-subword gathering and caseless masking, shared by both
|
| backends so that they cannot drift apart. Subclasses supply `_counts`
|
| (subwords per word), `_encode` (a padded batch plus word ids) and
|
| `_forward` (softmaxed posteriors as numpy)."""
|
|
|
| def _load_meta(self, path):
|
| with open(os.path.join(path, "vpunct_config.json"), encoding="utf-8") as f:
|
| self.meta = json.load(f)
|
| self.max_len = int(self.meta.get("max_len", 512))
|
|
|
| def posteriors(self, words, lang):
|
| n = len(words)
|
| pp = np.zeros((n, len(PUNCT_LABELS)), dtype=np.float32)
|
| pp[:, 0] = 1.0
|
| cp = np.zeros((n, len(CASE_LABELS)), dtype=np.float32)
|
| cp[:, 0] = 1.0
|
| if n == 0:
|
| return pp, cp
|
| wins = plan_windows(self._counts(list(words)), max_sub=self.max_len - 4,
|
| overlap_sub=256)
|
| for i in range(0, len(wins), self.batch_size):
|
| chunk = wins[i:i + self.batch_size]
|
| ids, mask, wids = self._encode([list(words[s:e]) for s, e, _, _ in chunk])
|
| p, c = self._forward(ids, mask)
|
| for b, (s, e, cs, ce) in enumerate(chunk):
|
| loc_p = np.zeros((e - s, len(PUNCT_LABELS)), dtype=np.float32)
|
| loc_p[:, 0] = 1.0
|
| loc_c = np.zeros((e - s, len(CASE_LABELS)), dtype=np.float32)
|
| loc_c[:, 0] = 1.0
|
| seen = set()
|
| for pos, w in enumerate(wids[b]):
|
| if w is None or w in seen or pos >= p.shape[1]:
|
| continue
|
| seen.add(w)
|
| loc_p[w] = p[b, pos]
|
| loc_c[w] = c[b, pos]
|
| pp[cs:ce] = loc_p[cs - s:ce - s]
|
| cp[cs:ce] = loc_c[cs - s:ce - s]
|
| if is_caseless(lang):
|
| cp[:] = 0.0
|
| cp[:, 0] = 1.0
|
| return pp, cp
|
|
|
|
|
| def _softmax(x):
|
| x = x - x.max(-1, keepdims=True)
|
| e = np.exp(x)
|
| return e / e.sum(-1, keepdims=True)
|
|
|
|
|
| class _TorchMember(_MemberBase):
|
| """One tagger run by torch, from its safetensors."""
|
|
|
| def __init__(self, path, device, dtype, batch_size=64):
|
| from safetensors.torch import load_file
|
| from transformers import AutoConfig, AutoTokenizer
|
| self._load_meta(path)
|
| cfg = AutoConfig.from_pretrained(path)
|
| model = PunctCaseModel(cfg, self.meta["n_punct"], self.meta["n_case"])
|
| model.load_state_dict(load_file(os.path.join(path, "model.safetensors")))
|
| self.model = model.to(device=device, dtype=dtype).eval()
|
| try:
|
| self.model.encoder.config._attn_implementation = "sdpa"
|
| except Exception:
|
| pass
|
| self.tok = AutoTokenizer.from_pretrained(path)
|
| self.device = device
|
| self.batch_size = batch_size
|
|
|
| def _counts(self, words):
|
| return [len(e) for e in self.tok(words, add_special_tokens=False)["input_ids"]]
|
|
|
| def _encode(self, batch):
|
| enc = self.tok(batch, is_split_into_words=True, truncation=True,
|
| max_length=self.max_len, padding=True, return_tensors="pt")
|
| return enc["input_ids"], enc["attention_mask"], \
|
| [enc.word_ids(b) for b in range(len(batch))]
|
|
|
| def _forward(self, ids, mask):
|
| with torch.no_grad():
|
| lp, lc = self.model(ids.to(self.device), mask.to(self.device))
|
| return (lp.float().softmax(-1).cpu().numpy(),
|
| lc.float().softmax(-1).cpu().numpy())
|
|
|
|
|
| class _OnnxMember(_MemberBase):
|
| """One tagger run by onnxruntime, tokenised by the `tokenizers` library.
|
| No torch and no transformers anywhere on this path."""
|
|
|
| def __init__(self, path, onnx_dir, providers=None, batch_size=8, threads=None):
|
| import onnxruntime as ort
|
| from tokenizers import Tokenizer
|
| self._load_meta(path)
|
| self.tok = Tokenizer.from_file(os.path.join(path, "tokenizer.json"))
|
| self.tok.no_padding()
|
|
|
|
|
| self.tok.enable_truncation(max_length=self.max_len)
|
| with open(os.path.join(path, "tokenizer_config.json"), encoding="utf-8") as f:
|
| tcfg = json.load(f)
|
| self.pad_id = self.tok.token_to_id(tcfg.get("pad_token", "<pad>")) or 0
|
| so = ort.SessionOptions()
|
| if threads:
|
| so.intra_op_num_threads = threads
|
| if providers is None:
|
| avail = ort.get_available_providers()
|
| providers = [p for p in ("CUDAExecutionProvider", "CoreMLExecutionProvider",
|
| "DmlExecutionProvider") if p in avail]
|
| providers.append("CPUExecutionProvider")
|
| self.sess = ort.InferenceSession(os.path.join(onnx_dir, "model.onnx"), so,
|
| providers=providers)
|
| self.batch_size = batch_size
|
|
|
| def _counts(self, words):
|
| return [len(e.ids) for e in
|
| self.tok.encode_batch(words, add_special_tokens=False)]
|
|
|
| def _encode(self, batch):
|
| encs = self.tok.encode_batch(batch, is_pretokenized=True)
|
| L = min(self.max_len, max(len(e.ids) for e in encs))
|
| ids = np.full((len(encs), L), self.pad_id, dtype=np.int64)
|
| mask = np.zeros((len(encs), L), dtype=np.int64)
|
| wids = []
|
| for b, e in enumerate(encs):
|
| n = min(L, len(e.ids))
|
| ids[b, :n] = e.ids[:n]
|
| mask[b, :n] = 1
|
| wids.append(list(e.word_ids[:n]))
|
| return ids, mask, wids
|
|
|
| def _forward(self, ids, mask):
|
| lp, lc = self.sess.run(None, {"input_ids": ids, "attention_mask": mask})
|
| return _softmax(lp.astype(np.float32)), _softmax(lc.astype(np.float32))
|
|
|
|
|
|
|
| class Punctuator:
|
| """Restore punctuation and case in unpunctuated, lowercased text."""
|
|
|
| MEMBERS = ("mmbert-base", "xlm-roberta-large")
|
|
|
| def __init__(self, path, members=None, backend="auto", device=None, dtype=None,
|
| use_gazetteer=True, batch_size=None, providers=None, threads=None):
|
| members = list(members or self.MEMBERS)
|
| self.backend = self._pick_backend(path, members, backend)
|
| if self.backend == "torch":
|
| if device is None:
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| if dtype is None:
|
| dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
|
| self.members = [_TorchMember(os.path.join(path, m), device, dtype,
|
| batch_size or 64) for m in members]
|
| else:
|
| self.members = [_OnnxMember(os.path.join(path, m),
|
| os.path.join(path, "onnx", m), providers,
|
| batch_size or 8, threads) for m in members]
|
| if len(members) == len(self.MEMBERS):
|
|
|
|
|
| with open(os.path.join(path, "ensemble_config.json"), encoding="utf-8") as f:
|
| ens = json.load(f)
|
| self.weights = ens["weights"]
|
| self.bias = ens["punct_bias"]
|
| self.bias_per_lang = ens["punct_bias_per_lang"]
|
| else:
|
| self.weights = [1.0 / len(members)] * len(members)
|
| meta = self.members[0].meta if len(members) == 1 else {}
|
| self.bias = meta.get("punct_bias")
|
| self.bias_per_lang = meta.get("punct_bias_per_lang") or {}
|
| gz = os.path.join(path, "gazetteer.json")
|
| self.gazetteer = {}
|
| if use_gazetteer and os.path.exists(gz):
|
| with open(gz, encoding="utf-8") as f:
|
| self.gazetteer = json.load(f)
|
|
|
| @staticmethod
|
| def _pick_backend(path, members, backend):
|
| if backend == "auto":
|
| has_st = all(os.path.exists(os.path.join(path, m, "model.safetensors"))
|
| for m in members)
|
| backend = "torch" if (torch is not None and has_st) else "onnx"
|
| if backend == "torch" and torch is None:
|
| raise ImportError("backend='torch' needs torch; install it, or use "
|
| "backend='onnx' (onnxruntime + tokenizers only)")
|
| if backend not in ("torch", "onnx"):
|
| raise ValueError("backend must be 'auto', 'torch' or 'onnx'")
|
| return backend
|
|
|
| @classmethod
|
| def from_pretrained(cls, repo_or_path="valkayuh/dewpoint", **kw):
|
| """A local directory, or a Hugging Face repo id to download. Only the
|
| files the chosen backend needs are fetched: the ONNX backend never
|
| downloads the safetensors, and the torch backend never downloads ONNX."""
|
| if os.path.isdir(repo_or_path):
|
| return cls(repo_or_path, **kw)
|
| from huggingface_hub import snapshot_download
|
| members = kw.get("members") or cls.MEMBERS
|
| backend = kw.get("backend", "auto")
|
| if backend == "auto":
|
| backend = "torch" if torch is not None else "onnx"
|
| kw["backend"] = backend
|
| allow = ["dewpoint.py", "ensemble_config.json", "gazetteer.json"]
|
| for m in members:
|
| allow += [f"{m}/vpunct_config.json", f"{m}/tokenizer.json",
|
| f"{m}/tokenizer_config.json"]
|
| allow += ([f"{m}/config.json", f"{m}/model.safetensors"]
|
| if backend == "torch" else [f"onnx/{m}/*"])
|
| return cls(snapshot_download(repo_or_path, allow_patterns=allow), **kw)
|
|
|
|
|
| def _bias_for(self, lang):
|
| b = None
|
| if self.bias_per_lang:
|
| b = self.bias_per_lang.get(lang) or self.bias_per_lang.get(base_lang(lang))
|
| return b if b is not None else self.bias
|
|
|
| def predict(self, words: Sequence[str], lang: str = "en") -> dict:
|
| """Per-word labels and averaged posteriors for an already-split word list.
|
|
|
| ``words`` should be lowercased and free of punctuation, as ASR output is.
|
| """
|
| words = list(words)
|
| acc_p = acc_c = None
|
| for w, m in zip(self.weights, self.members):
|
| p, c = m.posteriors(words, lang)
|
| acc_p = w * p if acc_p is None else acc_p + w * p
|
| acc_c = w * c if acc_c is None else acc_c + w * c
|
| lp = np.log(np.clip(acc_p, 1e-9, None))
|
| b = self._bias_for(lang)
|
| if b is not None:
|
| lp = lp + np.asarray(b, dtype=np.float32)[None, :]
|
| punct = [PUNCT_LABELS[i] for i in lp.argmax(1)] if len(words) else []
|
| case = [CASE_LABELS[i] for i in acc_c.argmax(1)] if len(words) else []
|
| return {"words": words, "punct": punct, "case": case,
|
| "punct_probs": acc_p, "case_probs": acc_c, "punct_scores": lp}
|
|
|
| @staticmethod
|
| def _close(punct, scores):
|
| """A complete text ends on a sentence mark: if the last word got none,
|
| take the likeliest of . ? ! for it. Applied to finished text only --
|
| predict() is left exactly as benchmarked."""
|
| if punct and punct[-1] not in _SENT_END:
|
| ends = [PUNCT_LABELS.index(x) for x in ("PERIOD", "QUESTION", "EXCLAM")]
|
| punct[-1] = PUNCT_LABELS[max(ends, key=lambda i: scores[-1][i])]
|
| return punct
|
|
|
| def _postprocess(self, words, punct, case, lang, is_start=True, prev_punct=None):
|
| """Gazetteer surface forms (iPhone, McDonald) and a capital after every
|
| sentence-final mark."""
|
| case = list(case)
|
| caseless = is_caseless(lang)
|
| gz = self.gazetteer.get(base_lang(lang), {})
|
| out = []
|
| for i, w in enumerate(words):
|
| s = gz.get(w)
|
| if s and not caseless:
|
| out.append(s)
|
| case[i] = "AS_IS"
|
| else:
|
| out.append(w)
|
| if not caseless and words:
|
| lead = prev_punct in _SENT_END if prev_punct is not None else False
|
| if case[0] == "LOWER" and (is_start or lead):
|
| case[0] = "CAP"
|
| for i in range(1, len(words)):
|
| if punct[i - 1] in _SENT_END and case[i] == "LOWER":
|
| case[i] = "CAP"
|
| return out, list(punct), case
|
|
|
| def _render(self, words, punct, case, lang):
|
| surf = surface_for(lang)
|
| sep = "" if is_no_space(lang) else " "
|
| return sep.join((w if c == "AS_IS" else apply_case(w, c)) + surf.get(p, "")
|
| for w, p, c in zip(words, punct, case))
|
|
|
| def restore(self, text: str, lang: str = "en", close: bool = True) -> str:
|
| """Unpunctuated text in, punctuated and truecased text out.
|
|
|
| ``close`` ends the text on a sentence mark; turn it off when ``text``
|
| is a fragment that continues elsewhere.
|
| """
|
| words = split_words(text, lang)
|
| if not words:
|
| return text
|
| r = self.predict(words, lang)
|
| punct = self._close(r["punct"], r["punct_scores"]) if close else r["punct"]
|
| w, p, c = self._postprocess(words, punct, r["case"], lang)
|
| return self._render(w, p, c, lang)
|
|
|
| def restore_batch(self, texts: Iterable[str], lang: str = "en") -> list[str]:
|
| return [self.restore(t, lang) for t in texts]
|
|
|
| def stream(self, lang="en", lag=6, context=180, stable_n=2):
|
| return StreamingPunctuator(self, lang, lag, context, stable_n)
|
|
|
|
|
| class StreamingPunctuator:
|
| """Commit-with-lag for live transcripts.
|
|
|
| Later words change earlier decisions -- "what do you think" only becomes a
|
| question at its last word -- so a word is released only once it has
|
| ``lag`` words of right context and the same label for ``stable_n``
|
| consecutive updates. Call ``push(words)`` as ASR emits, ``finish()`` at
|
| the end; each returns the newly committed text.
|
| """
|
|
|
| def __init__(self, punctuator, lang="en", lag=6, context=180, stable_n=2):
|
| self.p, self.lang = punctuator, lang
|
| self.lag, self.context, self.stable_n = lag, context, stable_n
|
| self.words, self.emitted = [], 0
|
| self._last, self._streak = {}, {}
|
|
|
| def push(self, new_words) -> str:
|
| if isinstance(new_words, str):
|
| new_words = split_words(new_words, self.lang)
|
| self.words.extend(w.lower() for w in new_words)
|
| return self._commit(final=False)
|
|
|
| def finish(self) -> str:
|
| return self._commit(final=True)
|
|
|
| def _commit(self, final):
|
| if not self.words:
|
| return ""
|
| off = max(0, len(self.words) - self.context)
|
| r = self.p.predict(self.words[off:], self.lang)
|
| punct, case = r["punct"], r["case"]
|
| if final:
|
| punct = self.p._close(punct, r["punct_scores"])
|
| limit = len(self.words) if final else len(self.words) - self.lag
|
| done, i = [], self.emitted
|
| while i < limit:
|
| j = i - off
|
| if j < 0:
|
| i += 1
|
| continue
|
| if not final:
|
| lab = (punct[j], case[j])
|
| if self._last.get(i) == lab:
|
| self._streak[i] = self._streak.get(i, 1) + 1
|
| else:
|
| self._last[i], self._streak[i] = lab, 1
|
| if self._streak[i] < self.stable_n:
|
| break
|
| done.append(i)
|
| i += 1
|
| if not done:
|
| return ""
|
| s, e = done[0], done[-1] + 1
|
| prev = punct[s - off - 1] if s - off - 1 >= 0 else None
|
| w, p, c = self.p._postprocess(self.words[s:e], punct[s - off:e - off],
|
| case[s - off:e - off], self.lang,
|
| is_start=(s == 0), prev_punct=prev)
|
| self.emitted = e
|
| out = self.p._render(w, p, c, self.lang)
|
| return out if s == 0 or is_no_space(self.lang) else " " + out
|
|
|
|
|
| if __name__ == "__main__":
|
| import argparse
|
| import sys
|
| ap = argparse.ArgumentParser(description="Restore punctuation and case.")
|
| ap.add_argument("text", nargs="*", help="text to restore (default: stdin)")
|
| ap.add_argument("--lang", default="en", help="ISO 639-1 code, e.g. en, de, zh")
|
| ap.add_argument("--model", default=os.path.dirname(os.path.abspath(__file__)))
|
| ap.add_argument("--single", action="store_true",
|
| help="use only the mmBERT-base member (307M, faster)")
|
| ap.add_argument("--backend", default="auto", choices=["auto", "torch", "onnx"])
|
| a = ap.parse_args()
|
| p = Punctuator.from_pretrained(a.model, backend=a.backend,
|
| members=["mmbert-base"] if a.single else None)
|
| src = [" ".join(a.text)] if a.text else sys.stdin.read().splitlines()
|
| for line in src:
|
| print(p.restore(line, a.lang))
|
|
|
