feat: add expressive offline Russian book TTS

tools/tts/compare_booktts_bert.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Compare teacher, FP32 student, and INT8 student embeddings on real text."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+
+from vosk_booktts_experiment import WordPieceTokenizer, load_corpus, normalize_text
+
+
+def cosine_rows(left: np.ndarray, right: np.ndarray) -> np.ndarray:
+    numerator = np.sum(left * right, axis=-1)
+    denominator = np.linalg.norm(left, axis=-1) * np.linalg.norm(right, axis=-1)
+    return numerator / np.maximum(denominator, 1e-12)
+
+
+def compare(args: argparse.Namespace) -> None:
+    tokenizer = WordPieceTokenizer(args.vocab)
+    token_map = np.load(args.token_map)
+    compact_unk = int(token_map[tokenizer.vocabulary["[UNK]"]])
+    sessions = {
+        "teacher": ort.InferenceSession(str(args.teacher), providers=["CPUExecutionProvider"]),
+        "student_fp32": ort.InferenceSession(str(args.student_fp32), providers=["CPUExecutionProvider"]),
+        "student_int8": ort.InferenceSession(str(args.student_int8), providers=["CPUExecutionProvider"]),
+    }
+    timings = {name: 0.0 for name in sessions}
+    teacher_cosines: list[np.ndarray] = []
+    quantized_cosines: list[np.ndarray] = []
+    teacher_squared_error = 0.0
+    quantized_squared_error = 0.0
+    elements = 0
+    tokens = 0
+    unseen_tokens = 0
+    phrases = load_corpus(args.corpus)[: args.limit or None]
+    for text in phrases:
+        ids, _ = tokenizer.encode(normalize_text(text))
+        shape = (1, ids.size)
+        feeds = {
+            "input_ids": ids.reshape(shape),
+            "attention_mask": np.ones(shape, dtype=np.int64),
+            "token_type_ids": np.zeros(shape, dtype=np.int64),
+        }
+        outputs: dict[str, np.ndarray] = {}
+        for name, session in sessions.items():
+            started = time.perf_counter()
+            outputs[name] = session.run(None, feeds)[0].astype(np.float32, copy=False)
+            timings[name] += time.perf_counter() - started
+        teacher = outputs["teacher"]
+        fp32 = outputs["student_fp32"]
+        int8 = outputs["student_int8"]
+        if teacher.shape != fp32.shape or fp32.shape != int8.shape:
+            raise ValueError(f"Output shape mismatch for {text!r}: {teacher.shape}, {fp32.shape}, {int8.shape}")
+        teacher_cosines.append(cosine_rows(teacher, fp32))
+        quantized_cosines.append(cosine_rows(fp32, int8))
+        teacher_squared_error += float(np.sum((teacher - fp32) ** 2))
+        quantized_squared_error += float(np.sum((fp32 - int8) ** 2))
+        elements += teacher.size
+        tokens += ids.size
+        unseen_tokens += int(np.sum((token_map[ids] == compact_unk) & (ids != tokenizer.vocabulary["[UNK]"])))
+    teacher_cosine = np.concatenate(teacher_cosines)
+    quantized_cosine = np.concatenate(quantized_cosines)
+    report = {
+        "phrases": len(phrases),
+        "tokens": tokens,
+        "unseen_tokens": unseen_tokens,
+        "unseen_percent": 100.0 * unseen_tokens / tokens,
+        "teacher_vs_student_fp32": {
+            "mean_token_cosine": float(np.mean(teacher_cosine)),
+            "p05_token_cosine": float(np.percentile(teacher_cosine, 5)),
+            "mse": teacher_squared_error / elements,
+        },
+        "student_fp32_vs_int8": {
+            "mean_token_cosine": float(np.mean(quantized_cosine)),
+            "p05_token_cosine": float(np.percentile(quantized_cosine, 5)),
+            "mse": quantized_squared_error / elements,
+        },
+        "total_inference_seconds": timings,
+        "milliseconds_per_phrase": {name: seconds * 1000.0 / len(phrases) for name, seconds in timings.items()},
+    }
+    if not all(math.isfinite(value) for value in (teacher_squared_error, quantized_squared_error)):
+        raise ValueError("Non-finite comparison result")
+    rendered = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
+    if args.output:
+        if args.output.exists():
+            raise FileExistsError(f"Refusing to overwrite {args.output}")
+        args.output.write_text(rendered, encoding="utf-8")
+    print(rendered, end="")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--teacher", type=Path, required=True)
+    parser.add_argument("--student-fp32", type=Path, required=True)
+    parser.add_argument("--student-int8", type=Path, required=True)
+    parser.add_argument("--vocab", type=Path, required=True)
+    parser.add_argument("--token-map", type=Path, required=True)
+    parser.add_argument("--corpus", type=Path, required=True)
+    parser.add_argument("--output", type=Path)
+    parser.add_argument("--limit", type=int, default=0)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    sys.exit(compare(parse_args()) or 0)