569 lines
26 KiB
Python
569 lines
26 KiB
Python
#!/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"(?<![^\W_])([+-]?)(\d{1,18})(?:[,.](\d+))?(?![^\W_])", re.UNICODE)
|
||
YEAR_IN_PREPOSITIONAL_PATTERN = re.compile(r"(?<![^\W_])(\d{4})\s+(году|г\.)", re.IGNORECASE | re.UNICODE)
|
||
|
||
SOFT_LETTERS = set("яёюиье")
|
||
SYLLABLE_STARTS = set("#ъьаёуюэеиы-")
|
||
SOFT_HARD_CONSONANTS = {
|
||
"б": "b", "в": "v", "г": "g", "д": "d", "з": "z", "к": "k", "л": "l",
|
||
"м": "m", "н": "n", "п": "p", "р": "r", "с": "s", "т": "t", "ф": "f", "х": "h",
|
||
}
|
||
OTHER_CONSONANTS = {"ж": "zh", "ц": "c", "ч": "ch", "ш": "sh", "щ": "sch", "й": "j"}
|
||
VOWELS = {"а": "a", "я": "a", "у": "u", "ю": "u", "о": "o", "ё": "o", "э": "e", "е": "e", "и": "i", "ы": "y"}
|
||
|
||
|
||
def scale_form(value: int, forms: tuple[str, str, str]) -> 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("<i2").tobytes()
|
||
with wave.open(str(path), "wb") as output:
|
||
output.setnchannels(1); output.setsampwidth(2); output.setframerate(SAMPLE_RATE); output.writeframes(pcm)
|
||
|
||
|
||
def command_compare(args: argparse.Namespace) -> 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())
|