| |
| """Vietnamese capitalization and punctuation restoration with the ONNX export of dragonSwing/xlm-roberta-capu. |
| |
| The reference for the Kotlin and Swift ports: it needs only numpy, onnxruntime and the files of the published repo |
| (onnx/*.onnx, unigram.tsv, labels.txt, number_words.tsv), and does every step itself, the way an app must. |
| It is a port of the model repo's GecBERTModel (gec_model.py, utils.py) with the model card's settings |
| (split_chunk=True): same chunking, the same up-to-3 correction rounds, the same chunk merging, so it gives the same |
| text (validate.py checks this). |
| |
| capu = Capu("path/to/repo", quantized=True) |
| capu("theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ") # "Theo đó, Thủ tướng dự kiến tiếp Bộ trưởng Nông nghiệp Mỹ." |
| |
| Usage: |
| python -I capu.py path/to/repo "text to correct" [--fp32] |
| """ |
|
|
| import re |
| import sys |
| import unicodedata |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
|
|
| START = "$START" |
| UNK, PAD = "@@UNKNOWN@@", "@@PADDING@@" |
| PUNCTUATION = (":", ".", ",", "?") |
| MAX_WORDS, MIN_WORDS, ITERATIONS = 64, 3, 3 |
| CHUNK_SIZE, OVERLAP_SIZE, MIN_WORDS_CUT = 48, 12, 6 |
| MAX_TOKENS = 512 |
|
|
|
|
| class Tokenizer: |
| """XLM-RoBERTa's SentencePiece unigram model (as in tokenizers' Unigram), one word at a time. |
| |
| Each word is NFKC-normalized, prefixed with "▁" and split by Viterbi search for the pieces with the highest |
| total log-probability. A character that no piece covers becomes <unk>, with consecutive <unk>s fused. |
| """ |
|
|
| def __init__(self, path: Path): |
| self.ids: dict[str, int] = {} |
| self.scores: list[float] = [] |
| with open(path, encoding="utf-8") as f: |
| for i, line in enumerate(f): |
| piece, score = line.rstrip("\n").split("\t") |
| self.ids.setdefault(piece, i) |
| self.scores.append(float(score)) |
| self.start_id = self.ids[START] |
| self.unk_id = self.ids["<unk>"] |
| self.unk_score = min(self.scores[4:self.start_id]) - 10.0 |
| self.max_len = max(len(p) for p in self.ids) |
|
|
| def encode_word(self, word: str) -> list[int]: |
| text = "▁" + unicodedata.normalize("NFKC", word) |
| n = len(text) |
| best = [-np.inf] * (n + 1) |
| back: list[tuple[int, int]] = [(0, 0)] * (n + 1) |
| best[0] = 0.0 |
| for start in range(n): |
| if best[start] == -np.inf: |
| continue |
| single = False |
| for end in range(start + 1, min(n, start + self.max_len) + 1): |
| token = self.ids.get(text[start:end]) |
| if token is None or token == self.start_id: |
| continue |
| single |= end == start + 1 |
| score = best[start] + self.scores[token] |
| if score > best[end]: |
| best[end], back[end] = score, (start, token) |
| if not single and best[start] + self.unk_score > best[start + 1]: |
| best[start + 1], back[start + 1] = best[start] + self.unk_score, (start, self.unk_id) |
| tokens, end = [], n |
| while end > 0: |
| start, token = back[end] |
| if not (token == self.unk_id and tokens and tokens[-1] == self.unk_id): |
| tokens.append(token) |
| end = start |
| return tokens[::-1] |
|
|
|
|
| def capitalize(s: str) -> str: |
| return s[:1].upper() + s[1:].lower() |
|
|
|
|
| def transform_case(token: str, label: str) -> str: |
| |
| if label.endswith("LOWER"): |
| return token.lower() |
| if label.endswith("UPPER"): |
| return token.upper() |
| if label.endswith("CAPITAL"): |
| return capitalize(token) |
| if label.endswith("CAPITAL_1"): |
| return token[0] + capitalize(token[1:]) |
| if label.endswith("UPPER_-1"): |
| return token[:-1].upper() + token[-1] |
| return token |
|
|
|
|
| class Capu: |
| def __init__(self, repo: str | Path, quantized: bool = True, threads: int = 4): |
| repo = Path(repo) |
| opts = ort.SessionOptions() |
| opts.intra_op_num_threads = threads |
| name = "model_quantized.onnx" if quantized else "model.onnx" |
| self.session = ort.InferenceSession(str(repo / "onnx" / name), opts, providers=["CPUExecutionProvider"]) |
| self.tokenizer = Tokenizer(repo / "unigram.tsv") |
| self.labels = (repo / "labels.txt").read_text(encoding="utf-8").splitlines() |
| self.number_words: dict[tuple[str, str], str] = {} |
| for line in (repo / "number_words.tsv").read_text(encoding="utf-8").splitlines(): |
| word, tag, replacement = line.split("\t") |
| self.number_words[(word, tag)] = replacement |
|
|
| def __call__(self, text: str) -> str: |
| chunks = split_chunks(text.split()) |
| merged: list[str] = [] |
| for chunk in chunks: |
| chunk = self.correct(chunk) |
| merged = chunk if len(chunks) == 1 else merge_chunks(merged, chunk) |
| return re.sub(r"\s+([:.,?])", r"\1", " ".join(merged)) |
|
|
| def correct(self, tokens: list[str]) -> list[str]: |
| """GecBERTModel.handle_batch for one chunk: up to 3 rounds, stopping when the text stops changing.""" |
| if len(tokens) < MIN_WORDS: |
| return tokens |
| seen = [tokens] |
| for _ in range(ITERATIONS): |
| new = self.apply(tokens, self.predict(tokens)) |
| if new == tokens: |
| break |
| tokens = new |
| if new in seen: |
| break |
| seen.append(new) |
| return tokens |
|
|
| def predict(self, tokens: list[str]) -> list[int]: |
| """Label id for $START and each of the first MAX_WORDS tokens.""" |
| ids, offsets = [self.tokenizer.start_id], [0] |
| for word in tokens[:MAX_WORDS]: |
| offsets.append(len(ids)) |
| ids += self.tokenizer.encode_word(word) |
| if len(ids) > MAX_TOKENS: |
| |
| words = sum(o < MAX_TOKENS for o in offsets) |
| ids, offsets = ids[:MAX_TOKENS], offsets[:words] |
| logits = self.session.run(["logits"], { |
| "input_ids": np.array([ids], np.int64), |
| "attention_mask": np.ones((1, len(ids)), np.int64), |
| "input_offsets": np.array([offsets], np.int64), |
| })[0][0] |
| labels = logits.argmax(-1).tolist() |
| return labels + [0] * (min(len(tokens), MAX_WORDS) + 1 - len(labels)) |
|
|
| def apply(self, tokens: list[str], labels: list[int]) -> list[str]: |
| """GecBERTModel.postprocess_batch and utils.get_target_sent_by_edits.""" |
| if max(labels) == 0: |
| return tokens |
| edits = [] |
| for i in range(min(len(tokens), MAX_WORDS) + 1): |
| label = self.labels[labels[i]] |
| if label in (UNK, PAD, "$KEEP"): |
| continue |
| if label.startswith("$APPEND_"): |
| edits.append((i, i, label[len("$APPEND_"):])) |
| elif label.startswith("$MERGE_"): |
| edits.append((i, i, label)) |
| else: |
| edits.append((i - 1, i, label)) |
|
|
| target, shift = tokens[:], 0 |
| for start, end, label in edits: |
| pos = start + shift |
| if start < 0: |
| continue |
| source = target[pos] if pos < len(target) else "" |
| if start == end: |
| if (pos < len(target) and target[pos] == label) or (pos > 0 and target[pos - 1] == label): |
| continue |
| target.insert(pos, label) |
| shift += 1 |
| elif label.startswith("$TRANSFORM_CASE"): |
| target[pos] = transform_case(source, label) |
| elif label.startswith("$TRANSFORM_VERB_"): |
| target[pos] = self.number_words.get((source, label[len("$TRANSFORM_VERB_"):]), source) |
| return replace_merges(target) |
|
|
|
|
| def replace_merges(tokens: list[str]) -> list[str]: |
| """utils.replace_merge_transforms: "5 $MERGE_SPACE G" -> "5G".""" |
| if not any(t.startswith("$MERGE_") for t in tokens): |
| return tokens |
| if tokens[0].startswith("$MERGE_"): |
| tokens = tokens[1:] |
| if tokens and tokens[-1].startswith("$MERGE_"): |
| tokens = tokens[:-1] |
| line = " ".join(tokens).replace(" $MERGE_HYPHEN ", "-").replace(" $MERGE_SPACE ", "") |
| return re.sub(r"([.,?:]\s+)+", r"\1", line).split() |
|
|
|
|
| def split_chunks(words: list[str]) -> list[list[str]]: |
| """GecBERTModel.split_chunks: overlapping windows of CHUNK_SIZE words.""" |
| n = len(words) |
| if n <= CHUNK_SIZE: |
| return [words] |
| if n < CHUNK_SIZE * 2 - OVERLAP_SIZE: |
| split = (n + OVERLAP_SIZE + 1) // 2 |
| return [words[:split], words[split - OVERLAP_SIZE:]] |
| return [words[i:i + CHUNK_SIZE] for i in range(0, n - OVERLAP_SIZE, CHUNK_SIZE - OVERLAP_SIZE)] |
|
|
|
|
| def merge_chunks(tokens: list[str], following: list[str]) -> list[str]: |
| """GecBERTModel.apply_chunk_merging: joins two corrected chunks inside their overlap. |
| |
| Aligns the last OVERLAP_SIZE words of one chunk with the first OVERLAP_SIZE of the next (punctuation ignored, |
| lower-cased) and cuts at the first common run at least MIN_WORDS_CUT words into the overlap. Where the original |
| finds no cut it raises, and its caller drops the following chunk; this keeps that behavior. |
| """ |
| if not tokens: |
| return following |
| source_idx, source = [], [] |
| i = 0 |
| while len(source_idx) < OVERLAP_SIZE and -i < len(tokens): |
| i -= 1 |
| if tokens[i] not in PUNCTUATION: |
| source_idx.insert(0, i) |
| source.insert(0, tokens[i].lower()) |
| target_idx, target = [], [] |
| i = 0 |
| while len(target_idx) < OVERLAP_SIZE and i < len(following): |
| if following[i] not in PUNCTUATION: |
| target_idx.append(i) |
| target.append(following[i].lower()) |
| i += 1 |
|
|
| keep = OVERLAP_SIZE - MIN_WORDS_CUT |
| for tag, i1, i2, j1, j2 in opcodes(source, target): |
| if tag == "equal": |
| if i1 >= keep: |
| return tokens[:source_idx[i1]] + following[target_idx[j1]:] |
| if i2 > keep: |
| j = j2 - i2 + keep |
| if j >= len(target_idx): |
| return tokens |
| return tokens[:source_idx[keep]] + following[target_idx[j]:] |
| elif tag == "delete" and i1 == 0: |
| keep += i2 // 2 |
| return tokens |
|
|
|
|
| def matching_blocks(a: list[str], b: list[str]) -> list[tuple[int, int, int]]: |
| """difflib.SequenceMatcher(None, a, b).get_matching_blocks() for short lists (no junk, no autojunk).""" |
|
|
| def longest(alo, ahi, blo, bhi): |
| |
| besti, bestj, bestsize = alo, blo, 0 |
| lengths: dict[int, int] = {} |
| for i in range(alo, ahi): |
| current = {} |
| for j in range(blo, bhi): |
| if a[i] == b[j]: |
| k = current[j] = lengths.get(j - 1, 0) + 1 |
| if k > bestsize: |
| besti, bestj, bestsize = i - k + 1, j - k + 1, k |
| lengths = current |
| return besti, bestj, bestsize |
|
|
| queue, blocks = [(0, len(a), 0, len(b))], [] |
| while queue: |
| alo, ahi, blo, bhi = queue.pop() |
| i, j, k = longest(alo, ahi, blo, bhi) |
| if k: |
| blocks.append((i, j, k)) |
| if alo < i and blo < j: |
| queue.append((alo, i, blo, j)) |
| if i + k < ahi and j + k < bhi: |
| queue.append((i + k, ahi, j + k, bhi)) |
| blocks.sort() |
| merged = [] |
| for i, j, k in blocks: |
| if merged and merged[-1][0] + merged[-1][2] == i and merged[-1][1] + merged[-1][2] == j: |
| merged[-1] = (merged[-1][0], merged[-1][1], merged[-1][2] + k) |
| else: |
| merged.append((i, j, k)) |
| merged.append((len(a), len(b), 0)) |
| return merged |
|
|
|
|
| def opcodes(a: list[str], b: list[str]) -> list[tuple[str, int, int, int, int]]: |
| """difflib.SequenceMatcher(None, a, b).get_opcodes().""" |
| i = j = 0 |
| result = [] |
| for ai, bj, size in matching_blocks(a, b): |
| tag = "replace" if i < ai and j < bj else "delete" if i < ai else "insert" if j < bj else "" |
| if tag: |
| result.append((tag, i, ai, j, bj)) |
| i, j = ai + size, bj + size |
| if size: |
| result.append(("equal", ai, i, bj, j)) |
| return result |
|
|
|
|
| if __name__ == "__main__": |
| if len(sys.argv) < 3: |
| sys.exit(f"usage: {sys.argv[0]} path/to/repo 'text' [--fp32]") |
| print(Capu(sys.argv[1], quantized="--fp32" not in sys.argv)(sys.argv[2])) |
|
|
