Add WER normalization

requirements.txt

@@ -15,4 +15,4 @@ datasets
 torchcodec==0.10
 torch==2.10
 torchaudio
-pyloudnorm
+pyloudnorm

utils.py

@@ -12,6 +12,8 @@ from PIL import Image
 
 from constants import TARGET_LOUDNESS, TARGET_TP, VAD_OFF, VAD_ON
 
+_WER_NORMALIZERS: dict[str, object] = {}
+
 
 def get_vad_labels(vad_timestamps: list[list[float]], length: float) -> list[dict]:
     subtitles = []
@@ -166,12 +168,41 @@ def spec_image(
     return Image.open(buf).convert("RGB")
 
 
-def compute_wer(reference: str, hypothesis: str) -> float:
+def _normalize_wer_text(text: str, normalizer: Optional[str] = "basic") -> str:
+    """Normalize transcript text before WER tokenization."""
+    if normalizer is None:
+        return text.strip()
+
+    if normalizer not in {"basic", "english"}:
+        raise ValueError(
+            f"Unsupported normalizer '{normalizer}'. Use 'basic', 'english', or None."
+        )
+
+    if normalizer == "english":
+        if normalizer not in _WER_NORMALIZERS:
+            from whisper_normalizer.english import EnglishTextNormalizer
+
+            _WER_NORMALIZERS[normalizer] = EnglishTextNormalizer()
+    else:
+        if normalizer not in _WER_NORMALIZERS:
+            from whisper_normalizer.basic import BasicTextNormalizer
+
+            _WER_NORMALIZERS[normalizer] = BasicTextNormalizer()
+
+    return _WER_NORMALIZERS[normalizer](text)
+
+
+def compute_wer(
+    reference: str, hypothesis: str, normalizer: Optional[str] = "basic"
+) -> float:
     """
     Compute Word Error Rate (WER) between reference and hypothesis transcripts.
+
+    Text is normalized before tokenization. Supported normalizers are
+    ``"basic"``, ``"english"``, or ``None`` to disable normalization.
     """
-    ref_words = reference.split()
-    hyp_words = hypothesis.split()
+    ref_words = _normalize_wer_text(reference, normalizer=normalizer).split()
+    hyp_words = _normalize_wer_text(hypothesis, normalizer=normalizer).split()
 
     d = np.zeros((len(ref_words) + 1, len(hyp_words) + 1), dtype=np.uint16)
 
@@ -237,4 +268,4 @@ def normalize_lufs(x: np.ndarray, sr: int) -> np.ndarray:
         return y.astype(np.float32)
     except Exception as e:
         warnings.warn(f"LUFS normalization failed, returning input unchanged: {e}")
-        return x.astype(np.float32)
+        return x.astype(np.float32)