Add int8 dynamic quantization support (#147)
# Add int8 dynamic quantization Closes #7 ## Summary This PR implements mixed-precision int8 dynamic quantization for pocket-tts, as requested in issue #7. It adds a `quantize=True` flag to `TTSModel.load_model()` that reduces runtime memory usage by ~48% and improves inference speed by ~27% on x86 and up to 23% on ARM, with no measurable impact on audio quality.
This commit is contained in:
nabil-tazi
@@ -0,0 +1,648 @@
|
||||
"""
|
||||
Evaluation harness for pocket-tts int8 quantization strategy.
|
||||
|
||||
For each quantization config:
|
||||
- Generates audio for eval sentences per voice
|
||||
- Measures Real-Time Speed (RTS): audio_duration / wall_clock_time
|
||||
- Measures quality: SNR vs baseline, PESQ, Whisper WER
|
||||
- Saves audio files for manual listening
|
||||
- Outputs a summary CSV and markdown report
|
||||
|
||||
Usage:
|
||||
uv run python scripts/evaluate_quantization.py --config all
|
||||
uv run python scripts/evaluate_quantization.py --all-configs
|
||||
uv run python scripts/evaluate_quantization.py --all-configs --skip-quality
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import statistics
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import scipy.io.wavfile
|
||||
import torch
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
from pocket_tts.quantization import apply_dynamic_int8
|
||||
|
||||
# Quantization configs for benchmarking. Each maps to a set of layer group keys.
|
||||
CONFIGS = {
|
||||
"baseline": set(),
|
||||
"all": {"attention", "ffn", "flow_net"},
|
||||
"attention_ffn": {"attention", "ffn"},
|
||||
"attention": {"attention"},
|
||||
"ffn": {"ffn"},
|
||||
"flow_net": {"flow_net"},
|
||||
"flow_net_attention": {"flow_net", "attention"},
|
||||
"ffn_flow_net": {"ffn", "flow_net"},
|
||||
}
|
||||
|
||||
|
||||
def get_model_size_mb(model: torch.nn.Module) -> float:
|
||||
"""Returns model size in MB by serializing to a byte buffer."""
|
||||
import io
|
||||
|
||||
buf = io.BytesIO()
|
||||
torch.save(model.state_dict(), buf)
|
||||
return buf.tell() / (1024**2)
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
VOICES = ["alba", "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma"]
|
||||
|
||||
# Main eval paragraph (prosody / naturalness)
|
||||
EVAL_TEXT = (
|
||||
"Sarah had always wondered whether the old lighthouse still worked. "
|
||||
"On clear mornings, she'd walk the narrow coastal path — past the fishing boats, "
|
||||
"the salt-weathered fences, the occasional curious gull — and stop at the cliff's "
|
||||
"edge to look. Three years ago, she'd climbed the iron staircase herself: one hundred "
|
||||
"and twelve steps, each one groaning under her boots. At the top, the lens was still "
|
||||
"there, enormous and dusty, its glass facets catching the pale winter light. Nobody "
|
||||
"had switched it on in decades. Still, on stormy nights, she sometimes thought she "
|
||||
"could see a faint glow from her bedroom window. She never mentioned this to anyone."
|
||||
)
|
||||
|
||||
# Diverse sentences for WER stress testing
|
||||
QUALITY_SENTENCES = [
|
||||
# Normal prose
|
||||
"Sarah had always wondered whether the old lighthouse still worked.",
|
||||
# Numbers
|
||||
"The package weighs one hundred and twelve grams and costs forty-three dollars.",
|
||||
# Rare words
|
||||
"The archaeologist meticulously catalogued the palaeolithic artefacts.",
|
||||
# Punctuation-heavy (prosody stress test)
|
||||
"Wait — are you certain? Because if not, we should stop, reconsider, and try again.",
|
||||
# Short sentence (trailing artifact detector)
|
||||
"Yes.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
voice: str
|
||||
config_id: str
|
||||
audio_duration_sec: float
|
||||
wall_clock_sec: float
|
||||
rts: float
|
||||
load_time_sec: float
|
||||
model_size_mb: float
|
||||
output_path: str
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class QualityResult:
|
||||
config_id: str
|
||||
voice: str
|
||||
sentence_idx: int
|
||||
sentence: str
|
||||
snr_db: Optional[float] = None
|
||||
pesq_score: Optional[float] = None
|
||||
wer: Optional[float] = None
|
||||
baseline_transcript: Optional[str] = None
|
||||
quantized_transcript: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfigSummary:
|
||||
config_id: str
|
||||
quantize_groups: list
|
||||
model_size_mb: float
|
||||
load_time_sec: float
|
||||
mean_rts: float
|
||||
min_rts: float
|
||||
max_rts: float
|
||||
rts_vs_baseline: Optional[float] = None
|
||||
mean_snr_db: Optional[float] = None
|
||||
mean_pesq: Optional[float] = None
|
||||
mean_wer_baseline: Optional[float] = None
|
||||
mean_wer_quantized: Optional[float] = None
|
||||
results: list = field(default_factory=list)
|
||||
quality_results: list = field(default_factory=list)
|
||||
|
||||
|
||||
def load_and_quantize_model(config_id: str):
|
||||
"""Load and quantize model. Returns (model, load_time, model_size_mb)."""
|
||||
quantize_groups = CONFIGS[config_id]
|
||||
|
||||
t0 = time.perf_counter()
|
||||
tts_model = TTSModel.load_model(temp=0.0)
|
||||
|
||||
if quantize_groups:
|
||||
logger.info("[%s] Applying int8 quantization: %s", config_id, quantize_groups)
|
||||
apply_dynamic_int8(tts_model.flow_lm, quantize_groups)
|
||||
|
||||
load_time = time.perf_counter() - t0
|
||||
model_size = get_model_size_mb(tts_model.flow_lm)
|
||||
|
||||
logger.info("[%s] Loaded in %.2fs, FlowLM size ~%.1fMB", config_id, load_time, model_size)
|
||||
return tts_model, load_time, model_size
|
||||
|
||||
|
||||
def save_audio(audio_tensor, sample_rate, output_path):
|
||||
"""Save audio tensor to WAV file."""
|
||||
audio_np = audio_tensor.numpy() if hasattr(audio_tensor, "numpy") else np.array(audio_tensor)
|
||||
audio_int16 = (audio_np * 32767).clip(-32768, 32767).astype(np.int16)
|
||||
scipy.io.wavfile.write(str(output_path), sample_rate, audio_int16)
|
||||
|
||||
|
||||
def generate_for_voice(tts_model, voice, config_id, output_dir):
|
||||
"""Generate audio for a single voice and return metrics."""
|
||||
logger.info("[%s] Generating for voice: %s", config_id, voice)
|
||||
|
||||
try:
|
||||
voice_state = tts_model.get_state_for_audio_prompt(voice)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
audio = tts_model.generate_audio(model_state=voice_state, text_to_generate=EVAL_TEXT)
|
||||
wall_clock = time.perf_counter() - t0
|
||||
|
||||
audio_duration = len(audio) / tts_model.sample_rate
|
||||
rts = audio_duration / wall_clock
|
||||
|
||||
output_path = output_dir / f"{config_id}__{voice}.wav"
|
||||
save_audio(audio, tts_model.sample_rate, output_path)
|
||||
|
||||
logger.info(
|
||||
"[%s][%s] audio=%.1fs wall=%.1fs RTS=%.2fx",
|
||||
config_id,
|
||||
voice,
|
||||
audio_duration,
|
||||
wall_clock,
|
||||
rts,
|
||||
)
|
||||
|
||||
return GenerationResult(
|
||||
voice=voice,
|
||||
config_id=config_id,
|
||||
audio_duration_sec=round(audio_duration, 2),
|
||||
wall_clock_sec=round(wall_clock, 2),
|
||||
rts=round(rts, 3),
|
||||
load_time_sec=0.0,
|
||||
model_size_mb=0.0,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("[%s][%s] Failed: %s", config_id, voice, e)
|
||||
return GenerationResult(
|
||||
voice=voice,
|
||||
config_id=config_id,
|
||||
audio_duration_sec=0,
|
||||
wall_clock_sec=0,
|
||||
rts=0,
|
||||
load_time_sec=0,
|
||||
model_size_mb=0,
|
||||
output_path="",
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Quality metrics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_snr(baseline_audio: torch.Tensor, quantized_audio: torch.Tensor) -> float:
|
||||
"""Compute Signal-to-Noise Ratio in dB between baseline and quantized audio."""
|
||||
# Align lengths (generation is non-deterministic, lengths may differ slightly)
|
||||
min_len = min(len(baseline_audio), len(quantized_audio))
|
||||
baseline = baseline_audio[:min_len].float()
|
||||
quantized = quantized_audio[:min_len].float()
|
||||
|
||||
noise = baseline - quantized
|
||||
signal_power = baseline.pow(2).mean()
|
||||
noise_power = noise.pow(2).mean()
|
||||
|
||||
if noise_power == 0:
|
||||
return float("inf")
|
||||
return (10 * torch.log10(signal_power / noise_power)).item()
|
||||
|
||||
|
||||
def compute_pesq(
|
||||
baseline_audio: np.ndarray, quantized_audio: np.ndarray, sample_rate: int
|
||||
) -> Optional[float]:
|
||||
"""Compute PESQ score. Returns None if pesq is not installed."""
|
||||
try:
|
||||
from pesq import pesq
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
# PESQ requires 16kHz or 8kHz
|
||||
target_sr = 16000
|
||||
if sample_rate != target_sr:
|
||||
import scipy.signal
|
||||
|
||||
baseline_audio = scipy.signal.resample(
|
||||
baseline_audio, int(len(baseline_audio) * target_sr / sample_rate)
|
||||
)
|
||||
quantized_audio = scipy.signal.resample(
|
||||
quantized_audio, int(len(quantized_audio) * target_sr / sample_rate)
|
||||
)
|
||||
|
||||
# Align lengths
|
||||
min_len = min(len(baseline_audio), len(quantized_audio))
|
||||
baseline_audio = baseline_audio[:min_len].astype(np.float32)
|
||||
quantized_audio = quantized_audio[:min_len].astype(np.float32)
|
||||
|
||||
# Normalize to [-1, 1]
|
||||
max_val = max(np.abs(baseline_audio).max(), np.abs(quantized_audio).max(), 1e-8)
|
||||
baseline_audio = baseline_audio / max_val
|
||||
quantized_audio = quantized_audio / max_val
|
||||
|
||||
try:
|
||||
return pesq(target_sr, baseline_audio, quantized_audio, "wb")
|
||||
except Exception as e:
|
||||
logger.warning("PESQ failed: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def compute_wer(audio_path: str, reference_text: str, whisper_model) -> tuple[float, str]:
|
||||
"""Compute Word Error Rate using Whisper. Returns (wer, transcript)."""
|
||||
try:
|
||||
from jiwer import wer
|
||||
except ImportError:
|
||||
return None, ""
|
||||
|
||||
result = whisper_model.transcribe(str(audio_path), language="en")
|
||||
transcript = result["text"].strip()
|
||||
error_rate = wer(reference_text, transcript)
|
||||
return error_rate, transcript
|
||||
|
||||
|
||||
def run_quality_eval(
|
||||
baseline_model, quantized_model, config_id, voice, output_dir, whisper_model=None
|
||||
):
|
||||
"""Run quality evaluation for a single voice across all quality sentences."""
|
||||
results = []
|
||||
voice_state_b = baseline_model.get_state_for_audio_prompt(voice)
|
||||
voice_state_q = quantized_model.get_state_for_audio_prompt(voice)
|
||||
|
||||
for i, sentence in enumerate(QUALITY_SENTENCES):
|
||||
logger.info(
|
||||
"[%s][%s] Quality sentence %d/%d", config_id, voice, i + 1, len(QUALITY_SENTENCES)
|
||||
)
|
||||
|
||||
baseline_audio = baseline_model.generate_audio(
|
||||
model_state=voice_state_b, text_to_generate=sentence
|
||||
)
|
||||
quantized_audio = quantized_model.generate_audio(
|
||||
model_state=voice_state_q, text_to_generate=sentence
|
||||
)
|
||||
|
||||
sr = baseline_model.sample_rate
|
||||
qr = QualityResult(config_id=config_id, voice=voice, sentence_idx=i, sentence=sentence)
|
||||
|
||||
# SNR
|
||||
qr.snr_db = round(compute_snr(baseline_audio, quantized_audio), 2)
|
||||
|
||||
# PESQ
|
||||
qr.pesq_score = compute_pesq(baseline_audio.numpy(), quantized_audio.numpy(), sr)
|
||||
if qr.pesq_score is not None:
|
||||
qr.pesq_score = round(qr.pesq_score, 3)
|
||||
|
||||
# WER (save temp files for Whisper)
|
||||
if whisper_model is not None:
|
||||
baseline_path = output_dir / f"quality_{config_id}__{voice}__s{i}_baseline.wav"
|
||||
quantized_path = output_dir / f"quality_{config_id}__{voice}__s{i}_quantized.wav"
|
||||
save_audio(baseline_audio, sr, baseline_path)
|
||||
save_audio(quantized_audio, sr, quantized_path)
|
||||
|
||||
wer_b, transcript_b = compute_wer(baseline_path, sentence, whisper_model)
|
||||
wer_q, transcript_q = compute_wer(quantized_path, sentence, whisper_model)
|
||||
|
||||
qr.wer = wer_q
|
||||
qr.baseline_transcript = transcript_b
|
||||
qr.quantized_transcript = transcript_q
|
||||
|
||||
logger.info(
|
||||
"[%s][%s][s%d] SNR=%.1fdB PESQ=%s WER_b=%.2f WER_q=%.2f",
|
||||
config_id,
|
||||
voice,
|
||||
i,
|
||||
qr.snr_db,
|
||||
qr.pesq_score or "N/A",
|
||||
wer_b or 0,
|
||||
wer_q or 0,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"[%s][%s][s%d] SNR=%.1fdB PESQ=%s",
|
||||
config_id,
|
||||
voice,
|
||||
i,
|
||||
qr.snr_db,
|
||||
qr.pesq_score or "N/A",
|
||||
)
|
||||
|
||||
results.append(qr)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run configs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_config(config_id, output_dir, voices):
|
||||
"""Run full evaluation for a single quantization config."""
|
||||
logger.info("=" * 60)
|
||||
logger.info("Evaluating config: %s (groups: %s)", config_id, CONFIGS[config_id] or "none")
|
||||
logger.info("=" * 60)
|
||||
|
||||
tts_model, load_time, model_size = load_and_quantize_model(config_id)
|
||||
|
||||
results = []
|
||||
for voice in voices:
|
||||
result = generate_for_voice(tts_model, voice, config_id, output_dir)
|
||||
result.load_time_sec = round(load_time, 2)
|
||||
result.model_size_mb = round(model_size, 1)
|
||||
results.append(result)
|
||||
|
||||
valid_rts = [r.rts for r in results if not r.error and r.rts > 0]
|
||||
mean_rts = statistics.mean(valid_rts) if valid_rts else 0
|
||||
min_rts = min(valid_rts) if valid_rts else 0
|
||||
max_rts = max(valid_rts) if valid_rts else 0
|
||||
|
||||
return ConfigSummary(
|
||||
config_id=config_id,
|
||||
quantize_groups=sorted(CONFIGS[config_id]),
|
||||
model_size_mb=round(model_size, 1),
|
||||
load_time_sec=round(load_time, 2),
|
||||
mean_rts=round(mean_rts, 3),
|
||||
min_rts=round(min_rts, 3),
|
||||
max_rts=round(max_rts, 3),
|
||||
results=results,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reporting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def write_csv(summaries, output_dir):
|
||||
csv_path = output_dir / "results.csv"
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(
|
||||
f,
|
||||
fieldnames=[
|
||||
"config_id",
|
||||
"quantize_groups",
|
||||
"voice",
|
||||
"audio_duration_sec",
|
||||
"wall_clock_sec",
|
||||
"rts",
|
||||
"load_time_sec",
|
||||
"model_size_mb",
|
||||
"error",
|
||||
"output_path",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
for summary in summaries:
|
||||
for r in summary.results:
|
||||
writer.writerow(
|
||||
{
|
||||
"config_id": r.config_id,
|
||||
"quantize_groups": "|".join(summary.quantize_groups) or "none",
|
||||
"voice": r.voice,
|
||||
"audio_duration_sec": r.audio_duration_sec,
|
||||
"wall_clock_sec": r.wall_clock_sec,
|
||||
"rts": r.rts,
|
||||
"load_time_sec": r.load_time_sec,
|
||||
"model_size_mb": r.model_size_mb,
|
||||
"error": r.error or "",
|
||||
"output_path": r.output_path,
|
||||
}
|
||||
)
|
||||
logger.info("CSV written: %s", csv_path)
|
||||
|
||||
|
||||
def write_quality_csv(summaries, output_dir):
|
||||
csv_path = output_dir / "quality_results.csv"
|
||||
rows = []
|
||||
for s in summaries:
|
||||
for qr in s.quality_results:
|
||||
rows.append(asdict(qr))
|
||||
if not rows:
|
||||
return
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
logger.info("Quality CSV written: %s", csv_path)
|
||||
|
||||
|
||||
def write_markdown_report(summaries, output_dir):
|
||||
report_path = output_dir / "report.md"
|
||||
lines = [
|
||||
"# pocket-tts int8 Quantization Evaluation Report",
|
||||
f"\nGenerated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
|
||||
f"\nEval text length: {len(EVAL_TEXT.split())} words",
|
||||
f"\nVoices tested: {', '.join(VOICES)}",
|
||||
"\n---\n",
|
||||
"## Speed Summary\n",
|
||||
"| Config | Quantized Groups | Model Size | Load Time | Mean RTS | Speedup vs Baseline |",
|
||||
"|--------|-----------------|------------|-----------|----------|---------------------|",
|
||||
]
|
||||
|
||||
for s in summaries:
|
||||
groups_str = ", ".join(s.quantize_groups) if s.quantize_groups else "none"
|
||||
speedup = f"{s.rts_vs_baseline:.2f}x" if s.rts_vs_baseline is not None else "---"
|
||||
lines.append(
|
||||
f"| `{s.config_id}` | {groups_str} | {s.model_size_mb:.1f} MB "
|
||||
f"| {s.load_time_sec:.1f}s | {s.mean_rts:.2f}x | {speedup} |"
|
||||
)
|
||||
|
||||
# Quality summary
|
||||
has_quality = any(s.quality_results for s in summaries)
|
||||
if has_quality:
|
||||
lines += ["\n---\n", "## Quality Summary\n"]
|
||||
lines.append(
|
||||
"| Config | Mean SNR (dB) | Mean PESQ | Mean WER (baseline) | Mean WER (quantized) |"
|
||||
)
|
||||
lines.append(
|
||||
"|--------|--------------|-----------|--------------------|--------------------|"
|
||||
)
|
||||
for s in summaries:
|
||||
if not s.quality_results:
|
||||
continue
|
||||
snr = f"{s.mean_snr_db:.1f}" if s.mean_snr_db is not None else "N/A"
|
||||
pesq_s = f"{s.mean_pesq:.3f}" if s.mean_pesq is not None else "N/A"
|
||||
wer_b = f"{s.mean_wer_baseline:.3f}" if s.mean_wer_baseline is not None else "N/A"
|
||||
wer_q = f"{s.mean_wer_quantized:.3f}" if s.mean_wer_quantized is not None else "N/A"
|
||||
lines.append(f"| `{s.config_id}` | {snr} | {pesq_s} | {wer_b} | {wer_q} |")
|
||||
|
||||
lines += ["\n---\n", "## Per-Voice Speed Results\n"]
|
||||
|
||||
for s in summaries:
|
||||
lines.append(f"### `{s.config_id}`\n")
|
||||
lines.append("| Voice | Audio Duration | Wall Clock | RTS | Notes |")
|
||||
lines.append("|-------|---------------|------------|-----|-------|")
|
||||
for r in s.results:
|
||||
note = f"ERROR: {r.error}" if r.error else "(listen)"
|
||||
lines.append(
|
||||
f"| {r.voice} | {r.audio_duration_sec:.1f}s "
|
||||
f"| {r.wall_clock_sec:.1f}s | {r.rts:.2f}x | {note} |"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
report_path.write_text("\n".join(lines))
|
||||
logger.info("Report written: %s", report_path)
|
||||
|
||||
|
||||
def write_json_summary(summaries, output_dir):
|
||||
json_path = output_dir / "summary.json"
|
||||
json_path.write_text(json.dumps([asdict(s) for s in summaries], indent=2))
|
||||
logger.info("JSON summary written: %s", json_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Evaluate pocket-tts quantization configs")
|
||||
parser.add_argument("--config", nargs="+", choices=list(CONFIGS.keys()))
|
||||
parser.add_argument("--all-configs", action="store_true")
|
||||
parser.add_argument("--voices", nargs="+", default=VOICES, choices=VOICES)
|
||||
parser.add_argument("--output-dir", default="./eval_output")
|
||||
parser.add_argument(
|
||||
"--skip-quality", action="store_true", help="Skip quality metrics (SNR/PESQ/WER)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quality-voices",
|
||||
nargs="+",
|
||||
default=["alba", "marius"],
|
||||
help="Voices to run quality eval on (default: alba, marius)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
configs_to_run = (
|
||||
list(CONFIGS.keys()) if args.all_configs else (args.config or ["baseline", "all"])
|
||||
)
|
||||
|
||||
# Ensure baseline is first
|
||||
configs_to_run = ["baseline"] + [c for c in configs_to_run if c != "baseline"]
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_dir = Path(args.output_dir) / timestamp
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Output directory: %s", output_dir)
|
||||
|
||||
# --- Speed evaluation ---
|
||||
summaries = []
|
||||
baseline_rts = None
|
||||
|
||||
for config_id in configs_to_run:
|
||||
summary = run_config(config_id, output_dir, args.voices)
|
||||
|
||||
if config_id == "baseline":
|
||||
baseline_rts = summary.mean_rts
|
||||
elif baseline_rts and baseline_rts > 0:
|
||||
summary.rts_vs_baseline = round(summary.mean_rts / baseline_rts, 3)
|
||||
|
||||
summaries.append(summary)
|
||||
|
||||
# --- Quality evaluation ---
|
||||
if not args.skip_quality:
|
||||
# Load Whisper model once
|
||||
whisper_model = None
|
||||
try:
|
||||
import whisper
|
||||
|
||||
logger.info("Loading Whisper model for WER evaluation...")
|
||||
whisper_model = whisper.load_model("base")
|
||||
logger.info("Whisper model loaded.")
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"openai-whisper not installed, skipping WER. Install with: pip install openai-whisper jiwer"
|
||||
)
|
||||
|
||||
# Load baseline model once for quality comparison
|
||||
baseline_model = TTSModel.load_model()
|
||||
|
||||
for summary in summaries:
|
||||
if summary.config_id == "baseline":
|
||||
continue # no need to compare baseline against itself
|
||||
|
||||
quantized_model, _, _ = load_and_quantize_model(summary.config_id)
|
||||
|
||||
all_quality = []
|
||||
for voice in args.quality_voices:
|
||||
if voice not in args.voices:
|
||||
continue
|
||||
quality_results = run_quality_eval(
|
||||
baseline_model,
|
||||
quantized_model,
|
||||
summary.config_id,
|
||||
voice,
|
||||
output_dir,
|
||||
whisper_model=whisper_model,
|
||||
)
|
||||
all_quality.extend(quality_results)
|
||||
|
||||
summary.quality_results = all_quality
|
||||
|
||||
# Aggregate quality metrics
|
||||
snrs = [q.snr_db for q in all_quality if q.snr_db is not None]
|
||||
pesqs = [q.pesq_score for q in all_quality if q.pesq_score is not None]
|
||||
wers_q = [q.wer for q in all_quality if q.wer is not None]
|
||||
|
||||
if snrs:
|
||||
summary.mean_snr_db = round(statistics.mean(snrs), 2)
|
||||
if pesqs:
|
||||
summary.mean_pesq = round(statistics.mean(pesqs), 3)
|
||||
if wers_q:
|
||||
summary.mean_wer_quantized = round(statistics.mean(wers_q), 3)
|
||||
|
||||
# Also compute baseline WER for comparison
|
||||
if whisper_model is not None:
|
||||
from jiwer import wer as compute_wer_score
|
||||
|
||||
baseline_wers = []
|
||||
for qr in all_quality:
|
||||
if qr.baseline_transcript is not None:
|
||||
baseline_wers.append(compute_wer_score(qr.sentence, qr.baseline_transcript))
|
||||
if baseline_wers:
|
||||
summary.mean_wer_baseline = round(statistics.mean(baseline_wers), 3)
|
||||
|
||||
# --- Write reports ---
|
||||
write_csv(summaries, output_dir)
|
||||
write_quality_csv(summaries, output_dir)
|
||||
write_markdown_report(summaries, output_dir)
|
||||
write_json_summary(summaries, output_dir)
|
||||
|
||||
# --- Print summary ---
|
||||
print("\n" + "=" * 60)
|
||||
print("EVALUATION COMPLETE")
|
||||
print("=" * 60)
|
||||
print(
|
||||
f"{'Config':<20} {'Mean RTS':>10} {'Speedup':>10} {'Size MB':>10} {'SNR dB':>10} {'PESQ':>8} {'WER':>8}"
|
||||
)
|
||||
print("-" * 86)
|
||||
for s in summaries:
|
||||
speedup = f"{s.rts_vs_baseline:.2f}x" if s.rts_vs_baseline else "---"
|
||||
snr = f"{s.mean_snr_db:.1f}" if s.mean_snr_db is not None else "---"
|
||||
pesq_s = f"{s.mean_pesq:.3f}" if s.mean_pesq is not None else "---"
|
||||
wer_s = f"{s.mean_wer_quantized:.3f}" if s.mean_wer_quantized is not None else "---"
|
||||
print(
|
||||
f"{s.config_id:<20} {s.mean_rts:>10.2f}x {speedup:>10} {s.model_size_mb:>10.1f} {snr:>10} {pesq_s:>8} {wer_s:>8}"
|
||||
)
|
||||
print(f"\nAudio files + report: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user