Fix num_examples computation in lf.eval.v2.

langfun/core/eval/v2/evaluation.py

@@ -108,8 +108,15 @@ def _example_input_by_id(self) -> dict[int, Any]:
   def num_examples(self) -> int:
     """Returns the number of examples from the inputs."""
     # NOTE(daiyip): setting `num_examples` of the input functor allows fast
-    # retrieval of number of examples without interating the whole dataset.
-    return getattr(self.inputs, 'num_examples', len(self.example_inputs))
+    # retrieval of number of examples without iterating the whole dataset.
+    num_examples = getattr(self.inputs, 'num_examples', None)
+    if not isinstance(num_examples, int):
+      it = self.example_inputs
+      if hasattr(it, '__len__'):
+        num_examples = len(it)
+      else:
+        num_examples = len(list(it))
+    return num_examples
 
   #
   # Evaluation logics.

langfun/core/eval/v2/evaluation_test.py

@@ -56,6 +56,18 @@ def test_hyper_evaluation(self):
         [set(['test_llm:0']), set(['test_llm:1']), set(['test_llm:2'])]
     )
 
+  def test_input(self):
+    exp = test_helper.TestEvaluation()
+    self.assertEqual(exp.num_examples, 10)
+    exp = test_helper.TestEvaluation(inputs=test_helper.test_inputs(None))
+    self.assertEqual(exp.num_examples, 20)
+    @pg.functor
+    def my_inputs():
+      yield pg.Dict(x=1, y=2)
+      yield pg.Dict(x=3, y=4)
+    exp = test_helper.TestEvaluation(inputs=my_inputs())
+    self.assertEqual(exp.num_examples, 2)
+
   def test_evaluate(self):
     exp = test_helper.TestEvaluation()
     example = exp.evaluate(Example(id=3))

langfun/core/eval/v2/test_helper.py

@@ -33,7 +33,9 @@
 
 
 @pg.functor()
-def test_inputs(num_examples: int = 10):
+def test_inputs(num_examples: int | None = 10):
+  if num_examples is None:
+    num_examples = 20
   return [
       pg.Dict(x=i, y=i ** 2, groundtruth=i + i ** 2)
       for i in range(num_examples)