Upload Arcee-VyLinh int4 ONNX with Android, iOS and Python code
11919c3 verified 9 days ago
44.7 kB
#!/usr/bin/env python3
"""Reference pipeline that turns a Vietnamese speech-recognition transcript into clean written text with
Qwen2.5-Instruct (ONNX): punctuation and capitals, misheard words fixed from context, fillers and stutters
removed, and lists written as "- " bullet points.
The Kotlin (kotlin/src/main/kotlin/QwenRefiner.kt) and Swift (swift/Sources/QwenRefiner/) ports follow this file
step by step, so keep the three in sync:
1. Tokenizer: Qwen2's byte-level BPE from vocab.txt and merges.txt. Text is NFC-normalized and split with Qwen2's
   pre-tokenizer regex, each piece is mapped to byte-level characters and merged by rank. Special tokens are never
   parsed out of the transcript; the prompt inserts their ids directly, so a transcript cannot inject a turn.
2. Numbers: spoken numbers of two or more words are written in digits, then dates, decimals, percentages, units,
   phone numbers and counted single digits the Vietnamese way (20/11/2026, 3,2, 20%, 12 km, 0912345678, 2 phòng),
   only where the words can mean nothing else. A small model adds numbers up wrong far more often than it copies digits.
3. Prompt: ChatML with a fixed system message and few-shot turns. Everything up to the transcript is the same for
   every call, so its KV cache is computed once when the model loads and reused.
4. Decoding: greedy, stopping at <|im_end|> or <|endoftext|>, at most max_new_tokens(n) tokens.
5. Merge: small models reword and drop words, so the output is aligned word by word with the transcript (LCS)
   and only trusted edits are kept: punctuation, capitals, line breaks and bullets; dropped fillers, stutters and
   list markers ("thứ hai là") before a bullet; a word replaced by a similar-sounding word (a misheard word
   fixed); spoken numbers written as digits with the same value, Vietnamese style (2.300.000, 3,5, 20%,
   0912345678, 20/11/2026, 4/2025); a spoken unit of measure written as its symbol (kg, mg, km, °C). Any other
   change is undone with the transcript's own words; when less than half of the transcript's words come through,
   the model answered or rewrote it, and the transcript is kept as it is.
   Last, hesitation sounds the model kept (ờ, ừm) are left out, and a yes/no question the model closed with a full
   stop (có ... không, được không, đã ... chưa) gets a question mark.
"""

import argparse
import json
import sys
import unicodedata
from pathlib import Path

import numpy as np
import onnxruntime as ort
import regex

IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
END_OF_TEXT = "<|endoftext|>"

# Qwen2's pre-tokenizer split (tokenizer.json, pre_tokenizer.pretokenizers[0].pattern.Regex).
PRETOKENIZE = regex.compile(
    r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
)

SYSTEM = """You correct Vietnamese speech-recognition transcripts. Copy the transcript word for word, and only:
- add punctuation and capital letters (sentence starts, names, places);
- replace a word the recognizer clearly misheard with the similar-sounding word the context needs (giáo việt → giáo viên, máy tín → máy tính);
- delete filler sounds (ờ, à, ừm, ừ) and words said twice by mistake;
- write numbers as digits: 3 việc, 15, 2026, 500.000 đồng, 3,5, 20%; a phone number as one run of digits; keep "một" when it means "a" (một chút, một số, một mình);
- write dates as day/month/year: ngày 20/11/2026, ngày 5/3, tháng 4/2025;
- write units of measure as symbols after the number: 5 kg, 200 mg, 3 m, 10 km, 2 l, 30°C;
- when the speaker lists several items, steps or points, put each one on its own line starting with "- ".
Never reword, shorten, translate, answer or add words. Keep names and foreign words you do not know as they are. Reply with the corrected transcript only."""

# (transcript, corrected) turns shown before every transcript. The transcripts are as the model sees them, after
# write_numbers.
EXAMPLES = [
    (
        "ờ hôm nay em cần làm 3 việc thứ nhất là đi chợ mua 2,5 kg gạo thứ hai là đón con lúc 5 giờ chiều và thứ "
        "ba là gọi điện cho cho mẹ",
        "Hôm nay em cần làm 3 việc:\n- Đi chợ mua 2,5 kg gạo.\n- Đón con lúc 5 giờ chiều.\n- Gọi điện cho mẹ.",
    ),
    (
        "bạn có biết bệnh viện bạch mai ở đâu không mình cần đến đó trước 8 giờ sáng ngày 12/3/2026",
        "Bạn có biết Bệnh viện Bạch Mai ở đâu không? Mình cần đến đó trước 8 giờ sáng ngày 12/3/2026.",
    ),
    (
        "chúng tôi đã ký hợp đồng với công ty vào thán trước và ừm sẽ giảm 10% nếu giao hàng trong tuần lày",
        "Chúng tôi đã ký hợp đồng với công ty vào tháng trước và sẽ giảm 10% nếu giao hàng trong tuần này.",
    ),
    (
        "nhà nghiên cứu jan rossby ở viện max planck cho biết từ tháng 6/2024 hệ số kờ rô càng lớn thì sao càng "
        "sáng",
        "Nhà nghiên cứu Jan Rossby ở Viện Max Planck cho biết từ tháng 6/2024 hệ số kờ rô càng lớn thì sao càng "
        "sáng.",
    ),
]

# Filler sounds the merge lets the model delete.
FILLERS = frozenset("ờ ơ à ừ ừm ừ hừm ưm ừa um uh ah ừhm".split())

# Words of a spoken list marker ("thứ hai là", "và ba là") the merge lets the model replace with a bullet.
LIST_MARKERS = frozenset("thứ nhất hai ba bốn năm sáu bảy tám chín mười một là và".split())

# Spoken digits, and the words that scale them, for checking a number the model wrote in digits.
DIGITS = {
    "không": 0, "một": 1, "mốt": 1, "hai": 2, "ba": 3, "bốn": 4, "tư": 4, "năm": 5, "lăm": 5, "nhăm": 5,
    "sáu": 6, "bảy": 7, "tám": 8, "chín": 9,
}
SCALES = {"nghìn": 1_000, "ngàn": 1_000, "triệu": 1_000_000, "tỷ": 1_000_000_000, "tỉ": 1_000_000_000}

# Units of measure the model may write as symbols after a number, by symbol: how they are said, with tone marks,
# spaces and hyphens removed (see plain()).
UNITS = {
    "kg": frozenset({"kilogam", "kilogram", "kilo", "ki", "can"}),
    "g": frozenset({"gam", "gram", "goram"}),
    "mg": frozenset({"miligam", "miligram"}),
    "km": frozenset({"kilomet", "cayso"}),
    "m": frozenset({"met"}),
    "cm": frozenset({"xentimet", "xangtimet", "centimet"}),
    "mm": frozenset({"milimet"}),
    "m²": frozenset({"metvuong"}),
    "km²": frozenset({"kilometvuong"}),
    "ha": frozenset({"hecta", "hectare"}),
    "l": frozenset({"lit"}),
    "ml": frozenset({"mililit"}),
    "°C": frozenset({"doc", "doxe"}),
    "km/h": frozenset({"kilomettrengio", "kilometmotgio", "caysomotgio"}),
}

def bytes_to_unicode() -> dict[int, str]:
    """GPT-2's byte-level alphabet: printable bytes map to themselves, the rest to chars from U+0100 up."""
    printable = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    chars, extra = {}, 0
    for b in range(256):
        if b in printable:
            chars[b] = chr(b)
        else:
            chars[b] = chr(256 + extra)
            extra += 1
    return chars


class Tokenizer:
    def __init__(self, repo: Path):
        tokens = (repo / "vocab.txt").read_text(encoding="utf-8").split("\n")[:-1]
        self.id_of = {t: i for i, t in enumerate(tokens)}
        self.tokens = tokens
        merges = (repo / "merges.txt").read_text(encoding="utf-8").split("\n")[:-1]
        self.rank = {tuple(m.split(" ")): r for r, m in enumerate(merges)}
        self.byte_char = bytes_to_unicode()
        self.char_byte = {c: b for b, c in self.byte_char.items()}
        self.cache: dict[str, list[int]] = {}

    def special(self, token: str) -> int:
        return self.id_of[token]

    def encode(self, text: str) -> list[int]:
        ids = []
        for piece in PRETOKENIZE.findall(unicodedata.normalize("NFC", text)):
            ids += self.bpe("".join(self.byte_char[b] for b in piece.encode("utf-8")))
        return ids

    def bpe(self, word: str) -> list[int]:
        cached = self.cache.get(word)
        if cached is not None:
            return cached
        parts = list(word)
        while len(parts) > 1:
            best, best_rank = -1, None
            for i in range(len(parts) - 1):
                r = self.rank.get((parts[i], parts[i + 1]))
                if r is not None and (best_rank is None or r < best_rank):
                    best, best_rank = i, r
            if best < 0:
                break
            merged = parts[best] + parts[best + 1]
            # Merge every occurrence of the best pair, left to right, as the reference BPE does.
            out, i = [], 0
            while i < len(parts):
                if i < len(parts) - 1 and parts[i] == parts[best] and parts[i + 1] == parts[best + 1]:
                    out.append(merged)
                    i += 2
                else:
                    out.append(parts[i])
                    i += 1
            parts = out
        ids = [self.id_of[p] for p in parts]
        self.cache[word] = ids
        return ids

    def decode(self, ids: list[int]) -> str:
        data = bytes(self.char_byte[c] for i in ids for c in self.tokens[i])
        return data.decode("utf-8", errors="replace")


def prompt_prefix(tok: Tokenizer) -> list[int]:
    """System message and few-shot turns, ending right before the transcript."""
    start, end, nl = tok.special(IM_START), tok.special(IM_END), tok.encode("\n")
    ids = [start] + tok.encode("system\n" + SYSTEM) + [end] + nl
    for transcript, edited in EXAMPLES:
        ids += [start] + tok.encode("user\n" + transcript) + [end] + nl
        ids += [start] + tok.encode("assistant\n" + edited) + [end] + nl
    return ids + [start] + tok.encode("user\n")


def prompt_suffix(tok: Tokenizer, transcript: str) -> list[int]:
    """The transcript and the opening of the assistant turn."""
    return tok.encode(transcript) + [tok.special(IM_END)] + tok.encode("\n") + [tok.special(IM_START)] + tok.encode("assistant\n")


def max_new_tokens(input_tokens: int) -> int:
    """Room for punctuation, bullet markers and line breaks on top of the input's length."""
    return min(input_tokens * 3 // 2 + 32, 1024)


def prepare(transcript: str) -> str:
    """NFC, whitespace collapsed. Whisper's punctuation is kept: it helps the model."""
    return " ".join(unicodedata.normalize("NFC", transcript).split())


class Unit:
    """A word of the text with what is written around it: punctuation before and after it, and whether it starts
    a line and a bullet. key is the word in lower case, for matching."""

    __slots__ = ("lead", "word", "trail", "newline", "bullet")

    def __init__(self, word: str, lead: str = "", trail: str = "", newline: bool = False, bullet: bool = False):
        self.word, self.lead, self.trail, self.newline, self.bullet = word, lead, trail, newline, bullet

    @property
    def key(self) -> str:
        return self.word.lower()


# A word: from its first to its last letter or digit (or a degree sign before one, as in °C); what is around it
# is its punctuation.
WORD_EDGES = regex.compile(r"^([^\p{L}\p{N}°]*)(.*?)([^\p{L}\p{N}]*)$", regex.DOTALL)
BULLETS = ("-", "*", "•", "–")


def units(text: str) -> list[Unit]:
    """Splits text into Units. Line breaks and bullet markers go to the next word, punctuation-only pieces to the
    next word's lead (or the last word's trail at the end)."""
    out: list[Unit] = []
    newline = bullet = False
    lead = ""
    for piece in regex.findall(r"\n|[^\s]+", text):
        if piece == "\n":
            if out:
                newline = True
            continue
        before, word, after = WORD_EDGES.match(piece).groups()
        if not word:
            if piece in BULLETS and not lead:
                bullet = True
            elif out and not newline and not bullet:
                out[-1].trail += piece
            else:
                lead += piece
            continue
        if before[:1] in BULLETS and len(before) == 1 and not lead:
            bullet, before = True, ""
        out.append(Unit(word, lead + before, after, newline, bullet))
        newline = bullet = False
        lead = ""
    return out


def render(items: list[Unit]) -> str:
    parts = []
    for i, u in enumerate(items):
        if i > 0:
            parts.append("\n" if u.newline else " ")
        parts.append(("- " if u.bullet else "") + u.lead + u.word + u.trail)
    return "".join(parts)


def plain(word: str) -> str:
    """Lower case without tone and vowel marks (đ is d), so similar-sounding words compare close."""
    decomposed = unicodedata.normalize("NFD", word.lower().replace("đ", "d"))
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def edit_distance(a: str, b: str) -> int:
    row = list(range(len(b) + 1))
    for i in range(1, len(a) + 1):
        prev, row[0] = row[0], i
        for j in range(1, len(b) + 1):
            cur = min(row[j] + 1, row[j - 1] + 1, prev + (a[i - 1] != b[j - 1]))
            prev, row[j] = row[j], cur
    return row[len(b)]


# Initial consonants Vietnamese speech merges, so a recognizer confuses them: gi-, d- and r- sound alike in the
# north, as do s- and x-, and tr- and ch-. Longest first.
SAME_SOUND_STARTS = (("gi", "d"), ("tr", "ch"), ("r", "d"), ("x", "s"))


def sound(word: str) -> str:
    """plain(word) with its first consonant folded to the one it sounds like: giục is duc, xông is song."""
    p = plain(word)
    for start, same in SAME_SOUND_STARTS:
        if p.startswith(start) and len(p) > len(start):
            return same + p[len(start):]
    return p


def sounds_alike(a: str, b: str) -> bool:
    """One edit apart without marks (and with gi/d/r, tr/ch and x/s folded) when either word is short, two when both
    are longer: lội/nội, thán/tháng, giục/dục, but not xông/đường."""
    pa, pb = sound(a), sound(b)
    return edit_distance(pa, pb) <= (1 if min(len(pa), len(pb)) <= 4 else 2)


def spoken_number(words: list[str]) -> int | None:
    """The value of a spoken Vietnamese number ("hai nghìn không trăm hai mươi sáu" is 2026), or None when the
    words are not one number."""
    total = group = 0  # group: the part below the current thousand
    digit: int | None = None
    for w in words:
        if w in DIGITS:
            if digit is not None:
                return None  # two digits in a row: a list of digits, not one number
            digit = DIGITS[w]
        elif w == "mười":
            if digit is not None:
                return None
            group += 10
        elif w == "mươi" and digit is not None:
            group, digit = group + digit * 10, None
        elif w == "trăm" and digit is not None:
            group, digit = group + digit * 100, None
        elif w in ("linh", "lẻ"):
            continue
        elif w in SCALES:
            total, group, digit = total + (group + (digit or 0)) * SCALES[w], 0, None
        else:
            return None
    if not words:
        return None
    return total + group + (digit or 0)


# Words a spoken number is built from, for writing it in digits before the model sees it.
HUNDREDS_DIGITS = frozenset("một hai ba bốn năm sáu bảy tám chín".split())  # before trăm, mươi, or a scale
TENS_DIGITS = HUNDREDS_DIGITS - {"một"}  # before mươi
AFTER_MUOI = frozenset("mốt một hai ba bốn tư lăm nhăm sáu bảy tám chín".split())  # năm after mươi is "year"
AFTER_MUOI_NUMBER = AFTER_MUOI - {"mốt", "tư"}  # after mười
NUMBER_WORDS = frozenset(DIGITS) | {"mười", "mươi", "trăm", "linh", "lẻ"} | frozenset(SCALES)
SCALE_WORDS = frozenset({"mươi", "trăm"}) | frozenset(SCALES)  # words that go on from a digit before them


def parse_below_thousand(words: list[str], i: int, after_scale: bool) -> tuple[int, int] | None:
    """(value, next index) of the longest number below 1,000 at words[i]: [digit trăm] then mười [digit],
    digit mươi [digit], or (after trăm) linh digit. "không trăm" only after a scale (2026 is "hai nghìn không
    trăm hai mươi sáu"). None when nothing parses."""
    value, j = 0, i
    hundreds = j + 1 < len(words) and words[j + 1] == "trăm" and (
        words[j] in HUNDREDS_DIGITS or (after_scale and words[j] == "không"))
    if hundreds:
        value, j = DIGITS[words[j]] * 100, j + 2
    if j < len(words) and words[j] == "mười":
        value, j = value + 10, j + 1
        if j < len(words) and words[j] in AFTER_MUOI_NUMBER:
            value, j = value + DIGITS[words[j]], j + 1
    elif j + 1 < len(words) and words[j] in TENS_DIGITS and words[j + 1] == "mươi":
        value, j = value + DIGITS[words[j]] * 10, j + 2
        if j < len(words) and words[j] in AFTER_MUOI:
            value, j = value + DIGITS[words[j]], j + 1
    elif hundreds and j + 1 < len(words) and words[j] in ("linh", "lẻ") and words[j + 1] in DIGITS:
        value, j = value + DIGITS[words[j + 1]], j + 2
    elif not hundreds and j < len(words) and words[j] in HUNDREDS_DIGITS and j + 1 < len(words) and words[j + 1] in SCALES:
        value, j = DIGITS[words[j]], j + 1  # a lone digit before nghìn, triệu, tỷ
    return (value, j) if j > i else None


def parse_number(words: list[str], i: int) -> tuple[int | None, int] | None:
    """(value, next index) of a spoken number of two or more words at words[i], with its scales falling (tỷ, then
    triệu, then nghìn), or None. The value is None when more number words follow that could belong to it ("hai
    triệu ba" is 2.3 million said short, "hai mươi ba mươi" is two numbers): the model sees those words."""
    total, j, last_scale, after_scale = 0, i, None, False
    while True:
        part = parse_below_thousand(words, j, after_scale)
        if part is None:
            break
        group, k = part
        if k < len(words) and words[k] in SCALES and (last_scale is None or SCALES[words[k]] < last_scale):
            last_scale = SCALES[words[k]]
            total, j, after_scale = total + group * last_scale, k + 1, True
            continue
        total, j = total + group, k
        break
    if j - i < 2:
        return None
    if j < len(words) and words[j] in NUMBER_WORDS:
        # "năm" after a number is "year" unless it is the digit of a number that goes on ("hai mươi năm mươi" is
        # 20-50): "tháng mười một năm 2026", "20 năm trước".
        if words[j] != "năm" or (j + 1 < len(words) and words[j + 1] in SCALE_WORDS):
            return None, j
    return total, j


def write_number(value: int) -> str:
    """Digits, with "." between thousands from five digits up (2026, 10.000, 2.300.000), as Vietnamese writes them."""
    return f"{value:,}".replace(",", ".") if value >= 10_000 else str(value)


def write_spoken_numbers(text: str) -> str:
    """Writes spoken numbers of two or more words in digits ("hai nghìn không trăm hai mươi sáu" is 2026) so the
    model only has to lay them out, which a small model does far more reliably than adding them up. Single words
    stay ("một chút", "ba việc"): only the model can tell from context whether they are numbers."""
    tokens = text.split(" ")
    edges = [WORD_EDGES.match(t).groups() for t in tokens]
    words = [w.lower() for _, w, _ in edges]
    out, i = [], 0
    while i < len(tokens):
        if words[i] not in NUMBER_WORDS:
            out.append(tokens[i])
            i += 1
            continue
        # Punctuation ends a number: "hai mươi, ba mươi" is two.
        stop = i + 1
        while stop < len(tokens) and not edges[stop - 1][2] and not edges[stop][0]:
            stop += 1
        parsed = parse_number(words[:stop], i)
        if parsed is None:
            out.append(tokens[i])
            i += 1
        elif parsed[0] is None:  # unclear where the number ends: keep the whole run of number words
            j = i
            while j < stop and words[j] in NUMBER_WORDS:
                j += 1
            out += tokens[i:j]
            i = j
        else:
            value, j = parsed
            out.append(edges[i][0] + write_number(value) + edges[j - 1][2])
            i = j
    return " ".join(out)


# A single spoken digit that is plainly a number next to tháng, phẩy, phần trăm or a unit, and its value.
SINGLE_NUMBERS = {w: v for w, v in DIGITS.items() if w not in ("mốt", "lăm", "nhăm")} | {"mười": 10}
FRACTION_DIGITS = {w: v for w, v in DIGITS.items() if w not in ("mốt", "tư")}  # read one by one after phẩy

# Spoken units of measure written as symbols after a number, as plain() forms. "kí" and "cân" are left to the
# model: alone they are often not units.
UNIT_WORDS = {form: symbol for symbol, forms in UNITS.items() for form in forms if form not in ("ki", "can")}
UNIT_MAX_WORDS = 5  # "ki lô mét trên giờ"
# Digits read one by one this many times or more in a row are a phone number or a code: "năm sáu người" (five or
# six people) is not.
PHONE_MIN_DIGITS = 7

# Words that count what a single spoken digit before them numbers: "hai phòng" is 2 phòng. Two-word counters are
# written with a space. Left out: "bạn" and "con" ("ba bạn" is also "your father"), and "năm" before "tháng" and
# "năm" ("năm tháng" is also "the years"). "một" never counts: it is also "a".
COUNTERS = frozenset("""người lần viên cái chiếc quyển cuốn phòng giờ phút giây ngày tuần tháng năm tuổi việc
    tầng lớp chai ly cốc bát hộp gói bộ đôi trang bài câu điểm muỗng thìa quả trái km kg""".split()) | {"học sinh", "sinh viên", "nhân viên"}
COUNT_DIGITS = {w: v for w, v in SINGLE_NUMBERS.items() if w not in ("một", "không", "tư")}


def number_token(word: str) -> int | None:
    """The value of a word that is a whole number on its own: digits without separators, or a single spoken digit."""
    if regex.fullmatch(r"\d+", word):
        return int(word)
    return SINGLE_NUMBERS.get(word.lower())


def write_measures(text: str) -> str:
    """After write_spoken_numbers: dates as day/month/year (ngày 20 tháng 11 năm 2026 is ngày 20/11/2026, tháng 4
    năm 2025 is tháng 4/2025), decimals with a comma (ba phẩy hai is 3,2), percentages (20 phần trăm is 20%) and
    units of measure as symbols (12 ki lô mét is 12 km), phone numbers (không chín một ... is 091...), and a single
    digit that counts something (hai phòng is 2 phòng). Each one only where the words can mean nothing else, so the model just copies them; punctuation inside a run stops it."""
    tokens = text.split(" ")
    edges = [WORD_EDGES.match(t).groups() for t in tokens]
    words = [w.lower() for _, w, _ in edges]

    def joined(i: int, j: int) -> bool:
        """Whether tokens[i:j] were said as one run: no punctuation between them."""
        return all(not edges[k][2] for k in range(i, j - 1)) and all(not edges[k][0] for k in range(i + 1, j))

    def date(i: int) -> tuple[str, int] | None:
        """D tháng M [năm YYYY] after ngày, mùng or mồng; M năm YYYY after tháng."""
        if i == 0 or words[i - 1] not in ("ngày", "mùng", "mồng", "tháng") or edges[i - 1][2]:
            return None
        first = number_token(edges[i][1])
        if first is None or i + 2 >= len(tokens):
            return None
        if words[i - 1] == "tháng":  # month/year
            year = edges[i + 2][1]
            if words[i + 1] == "năm" and regex.fullmatch(r"\d{4}", year) and 1 <= first <= 12 and joined(i, i + 3):
                return f"{first}/{year}", i + 3
            return None
        month = number_token(edges[i + 2][1])
        if words[i + 1] != "tháng" or month is None or not (1 <= first <= 31 and 1 <= month <= 12) or not joined(i, i + 3):
            return None
        if i + 4 < len(tokens) and words[i + 3] == "năm" and regex.fullmatch(r"\d{4}", edges[i + 4][1]) and joined(i, i + 5):
            return f"{first}/{month}/{edges[i + 4][1]}", i + 5
        return f"{first}/{month}", i + 3

    def amount(i: int) -> tuple[str, int] | None:
        """A number with what makes it plain: X phẩy Y, then phần trăm or a unit; or a number with phần trăm or a
        unit. A digits word alone is left as it is."""
        whole = number_token(edges[i][1])
        if whole is None and written_number(edges[i][1]) is None:
            return None
        value, j = edges[i][1] if whole is None else str(whole), i + 1
        if whole is not None and j + 1 < len(tokens) and words[j] == "phẩy" and joined(i, j + 2):
            if regex.fullmatch(r"\d+", edges[j + 1][1]):
                value, j = f"{value},{edges[j + 1][1]}", j + 2
            elif words[j + 1] in FRACTION_DIGITS:
                k = j + 1
                while k < len(tokens) and words[k] in FRACTION_DIGITS and joined(i, k + 1):
                    k += 1
                value, j = f"{value}," + "".join(str(FRACTION_DIGITS[w]) for w in words[j + 1:k]), k
        if j + 1 < len(tokens) and words[j:j + 2] == ["phần", "trăm"] and joined(i, j + 2):
            return value + "%", j + 2
        for k in range(min(len(tokens), j + UNIT_MAX_WORDS), j, -1):
            symbol = UNIT_WORDS.get(plain("".join(words[j:k]).replace("-", "")))
            # "một mét bảy" is 1.7 m said short: a number after the unit leaves it to the model.
            if symbol is not None and joined(i, k) and (k >= len(tokens) or number_token(edges[k][1]) is None
                                                        or edges[k - 1][2]):
                return value + ("" if symbol == "°C" else " ") + symbol, k
        if "," in value:
            return value, j
        return None

    def count(i: int) -> tuple[str, int] | None:
        """A single spoken digit before a counter (hai phòng is 2 phòng), or after "điểm" (điểm mười is điểm 10)."""
        value = COUNT_DIGITS.get(words[i])
        if value is None:
            return None
        if i > 0 and words[i - 1] == "điểm" and not edges[i - 1][2]:
            return str(value), i + 1
        if i + 1 >= len(tokens) or not joined(i, i + 2):
            return None
        after = words[i + 1]
        if words[i] == "năm" and after in ("tháng", "năm"):
            return None
        # Next to other digits it is a run of digits ("bốn năm sáu"), not a count.
        if (i > 0 and words[i - 1] in DIGITS and joined(i - 1, i + 1)) or (
                i + 2 < len(tokens) and words[i + 2] in DIGITS and joined(i + 1, i + 3)):
            return None
        two = after + " " + words[i + 2] if i + 2 < len(tokens) and joined(i + 1, i + 3) else ""
        if after in COUNTERS or two in COUNTERS:
            return str(value), i + 1
        return None

    def phone(i: int) -> tuple[str, int] | None:
        j = i
        while j < len(tokens) and words[j] in FRACTION_DIGITS and words[j] != "lăm" and joined(i, j + 1):
            j += 1
        if j - i < PHONE_MIN_DIGITS:
            return None
        return "".join(str(DIGITS[w]) for w in words[i:j]), j

    out, i = [], 0
    while i < len(tokens):
        found = date(i) or amount(i) or phone(i) or count(i)
        if found is None:
            out.append(tokens[i])
            i += 1
        else:
            written, j = found
            out.append(edges[i][0] + written + edges[j - 1][2])
            i = j
    return " ".join(out)


def write_numbers(text: str) -> str:
    """The transcript as the model sees it: spoken numbers, dates, decimals, percentages, units and phone numbers
    written the Vietnamese way wherever the words can mean nothing else."""
    return write_measures(write_spoken_numbers(text))


def written_number(word: str) -> int | None:
    """The value of digits the model wrote, with "." or "," as thousands separators: 15, 2026, 1.000."""
    if not regex.fullmatch(r"\d{1,3}([.,]\d{3})*|\d+", word):
        return None
    return int(word.replace(".", "").replace(",", ""))


def spoken_digits(words: list[str]) -> str | None:
    """Digits read one by one, as in a phone number ("không chín một" is "091"), or None."""
    if len(words) < 2 or any(w not in DIGITS or w in ("mốt", "lăm", "nhăm", "tư") for w in words):
        return None
    return "".join(str(DIGITS[w]) for w in words)


def number_matches(words: list[str], written: str) -> bool:
    """Whether the digits the model wrote say the same as the spoken words: a number (2026, 1.000), digits read
    one by one (0912), or a decimal with a comma ("ba phẩy năm" is 3,5)."""
    if not words:
        return False
    if len(words) == 1 and written_number(words[0]) is not None:  # digits written before the model ran
        return written_number(words[0]) == written_number(written)
    if regex.fullmatch(r"\d+", written) and spoken_digits(words) == written:
        return True
    value = written_number(written)
    if value is not None:
        return value == spoken_number(words)
    decimal = regex.fullmatch(r"(\d+),(\d+)", written)
    if decimal is None or "phẩy" not in words:
        return False
    k = words.index("phẩy")
    whole, fraction = words[:k], words[k + 1:]
    fraction_ok = spoken_digits(fraction) == decimal[2] or (
        decimal[2][0] != "0" and spoken_number(fraction) == int(decimal[2]))
    return spoken_number(whole) == int(decimal[1]) and fraction_ok


def split_matches(words: list[str], sep: str, parts: list[str]) -> bool:
    """Whether words are parts[0] sep parts[1] sep ... read aloud, trying every place sep could split them
    ("năm" is both "year" and five, so "tháng năm năm 2024" splits in more than one way)."""
    if len(parts) == 1:
        return number_matches(words, parts[0])
    return any(
        w == sep and number_matches(words[:k], parts[0]) and split_matches(words[k + 1:], sep, parts[1:])
        for k, w in enumerate(words)
    )


def date_matches(words: list[str], written: str) -> bool:
    """A spoken date written the Vietnamese way: day/month/year, day/month or (after "tháng") month/year.
    "hai mươi tháng mười một năm hai nghìn không trăm hai mươi sáu" is 20/11/2026."""
    match = regex.fullmatch(r"(\d{1,2})/(\d{1,2})(?:/(\d{4}))?", written)
    if match is None:
        return False
    if words and words[0] in ("mùng", "mồng"):
        words = words[1:]
    first, second, year = match[1], match[2], match[3]
    if not (1 <= int(first) <= 31 and 1 <= int(second) <= 12):
        return False
    if year is None:
        return split_matches(words, "tháng", [first, second])
    return any(
        w == "năm" and split_matches(words[:k], "tháng", [first, second]) and number_matches(words[k + 1:], year)
        for k, w in enumerate(words)
    )


def month_year_matches(words: list[str], written: str) -> bool:
    """A month and year after "tháng", written month/year: "tư năm hai nghìn không trăm hai mươi lăm" is 4/2025."""
    match = regex.fullmatch(r"(\d{1,2})/(\d{4})", written)
    return match is not None and 1 <= int(match[1]) <= 12 and split_matches(words, "năm", [match[1], match[2]])


def accept_number(old: list[str], new: Unit) -> bool:
    """One written token for spoken words: a number, a percentage, a date or a month and year."""
    if new.trail.startswith("%"):
        return old[-2:] == ["phần", "trăm"] and number_matches(old[:-2], new.word)
    return number_matches(old, new.word) or date_matches(old, new.word) or month_year_matches(old, new.word)


def accept_unit(old: list[str], symbol: str) -> bool:
    """A spoken unit of measure written as its symbol: "ki lô gam" is kg."""
    spoken = plain("".join(old).replace("-", ""))
    return spoken in UNITS.get(symbol, ())


def is_number(u: Unit | None) -> bool:
    """A word that ends a number: digits, or a spoken number word."""
    return u is not None and (u.key in NUMBER_WORDS or written_number(u.word) is not None)


def accept_replacement(old: list[Unit], new: list[Unit], before: Unit | None) -> bool:
    """A misheard word fixed one for one, a spoken number written as digits with the same value (a number,
    percentage, decimal, phone number or date), or a unit of measure written as its symbol after one. before is
    the transcript word before old: a unit alone is trusted only after a number ("kì" is not kg)."""
    if len(new) == 1 and new[0].word in UNITS:
        return is_number(before) and accept_unit([o.key for o in old], new[0].word)
    if len(old) == len(new) and all(sounds_alike(o.key, n.key) for o, n in zip(old, new)):
        return True
    keys = [o.key for o in old]
    if len(new) == 1:
        joined = regex.fullmatch(r"(\d[\d.,]*)(\D+)", new[0].word)  # 5kg, 30°C
        if joined is not None:
            new = [Unit(joined[1]), Unit(joined[2])]
        elif accept_number(keys, new[0]) or (is_number(before) and accept_unit(keys, new[0].word)):
            return True
    if len(new) == 2:
        return any(accept_number(keys[:k], new[0]) and accept_unit(keys[k:], new[1].word) for k in range(1, len(keys)))
    return False


def align(a: list[str], b: list[str]) -> list[tuple[str, int, int, int, int]]:
    """Longest common subsequence of a and b as difflib-style opcodes (tag, i1, i2, j1, j2); a match is taken as
    early as possible, and an unmatched run of a before one of b."""
    n, m = len(a), len(b)
    lcs = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n - 1, -1, -1):
        for j in range(m - 1, -1, -1):
            lcs[i][j] = lcs[i + 1][j + 1] + 1 if a[i] == b[j] else max(lcs[i + 1][j], lcs[i][j + 1])
    ops, i, j = [], 0, 0
    while i < n or j < m:
        if i < n and j < m and a[i] == b[j]:
            i2, j2 = i, j
            while i2 < n and j2 < m and a[i2] == b[j2]:
                i2, j2 = i2 + 1, j2 + 1
            ops.append(("equal", i, i2, j, j2))
            i, j = i2, j2
            continue
        i2, j2 = i, j
        while i2 < n or j2 < m:
            if i2 < n and j2 < m and a[i2] == b[j2]:
                break
            if j2 >= m or (i2 < n and lcs[i2 + 1][j2] >= lcs[i2][j2 + 1]):
                i2 += 1
            else:
                j2 += 1
        tag = "replace" if i2 > i and j2 > j else ("delete" if i2 > i else "insert")
        ops.append((tag, i, i2, j, j2))
        i, j = i2, j2
    return ops


def may_delete(old: list[Unit], start: int, end: int, next_unit: Unit | None) -> list[bool]:
    """For each transcript word in old[start:end] the model left out: whether leaving it out is trusted."""
    run = old[start:end]
    if next_unit is not None and next_unit.bullet and all(u.key in LIST_MARKERS for u in run):
        return [True] * len(run)
    keep = []
    for k, u in enumerate(run):
        i = start + k
        repeated = (i > 0 and old[i - 1].key == u.key) or (i + 1 < len(old) and old[i + 1].key == u.key)
        keep.append(u.key in FILLERS or repeated)
    return keep


def merge(transcript: str, edited: str) -> str:
    """The model's text, with every edit that is not trusted undone (see the module docstring)."""
    old, new = units(transcript), units(edited)
    out: list[Unit] = []
    carry: Unit | None = None  # layout of dropped words, for the next word written

    def emit(u: Unit) -> None:
        nonlocal carry
        if carry is not None:
            u.newline, u.bullet = u.newline or carry.newline, u.bullet or carry.bullet
            u.lead = carry.lead + u.lead
            carry = None
        out.append(u)

    def drop(dropped: list[Unit]) -> None:
        """Keeps the layout and the closing punctuation of words that are left out."""
        nonlocal carry
        for u in dropped:
            if u.newline or u.bullet or u.lead:
                carry = carry or Unit("")
                carry.newline, carry.bullet = carry.newline or u.newline, carry.bullet or u.bullet
                carry.lead += u.lead
            if u.trail and out and not out[-1].trail:
                out[-1].trail = u.trail

    def restore(words: list[Unit], like: list[Unit]) -> list[Unit]:
        """The transcript's words in place of the model's [like]: its layout, its first capital, its last mark."""
        back = [Unit(u.word) for u in words]
        if like and back:
            back[0].newline, back[0].bullet, back[0].lead = like[0].newline, like[0].bullet, like[0].lead
            back[-1].trail = like[-1].trail.lstrip("%°")  # a sign that belonged to the model's digits
            if like[0].word[:1].isupper():
                back[0].word = back[0].word[:1].upper() + back[0].word[1:]
        return back

    kept = 0  # transcript words the model kept or edited in a trusted way
    for tag, i1, i2, j1, j2 in align([u.key for u in old], [u.key for u in new]):
        if tag == "equal":
            kept += i2 - i1
            for k, u in enumerate(new[j1:j2]):
                # A % is trusted only for "phần trăm" the model left out right after the number.
                said = i1 + k + 1 == i2 and [o.key for o in old[i2:i2 + 2]] == ["phần", "trăm"]
                if "%" in u.trail and "%" not in old[i1 + k].trail and not said:
                    u.trail = u.trail.replace("%", "")
                emit(u)
        elif tag == "insert":
            drop(new[j1:j2])
        elif tag == "replace" and accept_replacement(old[i1:i2], new[j1:j2], old[i1 - 1] if i1 > 0 else None):
            kept += i2 - i1
            for u in new[j1:j2]:
                emit(u)
        elif tag == "replace":
            for u in restore(old[i1:i2], new[j1:j2]):
                emit(u)
        elif [u.key for u in old[i1:i2]] == ["phần", "trăm"] and out and "%" in out[-1].trail:
            drop(old[i1:i2])  # written as the % after the number
        else:  # delete
            next_unit = new[j1] if j1 < len(new) else None
            for u, ok in zip(old[i1:i2], may_delete(old, i1, i2, next_unit)):
                if not ok:
                    restored = Unit(u.word)
                    # The word belongs before the mark that ends the previous word's sentence or clause.
                    if out and out[-1].trail:
                        restored.trail, out[-1].trail = out[-1].trail, ""
                    emit(restored)
    if carry is not None and out and carry.lead:
        out[-1].trail += carry.lead
    if kept * 2 < len(old):
        return ""  # the model answered or rewrote instead of correcting: the caller keeps the transcript
    out = drop_hesitations(out)
    capitalize_sentence_starts(out)
    mark_questions(out)
    return render(out)


# Hesitation sounds the merge leaves out even when the model kept them. Not "ừ" (yes), "à" (also a question word)
# or "ơ" (surprise).
HESITATIONS = frozenset("ờ ừm ưm hừm ừhm um uh".split())


def drop_hesitations(items: list[Unit]) -> list[Unit]:
    """The units without hesitation sounds: their layout goes to the next word, a closing mark to the word before."""
    out: list[Unit] = []
    carry: Unit | None = None
    for u in items:
        if u.key in HESITATIONS:
            carry = carry or Unit("")
            carry.newline, carry.bullet = carry.newline or u.newline, carry.bullet or u.bullet
            carry.lead += u.lead
            if u.trail.rstrip(",;") and out and out[-1].trail in ("", ","):
                out[-1].trail = u.trail.rstrip(",;")
            elif out and out[-1].trail == "," and u.trail in ("", ","):
                out[-1].trail = ""  # "đi, ờ, chợ" is "đi chợ": the commas set off the sound
            continue
        if carry is not None:
            u.newline, u.bullet, u.lead = u.newline or carry.newline, u.bullet or carry.bullet, carry.lead + u.lead
            carry = None
        out.append(u)
    return out if out else items


# How a yes/no question ends: "có ... không", "... được không", "... phải không", "đã ... chưa". "là không" and
# "cũng không" end statements.
QUESTION_OPENERS = {"không": ("có",), "chưa": ("đã", "có")}
QUESTION_ENDINGS = (("được", "không"), ("phải", "không"))
NOT_QUESTIONS = frozenset({"là", "cũng", "đều", "vẫn", "chẳng"})


def mark_questions(items: list[Unit]) -> None:
    """A sentence the model ended with a full stop that is a yes/no question gets a question mark: small models
    often miss "Bạn có biết mấy giờ rồi không?"."""
    start = 0
    for i, u in enumerate(items):
        if i > start and (u.newline or u.bullet):
            start = i
        if not u.trail.endswith((".", "?", "!")):
            continue
        sentence = [v.key for v in items[start:i + 1]]
        start = i + 1
        if not u.trail.endswith(".") or u.trail.endswith("..") or len(sentence) < 3 or sentence[-2] in NOT_QUESTIONS:
            continue
        if tuple(sentence[-2:]) in QUESTION_ENDINGS or any(w in sentence[:-1] for w in QUESTION_OPENERS.get(sentence[-1], ())):
            u.trail = u.trail[:-1] + "?"


def capitalize_sentence_starts(items: list[Unit]) -> None:
    """Restored words keep the transcript's lower case; one that now starts a sentence, line or bullet gets a capital."""
    for i, u in enumerate(items):
        starts = i == 0 or u.newline or u.bullet or items[i - 1].trail.rstrip("\"')»”").endswith((".", "?", "!"))
        if starts and u.word[:1].islower() and not u.lead:
            u.word = u.word[:1].upper() + u.word[1:]


class Refiner:
    def __init__(self, repo: Path, threads: int = 4):
        self.tok = Tokenizer(repo)
        options = ort.SessionOptions()
        options.intra_op_num_threads = threads
        self.session = ort.InferenceSession(str(repo / "onnx" / "model.onnx"), options, providers=["CPUExecutionProvider"])
        names = [i.name for i in self.session.get_inputs()]
        self.past_names = [n for n in names if n.startswith("past_key_values.")]
        self.uses_position_ids = "position_ids" in names
        config = json.loads((repo / "onnx" / "genai_config.json").read_text(encoding="utf-8"))["model"]["decoder"]
        self.kv_heads, self.head_size = config["num_key_value_heads"], config["head_size"]
        self.stop = {self.tok.special(IM_END), self.tok.special(END_OF_TEXT)}
        self.prefix = prompt_prefix(self.tok)
        empty = np.zeros((1, self.kv_heads, 0, self.head_size), np.float32)
        _, self.prefix_cache = self.forward(self.prefix, {n: empty for n in self.past_names}, 0)

    def forward(self, ids: list[int], cache: dict[str, np.ndarray], past_len: int) -> tuple[np.ndarray, dict[str, np.ndarray]]:
        """Runs ids after past_len cached tokens; returns the last position's logits and the grown cache."""
        total = past_len + len(ids)
        feed = {
            "input_ids": np.array([ids], np.int64),
            "attention_mask": np.ones((1, total), np.int64),
            **cache,
        }
        if self.uses_position_ids:
            feed["position_ids"] = np.arange(past_len, total, dtype=np.int64)[None, :]
        outputs = self.session.run(None, feed)
        out_names = [o.name for o in self.session.get_outputs()]
        result = dict(zip(out_names, outputs))
        new_cache = {n: result[n.replace("past_key_values.", "present.")] for n in self.past_names}
        return result["logits"][0, -1], new_cache

    def generate(self, transcript: str) -> list[int]:
        suffix = prompt_suffix(self.tok, transcript)
        logits, cache = self.forward(suffix, self.prefix_cache, len(self.prefix))
        length = len(self.prefix) + len(suffix)
        out: list[int] = []
        for _ in range(max_new_tokens(len(suffix))):
            nxt = int(np.argmax(logits))
            if nxt in self.stop:
                break
            out.append(nxt)
            logits, cache = self.forward([nxt], cache, length)
            length += 1
        return out

    def refine(self, transcript: str) -> str:
        text = write_numbers(prepare(transcript))
        if not text:
            return transcript
        edited = self.tok.decode(self.generate(text)).strip()
        merged = merge(text, edited) if edited else ""
        return merged if merged.strip() else text


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("repo", type=Path)
    parser.add_argument("text", nargs="+")
    args = parser.parse_args()
    refiner = Refiner(args.repo)
    for text in args.text:
        print(refiner.refine(text))
        sys.stdout.flush()


if __name__ == "__main__":
    main()