Repository: google/metrax
Pull request #30: Add WER metric
Author: login: nikolasavic3
Created: 2025-03-07T21:29:11Z | Merged: 2025-03-14T21:25:35Z

Add Word Error Rate (WER) metric 
Fixes #28


Linked issue #28: Add WER metric

Diff:
diff --git a/src/metrax/__init__.py b/src/metrax/__init__.py
index 78da1a6..1fc5b5c 100644
--- a/src/metrax/__init__.py
+++ b/src/metrax/__init__.py
@@ -20,6 +20,7 @@
 )
 from metrax.nlp_metrics import (
     Perplexity,
+    WER
 )
 from metrax.ranking_metrics import (
     AveragePrecisionAtK,
@@ -40,4 +41,5 @@
     "Recall",
     "RMSE",
     "RSQUARED",
+    "WER",
 ]
diff --git a/src/metrax/nlp_metrics.py b/src/metrax/nlp_metrics.py
index 93d83ed..327b5a2 100644
--- a/src/metrax/nlp_metrics.py
+++ b/src/metrax/nlp_metrics.py
@@ -114,3 +114,104 @@ 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.Average):
+    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.
+
+    This implementation accepts both pre-tokenized inputs (lists of tokens) and untokenized
+    strings. When strings are provided, they are tokenized by splitting on whitespace.
+    """
+
+    @classmethod
+    def from_model_output(
+        cls,
+        predictions: list[str],
+        references: list[str],
+    ) -> "WER":
+        """Updates the metric.
+
+        Args:
+            prediction: Either a string or a list of tokens in the predicted sequence.
+            reference: Either a string or a list of tokens in the reference sequence.
+
+        Returns:
+            New WER metric instance.
+
+        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 isinstance(predictions, str):
+            predictions = predictions.split()
+        if isinstance(references, str):
+            references = references.split()
+
+        edit_distance = cls._levenshtein_distance(predictions, references)
+        reference_length = len(references)
+
+        return cls(
+            total=jnp.array(edit_distance, dtype=jnp.float32),
+            count=jnp.array(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]
diff --git a/src/metrax/nlp_metrics_test.py b/src/metrax/nlp_metrics_test.py
index bd196a2..e6955ee 100644
--- a/src/metrax/nlp_metrics_test.py
+++ b/src/metrax/nlp_metrics_test.py
@@ -30,6 +30,12 @@ def test_perplexity_empty(self):
     self.assertEqual(m.aggregate_crossentropy, jnp.array(0, jnp.float32))
     self.assertEqual(m.num_samples, jnp.array(0, jnp.float32))
 
+  def test_wer_empty(self):
+    """Tests the `empty` method of the `WER` class."""
+    m = metrax.WER.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.float32))
+
   @parameterized.named_parameters(
       (
           'basic',
@@ -68,6 +74,49 @@ def test_perplexity(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  def test_wer(self):
+    """Tests that WER metric computes correct values with tokenized and untokenized inputs."""
+    string_preds = [
+      "the cat sat on the mat",
+      "a quick brown fox jumps over the lazy dog",
+      "hello world"
+    ]
+    string_refs = [
+      "the cat sat on the hat",
+      "the quick brown fox jumps over the lazy dog",
+      "hello beautiful world"
+    ]
+    tokenized_preds = [sentence.split() for sentence in string_preds]
+    tokenized_refs = [sentence.split() for sentence in string_refs]
+
+    metrax_token_metric = None
+    keras_metric = keras_hub.metrics.EditDistance(normalize=True)
+    for pred, ref in zip(tokenized_preds, tokenized_refs):
+      metrax_update = metrax.WER.from_model_output(pred,ref)
+      keras_metric.update_state(ref, pred)
+      metrax_token_metric = metrax_update if metrax_token_metric is None else metrax_token_metric.merge(metrax_update)
+
+    np.testing.assert_allclose(
+        metrax_token_metric.compute(),
+        keras_metric.result(),
+        rtol=1e-05,
+        atol=1e-05,
+        err_msg="String-based WER should match keras_hub EditDistance"
+    )
+
+    metrax_string_metric = None
+    for pred, ref in zip(string_preds, string_refs):
+      update = metrax.WER.from_model_output(predictions=pred, references=ref)
+      metrax_string_metric = update if metrax_string_metric is None else metrax_string_metric.merge(update)
+
+    np.testing.assert_allclose(
+    metrax_string_metric.compute(),
+    metrax_token_metric.compute(),
+    rtol=1e-05,
+    atol=1e-05,
+    err_msg="String input and tokenized input should produce the same WER"
+    )
+
 
 if __name__ == '__main__':
   absltest.main()


Reviews:
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T22:03:41Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:38:00Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:41:26Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:47:01Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:51:00Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:58:22Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T21:20:45Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T21:21:28Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T21:50:14Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T21:51:12Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T22:01:52Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T22:06:07Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T22:30:36Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T22:44:48Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T14:43:37Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T14:53:22Z
user: login: nikolasavic3
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T16:04:36Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:45:44Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:46:15Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:50:29Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:50:43Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:52:37Z
user: login: jshin1394
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:52:49Z

Review comments:
diff_hunk: @@ -795,3 +795,140 @@ 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
+
+      # Calculate edit distance
+      edit_distance = cls._levenshtein_distance(pred_tokens, ref_tokens)
+
+      # Update totals
+      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.
+    """
+    # Create a matrix to store the edit distances
+    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)]
path: src/metrax/metrics.py
user: login: nikolasavic3
body: I implemented the Levenshtein distance calculation using pure Python rather than JAX operations because the distance calculation involves updating the matrix in-place, and if we use JAX, that would create a new copy of the array for every update. I also believe that typical WER use cases have relatively short sequences, so the performance difference is negligible. Only the final results are converted to JAX arrays, which follows the pattern used in other metrics in the library.
created_at: 2025-03-07T22:03:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
path: src/metrax/metrics_test.py
user: login: jshin1394
body: not sure why this shows up as part of this PR. This code already exists...?
created_at: 2025-03-10T18:37:59Z
author_association: COLLABORATOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        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",  # 1 substitution (mat->hat), 6 total words
+        "the quick brown fox jumps over the lazy dog",  # 1 substitution (a->the), 9 total words
+        "hello beautiful world"  # 1 insertion (beautiful), 3 total words
+    ]
+
+    # Expected individual WERs: 1/6, 1/9, 1/3
+    # Total edit distance: 1 + 1 + 1 = 3
+    # Total reference length: 6 + 9 + 3 = 18
+    # Expected WER: 3/18 = 0.1667
+
+    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 instead 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"]  # 1/6 WER
+
+    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"
+    ]  # (1+1)/(9+3) = 2/12 WER
+
+    # Create and compute first metric
+    metric1 = metrax.WER.from_model_output(
+        predictions=predictions1,
+        references=references1,
+    )
+
+    # Create and compute second metric
+    metric2 = metrax.WER.from_model_output(
+        predictions=predictions2,
+        references=references2,
+    )
+
+    # Merge and compute
+    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__':
path: src/metrax/metrics_test.py
user: login: jshin1394
body: could you try syncing your workspace seems like some of more recent updates are not reflected.
created_at: 2025-03-10T18:41:26Z
author_association: COLLABORATOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        atol=1e-05,
+    )
+
+  def test_wer_empty(self):
path: src/metrax/metrics_test.py
user: login: jshin1394
body: let's move all *_empty tests to the beginning of the file.
created_at: 2025-03-10T18:47:01Z
author_association: COLLABORATOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        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",  # 1 substitution (mat->hat), 6 total words
+        "the quick brown fox jumps over the lazy dog",  # 1 substitution (a->the), 9 total words
+        "hello beautiful world"  # 1 insertion (beautiful), 3 total words
+    ]
+
+    # Expected individual WERs: 1/6, 1/9, 1/3
+    # Total edit distance: 1 + 1 + 1 = 3
+    # Total reference length: 6 + 9 + 3 = 18
+    # Expected WER: 3/18 = 0.1667
+
+    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),
path: src/metrax/metrics_test.py
user: login: jshin1394
body: can we test against the result gained from using a different known library? for instance, we can compare against torchmetrics https://pytorch.org/torcheval/main/generated/torcheval.metrics.WordErrorRate.html
created_at: 2025-03-10T18:51:00Z
author_association: COLLABORATOR
diff_hunk: @@ -536,4 +536,141 @@ def compute(self) -> jax.Array:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.
-    return jnp.trapezoid(tp_rate, fp_rate) * -1
\ No newline at end of file
+    return jnp.trapezoid(tp_rate, fp_rate) * -1
+
+
+@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):
path: src/metrax/metrics.py
user: login: jshin1394
body: should we add a TODO to use a more involved tokenizer in the future?
created_at: 2025-03-10T18:58:22Z
author_association: COLLABORATOR
diff_hunk: @@ -795,3 +795,140 @@ 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
+
+      # Calculate edit distance
+      edit_distance = cls._levenshtein_distance(pred_tokens, ref_tokens)
+
+      # Update totals
+      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.
+    """
+    # Create a matrix to store the edit distances
+    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)]
path: src/metrax/metrics.py
user: login: jshin1394
body: sounds good to me, thank you!
created_at: 2025-03-10T21:20:45Z
author_association: COLLABORATOR
diff_hunk: @@ -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):
path: src/metrax/metrics.py
user: login: jshin1394
body: should your metric inherit from clu_metrics.Average?
created_at: 2025-03-10T21:21:28Z
author_association: COLLABORATOR
diff_hunk: @@ -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):
path: src/metrax/metrics.py
user: login: nikolasavic3
body: Hi Jiwon, that's a good question! Seems to me that we can subclass WER from clu_metrics.Average and still keep all current functionality. On the other side, by using clu_metric.Metric I believe the WER metric would be more upgradable in the future, for example,if we decide to implement aggregating metrics across devices implementing the abstract reduce method. 
created_at: 2025-03-10T21:50:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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):
path: src/metrax/metrics.py
user: login: nikolasavic3
body: What do you think about it?
created_at: 2025-03-10T21:51:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        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",  # 1 substitution (mat->hat), 6 total words
+        "the quick brown fox jumps over the lazy dog",  # 1 substitution (a->the), 9 total words
+        "hello beautiful world"  # 1 insertion (beautiful), 3 total words
+    ]
+
+    # Expected individual WERs: 1/6, 1/9, 1/3
+    # Total edit distance: 1 + 1 + 1 = 3
+    # Total reference length: 6 + 9 + 3 = 18
+    # Expected WER: 3/18 = 0.1667
+
+    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),
path: src/metrax/metrics_test.py
user: login: nikolasavic3
body: Yes, that's a good idea. Maybe in the future, we could open a PR for keras_hub to implement the WER metric there so we have less imports? It already has ROUGE and BLEU. Or maybe we could use a simpler library like https://github.com/jitsi/jiwer. 
created_at: 2025-03-10T22:01:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
path: src/metrax/metrics_test.py
user: login: nikolasavic3
body: Sorry, that snippet was accidentally duplicated while merging branches. I have removed it now.
created_at: 2025-03-10T22:06:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -536,4 +536,141 @@ def compute(self) -> jax.Array:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.
-    return jnp.trapezoid(tp_rate, fp_rate) * -1
\ No newline at end of file
+    return jnp.trapezoid(tp_rate, fp_rate) * -1
+
+
+@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):
path: src/metrax/metrics.py
user: login: nikolasavic3
body: That's a great point! To keep things in line with the single responsibility principle, I think we should remove the tokenization logic entirely and just accept pre-tokenized lists of strings as input, leaving the tokenizer implementation to the one who calls the metric.
created_at: 2025-03-10T22:30:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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):
path: src/metrax/metrics.py
user: login: jshin1394
body: Hi Nikola! We definitely want to use these metrics in distributed settings. However, I wanted to point out that clu_metrics.Average also inherits from the clu_metrics.Metric class. Therefore, inheriting from clu_metrics.Average should be just as upgradable as clu_metrics.Metric. Please let me know if I am missing anything. Otherwise, +1 to using Average as base class! :)
created_at: 2025-03-10T22:44:48Z
author_association: COLLABORATOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        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",  # 1 substitution (mat->hat), 6 total words
+        "the quick brown fox jumps over the lazy dog",  # 1 substitution (a->the), 9 total words
+        "hello beautiful world"  # 1 insertion (beautiful), 3 total words
+    ]
+
+    # Expected individual WERs: 1/6, 1/9, 1/3
+    # Total edit distance: 1 + 1 + 1 = 3
+    # Total reference length: 6 + 9 + 3 = 18
+    # Expected WER: 3/18 = 0.1667
+
+    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),
path: src/metrax/metrics_test.py
user: login: nikolasavic3
body: the test now use [keras_nlp.metrics.EditDistance](https://keras.io/api/keras_nlp/metrics/edit_distance/) with `normalize=True` which is the same as computing WER, as mentioned in the documentation. 
created_at: 2025-03-13T14:43:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -536,4 +536,141 @@ def compute(self) -> jax.Array:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.
-    return jnp.trapezoid(tp_rate, fp_rate) * -1
\ No newline at end of file
+    return jnp.trapezoid(tp_rate, fp_rate) * -1
+
+
+@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):
path: src/metrax/metrics.py
user: login: nikolasavic3
body: The metric now accepts only lists of strings. Do you think we should accept strings and implement a tokenizer?
created_at: 2025-03-13T14:53:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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):
path: src/metrax/metrics.py
user: login: nikolasavic3
body: You’re right. The class now inherits from Average.
created_at: 2025-03-13T16:04:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -114,3 +117,126 @@ 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.Average):
+  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.
+
+  This implementation only accepts pre-tokenized inputs (lists of tokens).
+
+  Attributes:
+      total_edit_distance: Sum of edit distances across all samples.
+      total_reference_length: Sum of reference lengths across all samples.
+  """
+  @property
+  def total_edit_distance(self):
+    return self.total
+  @property
+  def total_reference_length(self):
+    return self.count
+
+  @classmethod
path: src/metrax/nlp_metrics.py
user: login: jshin1394
body: no need to override this function as it already exists in `Average`
created_at: 2025-03-13T21:45:44Z
author_association: COLLABORATOR
diff_hunk: @@ -114,3 +117,126 @@ 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.Average):
+  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.
+
+  This implementation only accepts pre-tokenized inputs (lists of tokens).
+
+  Attributes:
+      total_edit_distance: Sum of edit distances across all samples.
+      total_reference_length: Sum of reference lengths across all samples.
+  """
+  @property
+  def total_edit_distance(self):
+    return self.total
+  @property
+  def total_reference_length(self):
+    return self.count
+
+  @classmethod
+  def empty(cls) -> 'WER':
+    return cls(
+        total=jnp.array(0, jnp.int32),
+        count=jnp.array(0, jnp.int32))
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: list[str],
+      references: list[str],
+      **kwargs
+  ) -> 'WER':
+    """Updates the metric.
+
+    Args:
+        prediction: A list of tokens in the predicted sequence.
+        reference: A list of tokens in the reference sequence.
+
+    Returns:
+        New WER metric instance.
+
+    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')
+
+    edit_distance = cls._levenshtein_distance(predictions, references)
+    reference_length = len(references)
+
+    return cls(
+        total=jnp.array(edit_distance, dtype=jnp.float32),
+        count=jnp.array(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':
path: src/metrax/nlp_metrics.py
user: login: jshin1394
body: no need to override this function anymore as it exists in `Average`
created_at: 2025-03-13T21:46:15Z
author_association: COLLABORATOR
diff_hunk: @@ -536,4 +536,141 @@ def compute(self) -> jax.Array:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.
-    return jnp.trapezoid(tp_rate, fp_rate) * -1
\ No newline at end of file
+    return jnp.trapezoid(tp_rate, fp_rate) * -1
+
+
+@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):
path: src/metrax/metrics.py
user: login: jshin1394
body: Hi Nikola! I agree that single responsibility principle is important. However, considering that it is more common for these types of metrics to take in a whole sentence, (for instance, this [torchMetrics WER](https://pytorch.org/torcheval/main/generated/torcheval.metrics.WordErrorRate.html)) let's keep the implementation such that we can take in sentences as input. 
created_at: 2025-03-13T21:50:28Z
author_association: COLLABORATOR
diff_hunk: @@ -320,6 +320,148 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=1e-05,
     )
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          None,
+      ),
+      (
+          'weighted',
+          np.random.randint(10, size=[2, 5, 10]),
+          np.random.uniform(size=(2, 5, 10, 20)),
+          np.random.randint(2, size=(2, 5, 10)).astype(np.float32),
+      ),
+  )
+  def test_perplexity(self, y_true, y_pred, sample_weights):
+    keras_metric = keras_hub.metrics.Perplexity()
+    metrax_metric = None
+    for index, (labels, logits) in enumerate(zip(y_true, y_pred)):
+      weights = sample_weights[index] if sample_weights is not None else None
+      keras_metric.update_state(labels, logits, sample_weight=weights)
+      update = metrax.Perplexity.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = update if metrax_metric is None else metrax_metric.merge(
+          update
+      )
+
+    expected = keras_metric.result()
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        expected,
+        rtol=1e-05,
+        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",  # 1 substitution (mat->hat), 6 total words
+        "the quick brown fox jumps over the lazy dog",  # 1 substitution (a->the), 9 total words
+        "hello beautiful world"  # 1 insertion (beautiful), 3 total words
+    ]
+
+    # Expected individual WERs: 1/6, 1/9, 1/3
+    # Total edit distance: 1 + 1 + 1 = 3
+    # Total reference length: 6 + 9 + 3 = 18
+    # Expected WER: 3/18 = 0.1667
+
+    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),
path: src/metrax/metrics_test.py
user: login: jshin1394
body: perfect. Thank you so much!
created_at: 2025-03-13T21:50:43Z
author_association: COLLABORATOR
diff_hunk: @@ -114,3 +117,126 @@ 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.Average):
+  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.
+
+  This implementation only accepts pre-tokenized inputs (lists of tokens).
+
+  Attributes:
+      total_edit_distance: Sum of edit distances across all samples.
+      total_reference_length: Sum of reference lengths across all samples.
+  """
+  @property
+  def total_edit_distance(self):
+    return self.total
+  @property
+  def total_reference_length(self):
+    return self.count
+
+  @classmethod
+  def empty(cls) -> 'WER':
+    return cls(
+        total=jnp.array(0, jnp.int32),
+        count=jnp.array(0, jnp.int32))
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: list[str],
+      references: list[str],
+      **kwargs
+  ) -> 'WER':
+    """Updates the metric.
+
+    Args:
+        prediction: A list of tokens in the predicted sequence.
+        reference: A list of tokens in the reference sequence.
+
+    Returns:
+        New WER metric instance.
+
+    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')
+
+    edit_distance = cls._levenshtein_distance(predictions, references)
+    reference_length = len(references)
+
+    return cls(
+        total=jnp.array(edit_distance, dtype=jnp.float32),
+        count=jnp.array(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=self.total + other.total,
+        count=self.count + other.count,
+    )
+
+  def compute(self) -> jax.Array:
path: src/metrax/nlp_metrics.py
user: login: jshin1394
body: no need to override this function as it exists in `Average`
created_at: 2025-03-13T21:52:36Z
author_association: COLLABORATOR
diff_hunk: @@ -114,3 +117,126 @@ 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.Average):
+  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.
+
+  This implementation only accepts pre-tokenized inputs (lists of tokens).
+
+  Attributes:
+      total_edit_distance: Sum of edit distances across all samples.
+      total_reference_length: Sum of reference lengths across all samples.
+  """
+  @property
path: src/metrax/nlp_metrics.py
user: login: jshin1394
body: these can be removed as well since
created_at: 2025-03-13T21:52:49Z
author_association: COLLABORATOR

PR comments:
user: login: jshin1394
created_at: 2025-03-10T19:27:41Z
body: Thank you Nikola for your PR! Left some comments. Feel free to sync the workspace as well! :)
author_association: COLLABORATOR
user: login: nikolasavic3
created_at: 2025-03-10T22:32:38Z
body: Thank you Jiwon for reviewing my PR!
author_association: CONTRIBUTOR
user: login: jshin1394
created_at: 2025-03-10T22:57:15Z
body: Thank you Nikola! LGTM given all the comments are resolved.
author_association: COLLABORATOR
user: login: jshin1394
created_at: 2025-03-11T22:06:03Z
body: Hi nikola, a bunch of PRs were submitted to split the metrics.py file into multiple modules. I think the WER class can be placed under the new `nlp_metrics.py`module.
author_association: COLLABORATOR
user: login: jshin1394
created_at: 2025-03-13T21:54:43Z
body: Thank you so much nikola! Almost there :) I left some comments so that we can 1) remove functions that already exist as part of [`Average`](https://github.com/google/CommonLoopUtils/blob/main/clu/metrics.py#L770) and 2) the metric can take in sentence as input following other popular metric libraries such as [torchMetrics](https://pytorch.org/torcheval/main/generated/torcheval.metrics.WordErrorRate.html).

Thank you again Nikola!
author_association: COLLABORATOR
user: login: nikolasavic3
created_at: 2025-03-14T19:03:01Z
body: Thank you for the detailed feedback, Jiwon! I've removed all the redundant function overrides (merge, compute, total_reference_length) and updated the implementation to accept untokenized strings.
author_association: CONTRIBUTOR
user: login: jshin1394
created_at: 2025-03-14T21:26:34Z
body: Thank you so much Nikola! We really appreciate your contribution to the Metrax codebase! :)
author_association: COLLABORATOR
user: login: nikolasavic3
created_at: 2025-03-15T09:14:17Z
body: Thank you so much, Jiwon! I really appreciate your detailed feedback and guidance throughout this process. Collaborating with you on adding the WER metric to the Metrax codebase was great. Looking forward to contributing more in the future!
author_association: CONTRIBUTOR