langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__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

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Checkpointing evaluation runs."""
 import abc
+import datetime
 import re
 import threading
 import traceback
@@ -29,12 +30,32 @@ 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'
+  ] = 'checkpoint.jsonl'
+
+  enable_inprogress_file: Annotated[
+      bool,
+      'If True, write file "<example_id>.inprogress" when example gets started.'
+  ] = True
+
+  max_ckpt_loading_threads: Annotated[
+      int,
+      'Max number of workers for loading checkpoint files at startup.'
+  ] = 128
 
   def on_experiment_start(
       self,
@@ -75,6 +96,24 @@ class Checkpointer(experiment_lib.Plugin):
           f'scratch. Example IDs: {example_ids_to_evaluate}.'
       )
 
+  def on_example_start(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    if self.enable_inprogress_file:
+      def _save_inprogress_file(example: Example):
+        inprogress_file = runner.current_run.output_path_for(
+            experiment, f'{example.id}.inprogress'
+        )
+        pg.io.writefile(
+            inprogress_file,
+            datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        )
+      runner.background_run(_save_inprogress_file, example)
+
   def on_example_complete(
       self,
       runner: Runner,
@@ -149,7 +188,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 +212,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 +282,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 +394,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,
   ):
@@ -90,7 +90,10 @@ class PerExampleCheckpointerTest(CheckpointerTest):
     root_dir = os.path.join(tempfile.mkdtemp(), 'per_example_checkpointer')
     experiment = eval_test_helper.test_experiment()
     checkpoint_filename = 'checkpoint.jsonl'
-    checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    checkpointer = checkpointing.PerExampleCheckpointer(
+        checkpoint_filename,
+        enable_inprogress_file=True
+    )
     collector = ExampleCollector()
     run = experiment.run(
         root_dir, 'new', runner='sequential', plugins=[checkpointer, collector]
@@ -102,6 +105,10 @@ class PerExampleCheckpointerTest(CheckpointerTest):
         example = collector.examples[i + 1]
         ckpt = run.output_path_for(leaf, f'checkpoint_{example.id}.jsonl')
         self.assertTrue(pg.io.path_exists(ckpt))
+        inprogress_file = run.output_path_for(
+            leaf, f'{example.id}.inprogress'
+        )
+        self.assertTrue(pg.io.path_exists(inprogress_file))
         with pg.io.open_sequence(ckpt) as f:
           examples_from_ckpt = list(iter(f))
           # `eval_test_helper.test_experiment` has two TestEvaluation with

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:
@@ -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,15 +121,22 @@ 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.')
 
 
+def test_evaluation(offset: int | pg.hyper.OneOf = 0):
+  """Returns a test evaluation."""
+  return TestEvaluation(lm=TestLLM(offset=offset))
+
+
 def test_experiment():
   """Returns a test experiment."""
   return Suite([
-      TestEvaluation(lm=TestLLM(offset=0)),
-      TestEvaluation(lm=TestLLM(offset=pg.oneof(range(5)))),
+      test_evaluation(),
+      test_evaluation(pg.oneof(range(5))),
   ])
 
 
@@ -135,3 +153,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
@@ -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."""
@@ -808,8 +880,9 @@ class EvaluationState:
       load_example_metadata: bool | Callable[
           [example_lib.Example], bool] = True,
       filter: Callable[[example_lib.Example], bool] | None = None,  # pylint: disable=redefined-builtin
-  ) -> None:
+  ) -> list[example_lib.Example]:
     """Loads the state from the example sequence file."""
+    examples = []
     for example in example_lib.Example.iter_ckpts(
         state_file,
         example_input_by_id=example_input_by_id,
@@ -819,6 +892,8 @@ class EvaluationState:
         continue
       example.newly_processed = False
       self._ckpt_examples[example.id] = example
+      examples.append(example)
+    return examples
 
   @property
   def evaluation_status(self) -> dict[int, ExampleStatus]:

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')