Raise ValueError in LegacyMetric.from_model_output on output count mismatch

seqio/evaluation_test.py

@@ -1301,9 +1301,7 @@ def test_metric_with_arbitrary_model_input_types(self):
         predict_metric_fns=[_sequence_accuracy_metric],
         score_metric_fns=[_sum_scores_metric],
     )
-    task.postprocess_fn = (
-        lambda x, example, is_target: x if is_target else x.replace("e6", "e7")
-    )
+    task.postprocess_fn = lambda x, example, is_target: x
     task.metric_objs = [
         metrics_lib.LegacyMetric.empty(mf, task.postprocess_fn)
         for mf in task.metric_fns
@@ -1338,7 +1336,8 @@ def test_metric_with_arbitrary_model_input_types(self):
             vocabulary=vocabularies.PassThroughVocabulary(size=4)
         )
     }
-    dummy_outputs = np.array([])
+    # One model output per input example.
+    dummy_outputs = np.array([[0]])
 
     for inputs in model_inputs:
       for metric_obj in task.metric_objs:

seqio/metrics.py

@@ -318,11 +318,28 @@ def from_model_output(  # pylint:disable=arguments-renamed  # pyrefly: ignore[ba
             predictions.append(predictions_for_one_example)
         self.targets_and_inferences["output"] = predictions
 
-      # Postprocesses the predictions here.
-      postprocessed_predictions = [
-          self.postprocess_fn(p, example=ex, is_target=False)
-          for ex, p in zip(inputs, predictions)  # pyrefly: ignore[unbound-name]
-      ]
+      # Postprocesses the predictions here. `inputs` may be an iterable without
+      # a known length (e.g. a tf.data dataset), so the examples are counted
+      # while iterating. A plain `zip()` would silently truncate to the shorter
+      # of the two when the number of model outputs doesn't match.
+      num_predictions = len(predictions)  # pyrefly: ignore[unbound-name]
+      num_inputs = 0
+      postprocessed_predictions = []
+      for ex in inputs:
+        if num_inputs < num_predictions:
+          postprocessed_predictions.append(
+              self.postprocess_fn(
+                  predictions[num_inputs],  # pyrefly: ignore[unbound-name]
+                  example=ex,
+                  is_target=False,
+              )
+          )
+        num_inputs += 1
+      if num_inputs != num_predictions:
+        raise ValueError(
+            f"Number of model outputs ({num_predictions}) does not match the "
+            f"number of input examples ({num_inputs})."
+        )
 
       self.metric_fn_kwargs["predictions"] = postprocessed_predictions
       self.targets_and_inferences["prediction"] = postprocessed_predictions

seqio/metrics_test.py

@@ -83,6 +83,71 @@ def test_passthrough_legacy_metric(self, do_external_decode: bool):
     )
     self.assertEqual(metric_value["accuracy"], 1.0)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="prediction_fewer_outputs",
+          metric_fn=lambda targets, predictions: {},
+          model_output=np.array([[1, 2], [2, 3]]),
+          expected_message=r"model outputs \(2\).*input examples \(3\)",
+      ),
+      dict(
+          testcase_name="prediction_more_outputs",
+          metric_fn=lambda targets, predictions: {},
+          model_output=np.array([[1, 2], [2, 3], [3, 4], [4, 5]]),
+          expected_message=r"model outputs \(4\).*input examples \(3\)",
+      ),
+      dict(
+          testcase_name="prediction_top_k_fewer_outputs",
+          metric_fn=lambda targets, predictions: {},
+          model_output=np.array([[[1, 2], [1, 3]], [[2, 3], [2, 4]]]),
+          expected_message=r"model outputs \(2\).*input examples \(3\)",
+      ),
+      dict(
+          testcase_name="prediction_with_aux_fewer_outputs",
+          metric_fn=lambda targets, predictions, aux_values: {},
+          model_output=(
+              np.array([[1, 2], [2, 3]]),
+              {"scores": np.array([0.1, 0.2])},
+          ),
+          expected_message=r"model outputs \(2\).*input examples \(3\)",
+      ),
+  )
+  def test_legacy_metric_model_output_length_mismatch(
+      self, metric_fn, model_output, expected_message
+  ):
+    metric_obj = metrics.LegacyMetric.empty(metric_fn, None)
+    inputs = [
+        {"targets": [1, 2]},
+        {"targets": [2, 3]},
+        {"targets": [3, 4]},
+    ]
+    features = {
+        "targets": dataset_providers.Feature(
+            vocabularies.PassThroughVocabulary(size=8)
+        )
+    }
+    with self.assertRaisesRegex(ValueError, expected_message):
+      metric_obj.from_model_output(inputs, model_output, features)
+
+  def test_legacy_metric_model_output_length_matches(self):
+    metric_obj = metrics.LegacyMetric.empty(
+        lambda targets, predictions: {"num_predictions": len(predictions)},
+        None,
+    )
+    inputs = [
+        {"targets": [1, 2]},
+        {"targets": [2, 3]},
+    ]
+    features = {
+        "targets": dataset_providers.Feature(
+            vocabularies.PassThroughVocabulary(size=8)
+        )
+    }
+    metric_instance = metric_obj.from_model_output(
+        inputs, np.array([[1, 2], [2, 3]]), features
+    )
+    self.assertEqual(metric_instance.compute(), {"num_predictions": 2})
+
 
 if __name__ == "__main__":
   absltest.main()