langfun 0.1.2.dev202510200805__py3-none-any.whl → 0.1.2.dev202511160804__py3-none-any.whl

langfun/core/eval/scoring.py

@@ -41,18 +41,19 @@ class Scoring(base.Evaluation):
 
   @property
   def score_rate(self) -> float:
-    """Returns the score rate."""
+    """Returns the rate of scored examples among the completed ones."""
     if self.num_completed == 0:
       return 0.0
     return self.num_scored / self.num_completed
 
   @property
   def scored_link(self) -> str:
-    """Returns the matches page."""
+    """Returns the scored examples page."""
     return self.link(os.path.join(self.dir, Scoring.SCORED_HTML))
 
   @property
   def avg_score(self) -> float:
+    """Returns the average score of scored examples."""
     if self.num_scored == 0:
       return 0
     return sum([i[2] for i in self._scored]) / self.num_scored
@@ -181,7 +182,7 @@ class Scoring(base.Evaluation):
     super()._render_summary_metrics(s)
 
   def _render_scored(self, s: io.StringIO) -> None:
-    """Formats the matched cases into html."""
+    """Formats the scored cases into html."""
     s.write('<h2> Scored </h2>')
     s.write('<div style="white-space:pre">\n')
     s.write(

langfun/core/eval/v2/__init__.py

@@ -38,6 +38,7 @@ from langfun.core.eval.v2 import runners
 from langfun.core.eval.v2.checkpointing import BulkCheckpointer
 from langfun.core.eval.v2.checkpointing import PerExampleCheckpointer
 from langfun.core.eval.v2.reporting import HtmlReporter
+from langfun.core.eval.v2.reporting import ExampleHtmlGenerator
 
 
 # pylint: enable=g-bad-import-order

langfun/core/eval/v2/checkpointing.py

@@ -29,13 +29,28 @@ Runner = experiment_lib.Runner
 
 
 class Checkpointer(experiment_lib.Plugin):
-  """Base class for checkpointing evaluation examples."""
+  """Base class for checkpointing evaluation examples.
+
+  `Checkpointer` is a plugin that saves the state of processed examples
+  incrementally during an experiment run, allowing the experiment to be resumed
+  later. When an experiment starts, the checkpointer loads any previously saved
+  examples from an earlier run (or a warm-start run) into `experiment.state`,
+  so the runner can skip processing them again.
+  Subclasses should implement `_list_checkpoint_filenames` to identify
+  checkpoint files to load, and `_save_example` to save a newly processed
+  example.
+  """
 
   checkpoint_filename: Annotated[
       str,
       'Checkpoint file pattern.'
   ] = 'checkpoint.bagz'
 
+  max_ckpt_loading_threads: Annotated[
+      int,
+      'Max number of workers for loading checkpoint files at startup.'
+  ] = 128
+
   def on_experiment_start(
       self,
       runner: Runner,
@@ -149,7 +164,10 @@ class Checkpointer(experiment_lib.Plugin):
 
     _ = list(
         lf.concurrent_map(
-            _load_state, ckpt_files, max_workers=16, silence_on_errors=None
+            _load_state,
+            ckpt_files,
+            max_workers=self.max_ckpt_loading_threads,
+            silence_on_errors=None
         )
     )
 
@@ -170,7 +188,12 @@ class Checkpointer(experiment_lib.Plugin):
 
 
 class PerExampleCheckpointer(Checkpointer):
-  """Checkpointer that saves each example to a separate file."""
+  """Checkpointer that saves each example to a separate file.
+
+  This checkpointer saves each processed example to its own checkpoint file,
+  named using the pattern `<checkpoint_filename_prefix>_<example_id>.<ext>`.
+  For example, `checkpoint_1.bagz`, `checkpoint_2.bagz`, etc.
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -235,7 +258,13 @@ class PerExampleCheckpointer(Checkpointer):
 
 
 class BulkCheckpointer(Checkpointer):
-  """Checkpointer that saves all examples to a single file."""
+  """Checkpointer that saves all examples of an evaluation to a single file.
+
+  This checkpointer appends newly processed examples of an evaluation to a
+  single sequence file (e.g., `checkpoint.bagz`). This is often more efficient
+  than `PerExampleCheckpointer` when dealing with a large number of examples
+  or when file system overhead is a concern.
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -341,7 +370,12 @@ class BulkCheckpointer(Checkpointer):
 
 
 class SequenceWriter:
-  """Thread safe sequence writer."""
+  """A thread-safe writer for sequence files (e.g., Bagz).
+
+  `SequenceWriter` wraps a `pg.io.SequenceWriter` to provide thread-safe
+  `add` and `close` operations, ensuring that examples can be written
+  concurrently from multiple threads without corrupting the sequence file.
+  """
 
   def __init__(self, path: str):
     self._lock = threading.Lock()

langfun/core/eval/v2/checkpointing_test.py

@@ -65,7 +65,7 @@ class ExampleCollector(experiment_lib.Plugin):
     return self._examples
 
   def on_example_complete(
-      self, runner: runners_lib.Runner,
+      self, runner: experiment_lib.Runner,
       experiment: experiment_lib.Experiment,
       example: example_lib.Example,
   ):

langfun/core/eval/v2/eval_test_helper.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 """Helper classes and functions for evaluation tests."""
 
+import threading
+import time
+
 from langfun.core import language_model
 from langfun.core import llms
 from langfun.core import message as message_lib
@@ -47,6 +50,8 @@ class TestLLM(llms.Fake):
 
   offset: int = 0
 
+  __test__ = False
+
   def _response_from(self, prompt: message_lib.Message) -> message_lib.Message:
     return message_lib.AIMessage(
         str(prompt.metadata.x + prompt.metadata.y + self.offset)
@@ -63,6 +68,8 @@ class TestEvaluation(Evaluation):
   metrics = [metrics_lib.Match()]
   lm: language_model.LanguageModel = TestLLM()
 
+  __test__ = False
+
   def process(self, example):
     v = example.input
     if v.x == 5:
@@ -75,7 +82,7 @@ class TestEvaluation(Evaluation):
 
 class BadJsonConvertible(pg.Object):
 
-  def to_json(self, *args, **kwargs):
+  def sym_jsonify(self, *args, **kwargs):
     raise ValueError('Cannot convert to JSON.')
 
 
@@ -84,6 +91,8 @@ class TestEvaluationWithExampleCheckpointingError(TestEvaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
+  __test__ = False
+
   def process(self, example):
     return 1, dict(
         x=BadJsonConvertible()
@@ -101,6 +110,8 @@ class TestEvaluationWithExampleHtmlGenerationError(Evaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
+  __test__ = False
+
   def process(self, example):
     return 1, dict(
         x=BadHtmlConvertible()
@@ -110,6 +121,8 @@ class TestEvaluationWithExampleHtmlGenerationError(Evaluation):
 class TestEvaluationWithIndexHtmlGenerationError(TestEvaluation):
   """Test evaluation class with bad index HTML generation."""
 
+  __test__ = False
+
   def _html_tree_view(self, *args, **kwargs):
     raise ValueError('Cannot render HTML.')
 
@@ -135,3 +148,86 @@ def test_experiment_with_example_html_generation_error():
 def test_experiment_with_index_html_generation_error():
   """Returns a test experiment with bad index HTML."""
   return TestEvaluationWithIndexHtmlGenerationError()
+
+
+class TestPlugin(experiment_lib.Plugin):
+  """Plugin for testing."""
+
+  started_experiments: list[experiment_lib.Experiment] = []
+  completed_experiments: list[experiment_lib.Experiment] = []
+  skipped_experiments: list[experiment_lib.Experiment] = []
+  started_example_ids: list[int] = []
+  completed_example_ids: list[int] = []
+  start_time: float | None = None
+  complete_time: float | None = None
+
+  __test__ = False
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._lock = threading.Lock()
+
+  def on_run_start(
+      self,
+      runner: experiment_lib.Runner,
+      root: experiment_lib.Experiment
+  ) -> None:
+    del root
+    with pg.notify_on_change(False), pg.allow_writable_accessors(True):
+      self.start_time = time.time()
+
+  def on_run_complete(
+      self,
+      runner: experiment_lib.Runner,
+      root: experiment_lib.Experiment
+  ) -> None:
+    del root
+    with pg.notify_on_change(False), pg.allow_writable_accessors(True):
+      self.complete_time = time.time()
+
+  def on_experiment_start(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.started_experiments.append(pg.Ref(experiment))
+
+  def on_experiment_skipped(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.skipped_experiments.append(pg.Ref(experiment))
+
+  def on_experiment_complete(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment
+  ) -> None:
+    del runner
+    with pg.notify_on_change(False), self._lock:
+      self.completed_experiments.append(pg.Ref(experiment))
+
+  def on_example_start(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment,
+      example: Example
+  ) -> None:
+    del runner, experiment
+    with pg.notify_on_change(False), self._lock:
+      self.started_example_ids.append(example.id)
+
+  def on_example_complete(
+      self,
+      runner: experiment_lib.Runner,
+      experiment: experiment_lib.Experiment,
+      example: Example
+  ) -> None:
+    del runner, experiment
+    with pg.notify_on_change(False), self._lock:
+      self.completed_example_ids.append(example.id)

langfun/core/eval/v2/evaluation.py

@@ -32,17 +32,63 @@ import pyglove as pg
 
 
 class Evaluation(experiment_lib.Experiment):
-  """Evaluation.
-
-  An evaluation can be a leaf node or a container of other evaluations,
-  depending on whether the current evaluation object is configured with
-  any `pg.oneof`.
-
-  For example, `MyEval(lm=pg.oneof([lf.llms.Gpt4(), lf.llms.Gemini1_5Pro()]))`
-  is a container of two sub-experiments, one for each LLM. In such case, the
-  evaluation object with `pg.oneof` is called a hyper evaluation, which
-  represents a search space of evaluations, and each sub-evaluation is called
-  a leaf evaluation, which will perform the actual evaluation.
+  """Base class for Langfun evaluations.
+
+  `lf.eval.Evaluation` is the base class for defining evaluation tasks in
+  Langfun. Users typically subclass it to implement custom evaluation logic by
+  overriding `inputs` and `process` methods.
+
+  An `Evaluation` object encapsulates:
+
+  *   **`inputs`**: A callable that returns an iterable of input examples to be
+      processed. This is usually provided by implementing an `inputs(self)`
+      method in the subclass, which yields input items for evaluation one by
+      one.
+  *   **`process(self, example)`**: An abstract method that processes one
+      example and returns the output, or a tuple of (output, metadata).
+      The output will be used for computing metrics.
+  *   **`metrics`**: A list of metrics (e.g., `lf.metrics.Accuracy`) to compute
+      based on the outputs from `process`. Some metrics may require users to
+      implement a `ground_truth(self, example)` method in the subclass to
+      compute metrics against ground truth.
+  *   **Hyperparameters**: Any other attributes of the class serve as
+      hyperparameters for the evaluation (e.g., the language model to use).
+
+  **Running Evaluations:**
+
+  Evaluations are executed via `lf.eval.Suite` or by calling the `.run()`
+  method on an `Evaluation` instance, which returns a `Run` object
+  containing the evaluation run information and results. If an evaluation
+  contains sweeable parameters (using `pg.oneof`), `.run()` will expand it
+  into multiple evaluation sub-tasks -- one for each combination of
+  hyperparameters -- all managed within the same `Run`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class MyEval(lf.eval.Evaluation):
+    lm: lf.LanguageModel
+    prompt: str = '1 + 1 = '
+
+    def inputs(self):
+      yield 2
+
+    def process(self, example: lf.eval.Example):
+      return int(lf.query(self.prompt, lm=self.lm))
+
+    def ground_truth(self, example: lf.eval.Example) -> int:
+      return example.input
+
+  # Run evaluation using two different LMs
+  evaluation = MyEval(
+      lm=pg.oneof([lf.llms.Gpt4(), lf.llms.Gemini()]),
+      metrics=[lf.metrics.Accuracy()]
+  )
+  run_info = evaluation.run()
+  ```
   """
 
   inputs: Annotated[
@@ -126,6 +172,20 @@ class Evaluation(experiment_lib.Experiment):
   # Evaluation logics.
   #
 
+  def setup(self) -> None:
+    """Sets up resources required by the evaluation.
+
+    Subclasses should always call the super().setup() method to ensure the
+    proper initialization of the evaluation.
+    """
+
+  def teardown(self) -> None:
+    """Tears down resources used by the evaluation.
+
+    Subclasses should always call the super().teardown() method to ensure the
+    proper cleanup of the evaluation.
+    """
+
   @abc.abstractmethod
   def process(
       self,
@@ -137,7 +197,7 @@ class Evaluation(experiment_lib.Experiment):
 
     Args:
       example: An example object to process. `example.input` is an object
-        returned from `Evaluable.inputs`.
+        yielded from `inputs()` method.
 
     Returns:
       A processed output. Or a tuple of (output, metadata).
@@ -150,6 +210,7 @@ class Evaluation(experiment_lib.Experiment):
       example: example_lib.Example | int,
       raise_if_has_error: bool = False,
       reevaluate_upon_previous_errors: bool = True,
+      force_recompute_metrics: bool = False
   ) -> example_lib.Example:
     """Evaluates a single example input.
 
@@ -158,6 +219,8 @@ class Evaluation(experiment_lib.Experiment):
       raise_if_has_error: Whether to raise an error if the example has error.
       reevaluate_upon_previous_errors: Whether to reevaluate the example if
         the previous checkpointed run has error.
+      force_recompute_metrics: If True, force recompute the metrics even if
+        metric metadata is already present from previous checkpoint.
 
     Returns:
       The evaluated example with the output and metric metadata populated.
@@ -206,6 +269,7 @@ class Evaluation(experiment_lib.Experiment):
         # Use the output and metadata obtained from the previous processing.
         example.output = checkpointed.output
         example.metadata = checkpointed.metadata
+        example.metric_metadata = checkpointed.metric_metadata
         example.error = checkpointed.error
         example.newly_processed = False
         example.execution_status = checkpointed.execution_status
@@ -225,8 +289,16 @@ class Evaluation(experiment_lib.Experiment):
         self.info(f'Starting metric computation for example {example.id}.')
         metric_metadata = {}
         for metric in self.metrics:
-          metric_metadata.update(metric.audit(example))
-        example.metric_metadata = metric_metadata
+          metric_metadata[metric.name] = metric.update(
+              example, force_recompute=force_recompute_metrics
+          )
+
+        if example.metric_metadata is None:
+          example.metric_metadata = metric_metadata
+        else:
+          # Accumulate the metric metadata as there might be existing metadata
+          # from previous metric computation runs.
+          example.metric_metadata.update(metric_metadata)
         self.info(f'Completed metric computation for example {example.id}.')
 
     # For previously processed examples, we keep the execution status for the
@@ -287,7 +359,7 @@ class Evaluation(experiment_lib.Experiment):
       A unique string representing the resource required.
     """
     return {
-        v.resource_id for _, v in self.sym_init_args.items()
+        v.resource_id for _, v in self.sym_init_args.sym_items()
         if isinstance(v, lf.LanguageModel)
     }
 
@@ -760,7 +832,7 @@ class Evaluation(experiment_lib.Experiment):
 
 
 class EvaluationState:
-  """Evaluation state."""
+  """In-memory state of an evaluation."""
 
   class ExampleStatus(pg.Object):
     """Example state."""

langfun/core/eval/v2/evaluation_test.py

@@ -88,7 +88,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.output, 6)
     self.assertIsNone(example.error)
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(match=True))
+    self.assertEqual(example.metric_metadata, dict(match=dict(is_correct=True)))
     self.assertIsNotNone(example.usage_summary)
     self.assertGreater(example.usage_summary.total.total_tokens, 0)
     self.assertEqual(example.usage_summary.total.num_requests, 1)
@@ -103,7 +103,10 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.output, 7)
     self.assertIsNone(example.error)
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(mismatch=True))
+    self.assertEqual(
+        example.metric_metadata,
+        dict(match=dict(is_correct=False))
+    )
 
     with self.assertRaisesRegex(ValueError, 'x should not be 5'):
       _ = exp.evaluate(6, raise_if_has_error=True)
@@ -113,7 +116,10 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(pg.MISSING_VALUE, example.output)
     self.assertEqual(example.error.tag, 'ValueError')
     self.assertEqual(example.metadata, {})
-    self.assertEqual(example.metric_metadata, dict(error='ValueError'))
+    self.assertEqual(
+        example.metric_metadata,
+        dict(match=dict(error='ValueError'))
+    )
 
   def test_evaluate_withstate(self):
     eval_dir = os.path.join(tempfile.mkdtemp(), 'test_eval')

langfun/core/eval/v2/example.py

@@ -22,19 +22,30 @@ import pyglove as pg
 
 @dataclasses.dataclass
 class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
-  """An item for the evaluation.
+  """An example for evaluation.
+
+  An evaluation example contains the input and output of an evaluation task,
+  as well as metadata about the evaluation process, such as execution time,
+  LLM usage, and metric results.
 
   Attributes:
-    id: The 1-based ID of the item in the evaluation set.
-    input: An element returned from the `Evaluable.inputs` functor.
-    output: The output of the `process` method. If `pg.MISSING_VALUE`, it has
-      not been processed yet.
-    metadata: The metadata of the item produced by the `process` method.
-    metric_metadata: The dictionary returned from `Metric.audit`.
-    start_time: The start time of the evaluation item.
-    end_time: The end time of the evaluation item.
-    usage_summary: The summary of LLM usages of the evaluation item.
-    execution_status: The timeit status of the evaluation item.
+    id: The 1-based ID of the example in the evaluation set.
+    input: An element returned from the `Evaluable.inputs` functor, which serves
+      as the input for `lf.Evaluable.process`.
+    output: The output of `lf.Evaluable.process` method. If `pg.MISSING_VALUE`,
+      it indicates the example has not been processed yet.
+    error: The error raised from `lf.Evaluable.process`. If None, it
+      indicates the process was successful.
+    metadata: The metadata of the example produced by `lf.Evaluable.process`.
+    metric_metadata: The dictionary returned from `Metric.audit`, which contains
+      metadata about metric computation for this example.
+    newly_processed: Whether this example is processed in the current run. If
+      False, it indicates the example was loaded from a checkpoint from previous
+      runs.
+    start_time: The start time of processing this example.
+    end_time: The end time of processing this example.
+    usage_summary: The summary of LLM usages for processing this example.
+    execution_status: The timeit status of processing this example.
   """
   id: int
   input: Any = pg.MISSING_VALUE
@@ -49,14 +60,6 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
   usage_summary: lf.UsageSummary | None = None
   execution_status: dict[str, pg.utils.TimeIt.Status] | None = None
 
-  def __post_init__(self):
-    if self.execution_status is not None:
-      for status in self.execution_status.values():
-        if status.has_error:
-          assert isinstance(status.error, pg.ErrorInfo)
-          self.error = status.error
-          break
-
   @property
   def is_processed(self) -> bool:
     """Returns whether the item has been processed."""
@@ -182,15 +185,23 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
     extra_flags = extra_flags or {}
     num_examples = extra_flags.get('num_examples', None)
 
-    def _metric_metadata_badge(key, value):
-      if isinstance(value, bool) and bool:
-        text = key
-      else:
-        text = f'{key}:{value}'
-      return pg.views.html.controls.Badge(
-          text,
-          css_classes=[pg.utils.camel_to_snake(key, '-')],
-      )
+    def _metric_label_group(metric_metadata: dict[str, Any] | None):
+      """Renders a label group for metric metadata."""
+      badges = []
+      if metric_metadata:
+        for metric_name, metadata in metric_metadata.items():
+          assert isinstance(metadata, dict), (metric_name, metadata)
+          for k, v in metadata.items():
+            css_class = k
+            if isinstance(v, bool):
+              css_class += '_true' if v else '_false'
+            badge = pg.views.html.controls.Badge(
+                f'{k}:{v}',
+                tooltip=f'{metric_name}: {k}',
+                css_classes=[css_class],
+            )
+            badges.append(badge)
+      return pg.views.html.controls.LabelGroup(badges)
 
     def _render_header():
       return pg.Html.element(
@@ -229,12 +240,7 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
                           extra_flags=dict(as_badge=True)
                       ) if self.usage_summary is not None else None,
                       # Metric metadata.
-                      pg.views.html.controls.LabelGroup(
-                          [   # pylint: disable=g-long-ternary
-                              _metric_metadata_badge(k, v)
-                              for k, v in self.metric_metadata.items()
-                          ] if self.metric_metadata else []
-                      ),
+                      _metric_label_group(self.metric_metadata)
                   ],
                   css_classes=['example-container'],
               )
@@ -305,18 +311,18 @@ class Example(pg.JSONConvertible, pg.views.HtmlTreeView.Extension):
           color: black;
         }
         /* Badge styles. */
-        .eval-example .badge.match {
+        .eval-example .badge.is_correct_true {
           color: green;
           background-color: #dcefbe;
         }
+        .eval-example .badge.is_correct_false {
+          color: orange;
+          background-color: #ffefc4;
+        }
         .eval-example .badge.error {
           color: red;
           background-color: #fdcccc;
         }
-        .eval-example .badge.mismatch {
-          color: orange;
-          background-color: #ffefc4;
-        }
         .eval-example .badge.score {
           color: blue;
           background-color: #c4dced;

langfun/core/eval/v2/example_test.py

@@ -32,9 +32,9 @@ class ExampleTest(unittest.TestCase):
             name='evaluation', elapse=1.0, error=error
         )
     })
-    self.assertEqual(ex.error, error)
+    self.assertIsNone(ex.error)
     self.assertFalse(ex.is_processed)
-    self.assertTrue(ex.has_error)
+    self.assertFalse(ex.has_error)
     self.assertEqual(ex.elapse, 1.0)
 
     ex = Example(id=2, output=1)
@@ -116,7 +116,7 @@ class ExampleTest(unittest.TestCase):
         input=pg.Dict(a=1, b=2),
         output=3,
         metadata=dict(sum=3),
-        metric_metadata=dict(match=True),
+        metric_metadata=dict(match=dict(match=True)),
     )
     self.assertNotIn(
         'next',