#!/usr/bin/env python3 """Prepare real Vosk feeds, quantize acoustic ONNX, and render A/B WAVs. The frontend mirrors VoskSpeechEngine.kt: number normalization, WordPiece positions, dictionary lookup and all five phone feature channels. """ from __future__ import annotations import argparse import bisect import json import math import re import sys import time import wave from collections import OrderedDict from dataclasses import dataclass from pathlib import Path from typing import Iterator import numpy as np SAMPLE_RATE = 22_050 SPEAKER_ID = 4 BERT_DIMENSIONS = 768 MAX_WORD_PIECE_CHARS = 100 CACHE_SIZE = 4_096 MULTISTREAM_PUNCTUATION = set("!(),-.:;?") BERT_EXCLUDED_PUNCTUATION = {"-", ",", ".", "?", "!", ";", ":", '"'} PHONEMES = [ "_", "^", "$", " ", "!", "'", "(", ")", ",", "-", ".", "...", ":", ";", "?", "a0", "a1", "b", "bj", "c", "ch", "d", "dj", "e0", "e1", "f", "fj", "g", "gj", "h", "hj", "i0", "i1", "j", "k", "kj", "l", "lj", "m", "mj", "n", "nj", "o0", "o1", "p", "pj", "r", "rj", "s", "sch", "sh", "sj", "t", "tj", "u0", "u1", "v", "vj", "y0", "y1", "z", "zh", "zj", ] PHONEME_IDS = {phone: index for index, phone in enumerate(PHONEMES)} DIGITS = ["ноль", "один", "два", "три", "четыре", "пять", "шесть", "семь", "восемь", "девять"] TEENS = [ "десять", "одиннадцать", "двенадцать", "тринадцать", "четырнадцать", "пятнадцать", "шестнадцать", "семнадцать", "восемнадцать", "девятнадцать", ] TENS = ["", "десять", "двадцать", "тридцать", "сорок", "пятьдесят", "шестьдесят", "семьдесят", "восемьдесят", "девяносто"] HUNDREDS = ["", "сто", "двести", "триста", "четыреста", "пятьсот", "шестьсот", "семьсот", "восемьсот", "девятьсот"] SCALES = [ None, ("тысяча", "тысячи", "тысяч"), ("миллион", "миллиона", "миллионов"), ("миллиард", "миллиарда", "миллиардов"), ("триллион", "триллиона", "триллионов"), ("квадриллион", "квадриллиона", "квадриллионов"), ] NUMBER_PATTERN = re.compile(r"(? str: if 11 <= value % 100 <= 14: return forms[2] return forms[0] if value % 10 == 1 else forms[1] if value % 10 in (2, 3, 4) else forms[2] def under_thousand_to_words(value: int, feminine: bool) -> str: words: list[str] = [] if value // 100: words.append(HUNDREDS[value // 100]) tail = value % 100 if 10 <= tail <= 19: words.append(TEENS[tail - 10]) else: if tail // 10: words.append(TENS[tail // 10]) units = tail % 10 if units: words.append("одна" if feminine and units == 1 else "две" if feminine and units == 2 else DIGITS[units]) return " ".join(words) def integer_to_words(value: int) -> str: if value == 0: return DIGITS[0] groups: list[int] = [] while value: groups.append(value % 1_000) value //= 1_000 result: list[str] = [] for group_index in range(len(groups) - 1, -1, -1): group = groups[group_index] if not group: continue words = under_thousand_to_words(group, group_index == 1) forms = SCALES[group_index] if group_index < len(SCALES) else None result.append(words if forms is None else f"{words} {scale_form(group, forms)}") return " ".join(result) def normalize_numbers(text: str) -> str: prepared = re.sub(r"(?<=\d)[\u00a0\u202f](?=\d{3}(?:\D|$))", "", text) prepared = re.sub(r"№\s*(?=\d)", "номер ", prepared) def replace(match: re.Match[str]) -> str: sign = "минус " if match.group(1) == "-" else "плюс " if match.group(1) == "+" else "" integer_digits = match.group(2) if len(integer_digits) > 1 and integer_digits.startswith("0"): integer = " ".join(DIGITS[int(char)] for char in integer_digits) else: integer = integer_to_words(int(integer_digits)) fraction = match.group(3) return f"{sign}{integer}" if not fraction else f"{sign}{integer} запятая {' '.join(DIGITS[int(char)] for char in fraction)}" prepared = YEAR_IN_PREPOSITIONAL_PATTERN.sub( lambda match: f"{year_in_prepositional_words(int(match.group(1)))} {match.group(2)}", prepared, ) return NUMBER_PATTERN.sub(replace, prepared) ORDINAL_UNITS_PREPOSITIONAL = ("", "первом", "втором", "третьем", "четвёртом", "пятом", "шестом", "седьмом", "восьмом", "девятом") ORDINAL_TEENS_PREPOSITIONAL = ("десятом", "одиннадцатом", "двенадцатом", "тринадцатом", "четырнадцатом", "пятнадцатом", "шестнадцатом", "семнадцатом", "восемнадцатом", "девятнадцатом") ORDINAL_TENS_PREPOSITIONAL = ("", "десятом", "двадцатом", "тридцатом", "сороковом", "пятидесятом", "шестидесятом", "семидесятом", "восьмидесятом", "девяностом") ORDINAL_HUNDREDS_PREPOSITIONAL = ("", "сотом", "двухсотом", "трёхсотом", "четырёхсотом", "пятисотом", "шестисотом", "семисотом", "восьмисотом", "девятисотом") EXACT_THOUSANDTH_PREPOSITIONAL = ("", "тысячном", "двухтысячном", "трёхтысячном", "четырёхтысячном", "пятитысячном", "шеститысячном", "семитысячном", "восьмитысячном", "девятитысячном") def ordinal_under_thousand_prepositional(value: int) -> str: last_two = value % 100 units = value % 10 if units and not 11 <= last_two <= 19: return f"{integer_to_words(value - units) if value > units else ''} {ORDINAL_UNITS_PREPOSITIONAL[units]}".strip() if 11 <= last_two <= 19: return f"{integer_to_words(value - last_two) if value > last_two else ''} {ORDINAL_TEENS_PREPOSITIONAL[last_two - 10]}".strip() if last_two >= 20: return f"{integer_to_words(value - last_two) if value > last_two else ''} {ORDINAL_TENS_PREPOSITIONAL[last_two // 10]}".strip() return ORDINAL_HUNDREDS_PREPOSITIONAL[value // 100] def year_in_prepositional_words(value: int) -> str: if not 1000 <= value <= 9999: raise ValueError(value) if value % 1000 == 0: return EXACT_THOUSANDTH_PREPOSITIONAL[value // 1000] thousands = value // 1000 thousands_words = under_thousand_to_words(thousands, True) thousands_prefix = "" if thousands_words == "одна" else thousands_words + " " return f"{thousands_prefix}{scale_form(thousands, SCALES[1])} {ordinal_under_thousand_prepositional(value % 1000)}" def normalize_text(text: str) -> str: stressed = re.sub(r"([А-Яа-яЁё])\u0301", lambda match: "+" + match.group(1), text) translated = normalize_numbers(stressed).lower().replace("—", "-").replace("–", "-").replace("…", "...") translated = translated.translate(str.maketrans({"«": '"', "»": '"', "“": '"', "”": '"'})) return re.sub(r"\s+", " ", translated).strip() class WordPieceTokenizer: def __init__(self, path: Path): self.tokens = path.read_text(encoding="utf-8").splitlines() self.vocabulary = {token: index for index, token in enumerate(self.tokens)} for special in ("[CLS]", "[SEP]", "[UNK]"): if special not in self.vocabulary: raise ValueError(f"Missing {special} in {path}") def encode(self, text: str) -> tuple[np.ndarray, list[int]]: ids = [self.vocabulary["[CLS]"]] positions = [0] for token in self._basic_tokens(text): punctuation = token in BERT_EXCLUDED_PUNCTUATION pieces = [self.vocabulary.get(token, self.vocabulary["[UNK]"])] if punctuation else self._word_pieces(token) if not punctuation: positions.append(len(ids)) ids.extend(pieces) ids.append(self.vocabulary["[SEP]"]) positions.append(len(ids) - 1) return np.asarray(ids, dtype=np.int64), positions @staticmethod def _basic_tokens(text: str) -> list[str]: result: list[str] = [] word: list[str] = [] for char in text: if char.isspace(): if word: result.append("".join(word)); word.clear() elif char.isalnum(): word.append(char.lower()) elif char == "+": continue else: if word: result.append("".join(word)); word.clear() result.append(char) if word: result.append("".join(word)) return result def _word_pieces(self, token: str) -> list[int]: if len(token) > MAX_WORD_PIECE_CHARS: return [self.vocabulary["[UNK]"]] result: list[int] = [] start = 0 while start < len(token): found: int | None = None found_end = start for end in range(len(token), start, -1): piece = token[start:end] if start == 0 else f"##{token[start:end]}" if piece in self.vocabulary: found = self.vocabulary[piece]; found_end = end; break if found is None: return [self.vocabulary["[UNK]"]] result.append(found); start = found_end return result class PronunciationDictionary: def __init__(self, dictionary_path: Path, index_path: Path): self.source = dictionary_path.open("rb") entries = [] for line in index_path.read_text(encoding="utf-8").splitlines(): word, offset = line.rsplit("\t", 1) entries.append((word, int(offset))) self.words = [entry[0] for entry in entries] self.offsets = [entry[1] for entry in entries] self.length = dictionary_path.stat().st_size self.cache: OrderedDict[str, str | None] = OrderedDict() def close(self) -> None: self.source.close() def find(self, word: str) -> str | None: if word in self.cache: value = self.cache.pop(word); self.cache[word] = value; return value index = max(0, bisect.bisect_right(self.words, word) - 1) end = self.offsets[index + 1] if index + 1 < len(self.offsets) else self.length self.source.seek(self.offsets[index]) value: str | None = None while self.source.tell() < end: raw = self.source.readline() if not raw: break entry_word, _, pronunciation = raw.decode("utf-8").rstrip("\r\n").partition("\t") if entry_word < word: continue if entry_word == word: value = pronunciation.strip() or None break self.cache[word] = value if len(self.cache) > CACHE_SIZE: self.cache.popitem(last=False) return value def russian_g2p(word: str) -> str: marked: list[tuple[str, int]] = [] stress = 0 for char in f"#{word}#": if char == "+": stress = 1 else: marked.append((char, stress)); stress = 0 phones: list[str] = [] for index, (char, accent) in enumerate(marked): previous = marked[index - 1][0] if index else "#" following = marked[index + 1][0] if index + 1 < len(marked) else None if char in SOFT_HARD_CONSONANTS: phones.append(SOFT_HARD_CONSONANTS[char] + ("j" if following in SOFT_LETTERS else "")) elif char in OTHER_CONSONANTS: phones.append(OTHER_CONSONANTS[char]) elif char in VOWELS: if previous in SYLLABLE_STARTS and char in set("яюеё"): phones.append("j") phones.append(f"{VOWELS[char]}{accent}") return " ".join(phones) @dataclass class RawPhone: phone: str punctuation: list[str] in_quote: int bert_word_index: int def build_multistream_phones(text: str, dictionary: PronunciationDictionary) -> tuple[np.ndarray, list[int]]: raw = [RawPhone("^", [], 0, 0)] word: list[str] = [] pending: list[str] = [] in_quote = 0 bert_word_index = 1 def flush_word() -> None: nonlocal bert_word_index if not word: return value = "".join(word) pronunciation = dictionary.find(value) or russian_g2p(value) raw.extend(RawPhone(phone, [], in_quote, bert_word_index) for phone in pronunciation.split() if phone) word.clear(); bert_word_index += 1 def append_space() -> None: raw.append(RawPhone(" ", list(pending), in_quote, bert_word_index)); pending.clear() index = 0 while index < len(text): if text.startswith("...", index): flush_word(); pending.append("..."); index += 3; continue char = text[index] if char in {'"', "«", "»", "“", "”"}: flush_word(); in_quote = 1 - in_quote elif char.isspace(): flush_word(); append_space() elif char == "-" and index > 0 and index + 1 < len(text) and text[index - 1].isalnum() and text[index + 1].isalnum(): flush_word() elif char in MULTISTREAM_PUNCTUATION: flush_word(); pending.append(char) elif char.isalnum() or char == "+": word.append(char) index += 1 flush_word(); append_space(); raw.append(RawPhone("$", [], 0, bert_word_index)) last_punctuation = " " last_sentence_punctuation = " " reversed_features: list[list[int]] = [] reversed_positions: list[int] = [] for phone in reversed(raw): for candidate in ("...", ".", "!", "?", "-"): if candidate in phone.punctuation: last_sentence_punctuation = candidate; break if phone.punctuation: last_punctuation = phone.punctuation[0] current = phone.punctuation[0] if phone.punctuation else "_" values = [PHONEME_IDS[phone.phone], PHONEME_IDS[current], phone.in_quote, PHONEME_IDS[last_punctuation], PHONEME_IDS[last_sentence_punctuation]] reversed_features.append(values); reversed_positions.append(phone.bert_word_index) features = np.asarray(list(reversed(reversed_features)), dtype=np.int64).T[np.newaxis, :, :] return features, list(reversed(reversed_positions)) class Frontend: def __init__(self, assets: Path, bert_model: Path | None = None): import onnxruntime as ort self.tokenizer = WordPieceTokenizer(assets / "vocab.txt") self.dictionary = PronunciationDictionary(assets / "dict.tsv", assets / "dictionary.index") self.bert = ort.InferenceSession(str(bert_model or assets / "bert.int8.onnx"), providers=["CPUExecutionProvider"]) def close(self) -> None: self.dictionary.close() def prepare(self, text: str, articulation: float = 1.0) -> dict[str, np.ndarray]: normalized = normalize_text(text) ids, embedding_positions = self.tokenizer.encode(normalized) shape = (1, ids.size) token_embeddings = self.bert.run(None, { "input_ids": ids.reshape(shape), "attention_mask": np.ones(shape, dtype=np.int64), "token_type_ids": np.zeros(shape, dtype=np.int64), })[0] selected = token_embeddings[np.asarray(embedding_positions, dtype=np.int64)] features, phone_positions = build_multistream_phones(normalized, self.dictionary) phone_embeddings = selected[np.clip(np.asarray(phone_positions), 0, len(selected) - 1)] bert = phone_embeddings.T[np.newaxis, :, :].astype(np.float32, copy=False) time_steps = features.shape[2] if bert.shape != (1, BERT_DIMENSIONS, time_steps): raise ValueError(f"Unexpected BERT shape {bert.shape} for {text!r}") return { "input": features, "input_lengths": np.asarray([time_steps], dtype=np.int64), "scales": np.asarray([0.8, articulation, 0.8], dtype=np.float32), "sid": np.asarray([SPEAKER_ID], dtype=np.int64), "bert": bert, } def load_corpus(path: Path) -> list[str]: lines = [line.strip() for line in path.read_text(encoding="utf-8-sig").splitlines()] result = [line for line in lines if line and not line.startswith("#")] if not result: raise ValueError(f"No phrases in {path}") return result def command_prepare(args: argparse.Namespace) -> None: args.output.mkdir(parents=True, exist_ok=False) frontend = Frontend(args.assets, args.bert_model) manifest = [] try: for index, text in enumerate(load_corpus(args.corpus)): feeds = frontend.prepare(text, args.articulation) name = f"sample-{index:04d}.npz" np.savez_compressed(args.output / name, **feeds) manifest.append({"file": name, "text": text, "time_steps": int(feeds["input_lengths"][0])}) print(f"prepared {index + 1}: phones={manifest[-1]['time_steps']} text={text}") finally: frontend.close() (args.output / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print(f"samples={len(manifest)} output={args.output}") class NpzCalibrationReader: def __init__(self, directory: Path): self.files = sorted(directory.glob("sample-*.npz")) self.iterator: Iterator[Path] | None = None if not self.files: raise ValueError(f"No sample-*.npz in {directory}") def get_next(self) -> dict[str, np.ndarray] | None: if self.iterator is None: self.iterator = iter(self.files) try: path = next(self.iterator) except StopIteration: return None with np.load(path) as data: return {name: data[name] for name in data.files} def rewind(self) -> None: self.iterator = iter(self.files) def command_quantize(args: argparse.Namespace) -> None: from onnxruntime.quantization import QuantFormat, QuantType, quantize_static if args.output.exists(): raise FileExistsError(f"Refusing to overwrite {args.output}") started = time.time() quantize_static( str(args.model), str(args.output), NpzCalibrationReader(args.calibration), quant_format=QuantFormat.QDQ, op_types_to_quantize=["Conv"], per_channel=True, activation_type=QuantType.QInt8, weight_type=QuantType.QInt8, ) source_size = args.model.stat().st_size target_size = args.output.stat().st_size print(json.dumps({ "source_bytes": source_size, "target_bytes": target_size, "saving_bytes": source_size - target_size, "saving_percent": round((source_size - target_size) * 100 / source_size, 4), "seconds": round(time.time() - started, 3), }, indent=2)) def write_wav(path: Path, samples: np.ndarray) -> None: pcm = (np.clip(samples.reshape(-1), -1.0, 1.0) * 32767.0).astype(" None: import onnxruntime as ort args.output.mkdir(parents=True, exist_ok=False) options = ort.SessionOptions(); options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL baseline = ort.InferenceSession(str(args.baseline), sess_options=options, providers=["CPUExecutionProvider"]) candidate = ort.InferenceSession(str(args.candidate), sess_options=options, providers=["CPUExecutionProvider"]) manifest = json.loads((args.calibration / "manifest.json").read_text(encoding="utf-8")) report = [] for item in manifest[: args.limit]: with np.load(args.calibration / item["file"]) as data: feeds = {name: data[name] for name in data.files} feeds["scales"] = np.asarray([0.0, float(feeds["scales"][1]), 0.0], dtype=np.float32) base_wav, base_length = baseline.run(None, feeds) test_wav, test_length = candidate.run(None, feeds) common = min(base_wav.size, test_wav.size) difference = base_wav.reshape(-1)[:common] - test_wav.reshape(-1)[:common] entry = { "text": item["text"], "baseline_samples": int(base_wav.size), "candidate_samples": int(test_wav.size), "baseline_length": np.asarray(base_length).tolist(), "candidate_length": np.asarray(test_length).tolist(), "mae": float(np.mean(np.abs(difference))), "rmse": float(math.sqrt(np.mean(difference * difference))), } report.append(entry) number = len(report) write_wav(args.output / f"{number:02d}-baseline.wav", base_wav) write_wav(args.output / f"{number:02d}-candidate.wav", test_wav) print(f"compared {number}: mae={entry['mae']:.6f} rmse={entry['rmse']:.6f}") (args.output / "report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") def command_compare_feeds(args: argparse.Namespace) -> None: import onnxruntime as ort args.output.mkdir(parents=True, exist_ok=args.resume) options = ort.SessionOptions(); options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL acoustic = ort.InferenceSession(str(args.model), sess_options=options, providers=["CPUExecutionProvider"]) baseline_manifest = json.loads((args.baseline_calibration / "manifest.json").read_text(encoding="utf-8")) candidate_manifest = json.loads((args.candidate_calibration / "manifest.json").read_text(encoding="utf-8")) if [item["text"] for item in baseline_manifest] != [item["text"] for item in candidate_manifest]: raise ValueError("Baseline and candidate manifests contain different texts") report = [] for baseline_item, candidate_item in zip(baseline_manifest[: args.limit or None], candidate_manifest[: args.limit or None]): with np.load(args.baseline_calibration / baseline_item["file"]) as data: baseline_feeds = {name: data[name] for name in data.files} with np.load(args.candidate_calibration / candidate_item["file"]) as data: candidate_feeds = {name: data[name] for name in data.files} baseline_feeds["scales"] = np.asarray([0.0, float(baseline_feeds["scales"][1]), 0.0], dtype=np.float32) candidate_feeds["scales"] = np.asarray([0.0, float(candidate_feeds["scales"][1]), 0.0], dtype=np.float32) baseline_wav, baseline_length = acoustic.run(None, baseline_feeds) candidate_wav, candidate_length = acoustic.run(None, candidate_feeds) common = min(baseline_wav.size, candidate_wav.size) difference = baseline_wav.reshape(-1)[:common] - candidate_wav.reshape(-1)[:common] entry = { "text": baseline_item["text"], "baseline_samples": int(baseline_wav.size), "candidate_samples": int(candidate_wav.size), "baseline_length": np.asarray(baseline_length).tolist(), "candidate_length": np.asarray(candidate_length).tolist(), "mae": float(np.mean(np.abs(difference))), "rmse": float(math.sqrt(np.mean(difference * difference))), } report.append(entry) number = len(report) baseline_path = args.output / f"{number:03d}-baseline.wav" candidate_path = args.output / f"{number:03d}-student-int8.wav" if not baseline_path.exists(): write_wav(baseline_path, baseline_wav) if not candidate_path.exists(): write_wav(candidate_path, candidate_wav) temporary_report = args.output / "report.json.tmp" temporary_report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") temporary_report.replace(args.output / "report.json") print(f"rendered {number}: mae={entry['mae']:.6f} rmse={entry['rmse']:.6f}", flush=True) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() commands = parser.add_subparsers(dest="command", required=True) prepare = commands.add_parser("prepare") prepare.add_argument("--assets", type=Path, required=True) prepare.add_argument("--corpus", type=Path, required=True) prepare.add_argument("--output", type=Path, required=True) prepare.add_argument("--articulation", type=float, default=1.0) prepare.add_argument("--bert-model", type=Path) prepare.set_defaults(handler=command_prepare) quantize = commands.add_parser("quantize") quantize.add_argument("--model", type=Path, required=True) quantize.add_argument("--calibration", type=Path, required=True) quantize.add_argument("--output", type=Path, required=True) quantize.set_defaults(handler=command_quantize) compare = commands.add_parser("compare") compare.add_argument("--baseline", type=Path, required=True) compare.add_argument("--candidate", type=Path, required=True) compare.add_argument("--calibration", type=Path, required=True) compare.add_argument("--output", type=Path, required=True) compare.add_argument("--limit", type=int, default=3) compare.set_defaults(handler=command_compare) compare_feeds = commands.add_parser("compare-feeds") compare_feeds.add_argument("--model", type=Path, required=True) compare_feeds.add_argument("--baseline-calibration", type=Path, required=True) compare_feeds.add_argument("--candidate-calibration", type=Path, required=True) compare_feeds.add_argument("--output", type=Path, required=True) compare_feeds.add_argument("--limit", type=int, default=0) compare_feeds.add_argument("--resume", action="store_true") compare_feeds.set_defaults(handler=command_compare_feeds) return parser.parse_args() def main() -> int: args = parse_args() args.handler(args) return 0 if __name__ == "__main__": sys.exit(main())