Add WER metric (#28)

Author: nikolasavic3

src/metrax/__init__.py

@@ -21,6 +21,7 @@
     RMSE,
     RSQUARED,
     Recall,
+    WER,
 )
 
 __all__ = [
@@ -32,4 +33,5 @@
     "Recall",
     "AUCPR",
     "AUCROC",
+    "WER",
 ]

src/metrax/metrics.py

@@ -165,7 +165,7 @@ class RSQUARED(clu_metrics.Metric):
   count: jax.Array
   sum_of_squared_error: jax.Array
   sum_of_squared_label: jax.Array
-  
+
 
   @classmethod
   def empty(cls) -> 'RSQUARED':
@@ -436,7 +436,7 @@ class AUCPR(clu_metrics.Metric):
   false_positives: jax.Array
   false_negatives: jax.Array
   num_thresholds: int
-  
+
   @classmethod
   def empty(cls) -> 'AUCPR':
     return cls(
@@ -795,3 +795,137 @@ def merge(self, other: 'Perplexity') -> 'Perplexity':
 
   def compute(self) -> jax.Array:
     return jnp.exp(self.aggregate_crossentropy / self.num_samples)
+
+
+@flax.struct.dataclass
+class WER(clu_metrics.Metric):
+  r"""Computes Word Error Rate (WER) for speech recognition or text generation tasks.
+
+  Word Error Rate measures the edit distance between reference texts and predictions,
+  normalized by the length of the reference texts. It is calculated as:
+
+  .. math::
+      WER = \frac{S + D + I}{N}
+
+  where:
+      - S is the number of substitutions
+      - D is the number of deletions
+      - I is the number of insertions
+      - N is the number of words in the reference
+
+  A lower WER indicates better performance, with 0 being perfect.
+
+  Attributes:
+      total_edit_distance: Sum of edit distances across all samples.
+      total_reference_length: Sum of reference lengths across all samples.
+  """
+
+  total_edit_distance: jax.Array
+  total_reference_length: jax.Array
+
+  @classmethod
+  def empty(cls) -> 'WER':
+    return cls(
+        total_edit_distance=jnp.array(0, jnp.float32),
+        total_reference_length=jnp.array(0, jnp.float32))
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: list[str] | list[list],
+      references: list[str] | list[list],
+  ) -> 'WER':
+    """Updates the metric.
+
+    Args:
+        predictions: A list of predicted texts/transcriptions or tokenized sequences.
+        references: A list of reference texts/transcriptions or tokenized sequences.
+
+    Returns:
+        Updated WER metric.
+
+    Raises:
+        ValueError: If inputs are not properly formatted or are empty.
+    """
+    if not predictions or not references:
+      raise ValueError('predictions and references must not be empty')
+
+    if len(predictions) != len(references):
+      raise ValueError(
+          f'Length mismatch: predictions has {len(predictions)} items, '
+          f'but references has {len(references)} items'
+      )
+
+    # Determine if inputs are strings that need tokenization or already tokenized
+    total_edit_distance = 0
+    total_reference_length = 0
+
+    for pred, ref in zip(predictions, references):
+      # Convert to tokens if needed
+      pred_tokens = pred.split() if isinstance(pred, str) else pred
+      ref_tokens = ref.split() if isinstance(ref, str) else ref
+
+      edit_distance = cls._levenshtein_distance(pred_tokens, ref_tokens)
+
+      total_edit_distance += edit_distance
+      total_reference_length += len(ref_tokens)
+
+    return cls(
+        total_edit_distance=jnp.array(total_edit_distance, dtype=jnp.float32),
+        total_reference_length=jnp.array(total_reference_length, dtype=jnp.float32),
+    )
+
+  @staticmethod
+  def _levenshtein_distance(prediction: list, reference: list) -> int:
+    """Computes the Levenshtein (edit) distance between two token sequences.
+
+    Args:
+        prediction: List of tokens in the predicted sequence.
+        reference: List of tokens in the reference sequence.
+
+    Returns:
+        The minimum number of edits needed to transform prediction into reference.
+    """
+    m, n = len(prediction), len(reference)
+
+    # Handle edge cases
+    if m == 0:
+      return n
+    if n == 0:
+      return m
+
+    # Create distance matrix
+    distance_matrix = [[0 for _ in range(n+1)] for _ in range(m+1)]
+
+    # Initialize first row and column
+    for i in range(m+1):
+      distance_matrix[i][0] = i
+    for j in range(n+1):
+      distance_matrix[0][j] = j
+
+    # Fill the matrix
+    for i in range(1, m+1):
+      for j in range(1, n+1):
+        if prediction[i-1] == reference[j-1]:
+          cost = 0
+        else:
+          cost = 1
+
+        distance_matrix[i][j] = min(
+            distance_matrix[i-1][j] + 1,      # deletion
+            distance_matrix[i][j-1] + 1,      # insertion
+            distance_matrix[i-1][j-1] + cost  # substitution
+        )
+
+    return distance_matrix[m][n]
+
+  def merge(self, other: 'WER') -> 'WER':
+    return type(self)(
+        total_edit_distance=self.total_edit_distance + other.total_edit_distance,
+        total_reference_length=self.total_reference_length + other.total_reference_length,
+    )
+
+  def compute(self) -> jax.Array:
+    return _divide_no_nan(
+        self.total_edit_distance, self.total_reference_length
+    )

src/metrax/metrics_test.py

@@ -455,6 +455,103 @@ def test_perplexity(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  def test_wer_empty(self):
+    """Tests the `empty` method of the `WER` class."""
+    m = metrax.WER.empty()
+    self.assertEqual(m.total_edit_distance, jnp.array(0, jnp.float32))
+    self.assertEqual(m.total_reference_length, jnp.array(0, jnp.float32))
+
+  def test_wer(self):
+    """Tests that WER metric computes correct values."""
+    # Test with string inputs
+    predictions = [
+        "the cat sat on the mat",
+        "a quick brown fox jumps over the lazy dog",
+        "hello world"
+    ]
+    references = [
+        "the cat sat on the hat",
+        "the quick brown fox jumps over the lazy dog",
+        "hello beautiful world"
+    ]
+
+    metric = None
+    for pred, ref in zip(predictions, references):
+      update = metrax.WER.from_model_output(
+          predictions=[pred],
+          references=[ref],
+      )
+      metric = update if metric is None else metric.merge(update)
+
+    np.testing.assert_allclose(
+        metric.compute(),
+        jnp.array(3/18, dtype=jnp.float32),
+        rtol=1e-05,
+        atol=1e-05,
+    )
+
+  def test_wer_with_tokens(self):
+    """Tests that WER metric computes correct values with tokenized inputs."""
+    # Test with token inputs (lists of strings)
+    tokenized_preds = [
+        ["the", "cat", "sat", "on", "the", "mat"],
+        ["a", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
+        ["hello", "world"]
+    ]
+    tokenized_refs = [
+        ["the", "cat", "sat", "on", "the", "hat"],
+        ["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
+        ["hello", "beautiful", "world"]
+    ]
+
+    metric = None
+    for pred, ref in zip(tokenized_preds, tokenized_refs):
+      update = metrax.WER.from_model_output(
+          predictions=[pred],
+          references=[ref],
+      )
+      metric = update if metric is None else metric.merge(update)
+
+    np.testing.assert_allclose(
+        metric.compute(),
+        jnp.array(3/18, dtype=jnp.float32),
+        rtol=1e-05,
+        atol=1e-05,
+    )
+
+  def test_wer_merge(self):
+    """Tests the merge functionality of the WER metric."""
+    predictions1 = ["the cat sat on the mat"]
+    references1 = ["the cat sat on the hat"]
+
+    predictions2 = [
+        "a quick brown fox jumps over the lazy dog",
+        "hello world"
+    ]
+    references2 = [
+        "the quick brown fox jumps over the lazy dog",
+        "hello beautiful world"
+    ]
+
+    metric1 = metrax.WER.from_model_output(
+        predictions=predictions1,
+        references=references1,
+    )
+
+    metric2 = metrax.WER.from_model_output(
+        predictions=predictions2,
+        references=references2,
+    )
+
+    merged_metric = metric1.merge(metric2)
+
+    np.testing.assert_allclose(
+        merged_metric.compute(),
+        jnp.array(3/18, dtype=jnp.float32),
+        rtol=1e-05,
+        atol=1e-05,
+    )
+
 
 if __name__ == '__main__':
   os.environ['XLA_FLAGS'] = (