langfun 0.1.2.dev202503110804__tar.gz → 0.1.2.dev202503130804__tar.gz

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202503110804
+Version: 0.1.2.dev202503130804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/agentic/action.py

@@ -674,6 +674,10 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   """Session for performing an agentic task."""
 
   root: ActionInvocation = ActionInvocation(RootAction())
+  id: Annotated[
+      str | None,
+      'An optional identifier for the sessin, which will be used for logging.'
+  ] = None
 
   def _on_bound(self):
     super()._on_bound()

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/agentic/action_eval.py

@@ -31,9 +31,10 @@ class ActionEval(lf.eval.v2.Evaluation):
       'Arguments to call the action.'
   ] = {}
 
-  def process(self, example: pg.Dict) -> tuple[str, dict[str, Any]]:
-    action = example.action
-    session = action_lib.Session()
+  def process(self, example: lf.eval.v2.Example) -> tuple[str, dict[str, Any]]:
+    example_input = example.input
+    action = example_input.action
+    session = action_lib.Session(id=str(example.id))
     with lf.logging.use_log_level('fatal'):
       action(session=session, **self.action_args)
     return session.final_result, dict(session=session)
@@ -68,7 +69,7 @@ class ActionEvalV1(lf_eval.Matching):
 
   def process(self, example: pg.Dict, **kwargs):
     action = example.action
-    session = action_lib.Session()
+    session = action_lib.Session(id=str(getattr(example, 'id', '<empty>')))
     action(session=session, lm=self.lm, **kwargs)
     return session.as_message()
 

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/checkpointing.py

@@ -53,14 +53,14 @@ class Checkpointer(experiment_lib.Plugin):
       self._load_experiment(runner, experiment)
 
     example_ids_to_evaluate = current_run.examples_to_evaluate(experiment)
-    if experiment.state.evaluated_examples:
+    if experiment.state.ckpt_examples:
       loaded_example_ids = list(
-          sorted(experiment.state.evaluated_examples.keys())
+          sorted(experiment.state.ckpt_examples.keys())
       )
       example_ids_to_evaluate -= set(loaded_example_ids)
       example_ids_to_evaluate = list(sorted(example_ids_to_evaluate))
       experiment.info(
-          f'{len(experiment.state.evaluated_examples)} examples '
+          f'{len(experiment.state.ckpt_examples)} examples '
           'loaded from checkpoint files. Their outputs will be used '
           f'for recomputing metrics. Example IDs: {loaded_example_ids}.'
       )
@@ -316,7 +316,7 @@ class BulkCheckpointer(Checkpointer):
         writer = self._sequence_writer.pop(experiment.id)
         writer.close()
         experiment.info(
-            f'{len(experiment.state.evaluated_examples)} examples are '
+            f'{len(experiment.state.evaluation_status)} examples are '
             f'checkpointed to {writer.path}.'
         )
 

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/checkpointing_test.py

@@ -18,6 +18,7 @@ import unittest
 from langfun.core.eval.v2 import checkpointing
 from langfun.core.eval.v2 import eval_test_helper
 from langfun.core.eval.v2 import example as example_lib
+from langfun.core.eval.v2 import experiment as experiment_lib
 from langfun.core.eval.v2 import runners as runners_lib  # pylint: disable=unused-import
 import pyglove as pg
 
@@ -52,6 +53,26 @@ class SequenceWriterTest(unittest.TestCase):
       self.assertEqual(len(list(iter(f))), 1)
 
 
+class ExampleCollector(experiment_lib.Plugin):
+  """Collects all examples."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._examples = {}
+
+  @property
+  def examples(self) -> dict[int, example_lib.Example]:
+    return self._examples
+
+  def on_example_complete(
+      self, runner: runners_lib.Runner,
+      experiment: experiment_lib.Experiment,
+      example: example_lib.Example,
+  ):
+    assert experiment.is_leaf, None
+    self._examples[example.id] = example
+
+
 class CheckpointerTest(unittest.TestCase):
 
   def assert_found_in_log(self, experiment, message):
@@ -70,13 +91,15 @@ class PerExampleCheckpointerTest(CheckpointerTest):
     experiment = eval_test_helper.test_experiment()
     checkpoint_filename = 'checkpoint.jsonl'
     checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    collector = ExampleCollector()
     run = experiment.run(
-        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+        root_dir, 'new', runner='sequential', plugins=[checkpointer, collector]
     )
     num_processed = {}
     for leaf in experiment.leaf_nodes:
       for i in range(leaf.num_examples):
-        example = leaf.state.get(i + 1)
+        self.assertIn(i + 1, collector.examples)
+        example = collector.examples[i + 1]
         ckpt = run.output_path_for(leaf, f'checkpoint_{example.id}.jsonl')
         if example.has_error:
           self.assertFalse(pg.io.path_exists(ckpt))
@@ -134,12 +157,15 @@ class PerExampleCheckpointerTest(CheckpointerTest):
     experiment = eval_test_helper.TestEvaluation()
     checkpoint_filename = 'checkpoint.jsonl'
     checkpointer = checkpointing.PerExampleCheckpointer(checkpoint_filename)
+    collector = ExampleCollector()
+
     run = experiment.run(
-        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+        root_dir, 'new', runner='sequential', plugins=[checkpointer, collector]
     )
     num_processed = {}
     for i in range(experiment.num_examples):
-      example = experiment.state.get(i + 1)
+      self.assertIn(i + 1, collector.examples)
+      example = collector.examples[i + 1]
       ckpt = run.output_path_for(experiment, f'checkpoint_{example.id}.jsonl')
       if not example.has_error:
         self.assertTrue(pg.io.path_exists(ckpt))

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/eval_test_helper.py

@@ -63,7 +63,8 @@ class TestEvaluation(Evaluation):
   metrics = [metrics_lib.Match()]
   lm: language_model.LanguageModel = TestLLM()
 
-  def process(self, v):
+  def process(self, example):
+    v = example.input
     if v.x == 5:
       raise ValueError('x should not be 5')
     return structured.query(
@@ -83,7 +84,7 @@ class TestEvaluationWithExampleCheckpointingError(TestEvaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
-  def process(self, v):
+  def process(self, example):
     return 1, dict(
         x=BadJsonConvertible()
     )
@@ -100,7 +101,7 @@ class TestEvaluationWithExampleHtmlGenerationError(Evaluation):
   inputs = test_inputs()
   metrics = [metrics_lib.Match()]
 
-  def process(self, v):
+  def process(self, example):
     return 1, dict(
         x=BadHtmlConvertible()
     )

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/evaluation.py

@@ -127,13 +127,17 @@ class Evaluation(experiment_lib.Experiment):
   #
 
   @abc.abstractmethod
-  def process(self, example_input: Any) -> Any | tuple[Any, dict[str, Any]]:
+  def process(
+      self,
+      example: example_lib.Example
+  ) -> Any | tuple[Any, dict[str, Any]]:
     """Processes a single example from the evaluation set.
 
     Users should override this method to implement the evaluation logic.
 
     Args:
-      example_input: An object returned from `Evaluable.inputs`.
+      example: An example object to process. `example.input` is an object
+        returned from `Evaluable.inputs`.
 
     Returns:
       A processed output. Or a tuple of (output, metadata).
@@ -162,25 +166,24 @@ class Evaluation(experiment_lib.Experiment):
     if pg.MISSING_VALUE == example.input:
       example.input = self.example_input_by_id(example.id)
 
-    cached = self._state.get(example.id)
-
+    checkpointed = self._state.ckpt_example(example.id)
     with pg.timeit('evaluate') as timeit, lf.track_usages() as usage_summary:
-      if cached is None or cached.has_error:
+      if checkpointed is None or checkpointed.has_error:
         example.start_time = time.time()
         self._process(example, raise_if_has_error=raise_if_has_error)
       else:
-        example.start_time = cached.start_time
+        example.start_time = checkpointed.start_time
 
-        # Use cached output and metadata obtained from the previous processing.
-        example.output = cached.output
-        example.metadata = cached.metadata
+        # Use the output and metadata obtained from the previous processing.
+        example.output = checkpointed.output
+        example.metadata = checkpointed.metadata
         example.newly_processed = False
 
         # For previously processed examples, we merge previous usages as
         # cached, so the usage summary will account previous usages, but as
         # cached.
-        assert cached.usage_summary is not None
-        usage_summary.merge(cached.usage_summary, as_cached=True)
+        assert checkpointed.usage_summary is not None
+        usage_summary.merge(checkpointed.usage_summary, as_cached=True)
 
       # Recompute the metrics and metadata for the example even its processed
       # output and metadata were from the cache.
@@ -221,7 +224,7 @@ class Evaluation(experiment_lib.Experiment):
     ):
       try:
         with pg.timeit('process'):
-          output = self.process(example.input)
+          output = self.process(example)
         if (isinstance(output, tuple)
             and len(output) == 2
             and isinstance(output[1], dict)):
@@ -687,9 +690,29 @@ class Evaluation(experiment_lib.Experiment):
 class EvaluationState:
   """Evaluation state."""
 
+  class ExampleStatus(pg.Object):
+    """Example state."""
+    evaluated: Annotated[
+        bool,
+        'Whether the example is evaluated.'
+    ] = False
+
+    newly_processed: Annotated[
+        bool,
+        'Whether the example is newly processed.'
+    ] = False
+
+    has_error: Annotated[
+        bool,
+        'Whether the example has error.'
+    ] = False
+
   def __init__(self):
     super().__init__()
-    self._evaluated_examples: dict[int, example_lib.Example] = {}
+    self._ckpt_examples: dict[int, example_lib.Example] = {}
+    self._evaluation_status: dict[
+        int, EvaluationState.ExampleStatus
+    ] = {}
 
   def load(
       self,
@@ -711,17 +734,41 @@ class EvaluationState:
         assert isinstance(example, example_lib.Example), example
         if filter is not None and not filter(example):
           continue
-        self._evaluated_examples[example.id] = example
+        example.newly_processed = False
+        self._ckpt_examples[example.id] = example
 
   @property
-  def evaluated_examples(self) -> dict[int, example_lib.Example]:
-    """Returns the examples in the state."""
-    return self._evaluated_examples
+  def evaluation_status(self) -> dict[int, ExampleStatus]:
+    """Returns the evaluation status of the examples."""
+    return self._evaluation_status
 
-  def get(self, example_id: int) -> example_lib.Example | None:
-    """Returns the example with the given ID."""
-    return self._evaluated_examples.get(example_id)
+  @property
+  def ckpt_examples(self) -> dict[int, example_lib.Example]:
+    """Returns the unevaluated examples from checkpoints."""
+    return self._ckpt_examples
+
+  def ckpt_example(self, example_id: int) -> example_lib.Example | None:
+    """Returns the unevaluated example from checkpoints for a given ID."""
+    return self._ckpt_examples.get(example_id)
+
+  def get_status(self, example_id: int) -> ExampleStatus:
+    """Returns the evaluation status of the example."""
+    return self._evaluation_status.get(
+        example_id, EvaluationState.ExampleStatus()
+    )
 
   def update(self, example: example_lib.Example) -> None:
     """Updates the state with the given example."""
-    self._evaluated_examples[example.id] = example
+    self._update_status(example)
+    # Processed examples will be removed once it's done.
+    self._ckpt_examples.pop(example.id, None)
+
+  def _update_status(self, example: example_lib.Example) -> None:
+    """Updates the evaluation status of the example."""
+    self._evaluation_status[example.id] = (
+        EvaluationState.ExampleStatus(
+            evaluated=example.output != pg.MISSING_VALUE,
+            newly_processed=example.newly_processed,
+            has_error=example.has_error,
+        )
+    )

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/evaluation_test.py

@@ -78,7 +78,9 @@ class EvaluationTest(unittest.TestCase):
   def test_evaluate(self):
     exp = eval_test_helper.TestEvaluation()
     example = exp.evaluate(Example(id=3))
-    self.assertIs(exp.state.get(3), example)
+    self.assertTrue(exp.state.get_status(3).evaluated)
+    self.assertTrue(exp.state.get_status(3).newly_processed)
+    self.assertFalse(exp.state.get_status(3).has_error)
     self.assertTrue(example.newly_processed)
     self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
     self.assertEqual(example.output, 6)
@@ -111,7 +113,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(example.metadata, {})
     self.assertEqual(example.metric_metadata, dict(error='ValueError'))
 
-  def test_evaluate_with_state(self):
+  def test_evaluate_withstate(self):
     eval_dir = os.path.join(tempfile.gettempdir(), 'test_eval')
     pg.io.mkdirs(eval_dir, exist_ok=True)
     state_file = os.path.join(eval_dir, 'state.jsonl')
@@ -121,13 +123,14 @@ class EvaluationTest(unittest.TestCase):
       self.assertTrue(example.newly_processed)
       self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
       self.assertEqual(example.output, 6)
-      self.assertEqual(len(exp._state.evaluated_examples), 1)
+      self.assertEqual(len(exp.state.evaluation_status), 1)
       f.add(pg.to_json_str(example))
 
     exp.reset()
-    self.assertEqual(len(exp._state.evaluated_examples), 0)
+    self.assertEqual(len(exp.state.ckpt_examples), 0)
     exp.load_state(state_file)
-    self.assertEqual(len(exp._state.evaluated_examples), 1)
+    self.assertEqual(len(exp.state.ckpt_examples), 1)
+    self.assertEqual(len(exp.state.evaluation_status), 0)
     example = exp.evaluate(3)
     self.assertFalse(example.newly_processed)
     self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
@@ -140,14 +143,14 @@ class EvaluationTest(unittest.TestCase):
 
     # Test load_state with filter.
     exp.reset()
-    self.assertEqual(len(exp._state.evaluated_examples), 0)
+    self.assertEqual(len(exp.state.ckpt_examples), 0)
     exp.load_state(state_file, filter=lambda x: x.id == 3)
-    self.assertEqual(len(exp._state.evaluated_examples), 1)
+    self.assertEqual(len(exp.state.ckpt_examples), 1)
 
     exp.reset()
-    self.assertEqual(len(exp._state.evaluated_examples), 0)
+    self.assertEqual(len(exp.state.ckpt_examples), 0)
     exp.load_state(state_file, filter=lambda x: x.id == 1)
-    self.assertEqual(len(exp._state.evaluated_examples), 0)
+    self.assertEqual(len(exp.state.ckpt_examples), 0)
 
   def test_html_view(self):
     exp = eval_test_helper.TestEvaluation()

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun/core/eval/v2/runners.py

@@ -181,8 +181,8 @@ class RunnerBase(Runner):
     )
     num_from_checkpoint, num_processed = 0, 0
     for example_id in example_ids:
-      example = experiment.state.get(example_id)
-      if example.newly_processed:
+      status = experiment.state.get_status(example_id)
+      if status.newly_processed:
         num_processed += 1
       else:
         num_from_checkpoint += 1
@@ -358,7 +358,8 @@ class RunnerBase(Runner):
     """Runs the evaluation example."""
     self.on_example_start(evaluation, item)
     item = evaluation.evaluate(
-        item, raise_if_has_error=self.current_run.raise_if_has_error
+        item,
+        raise_if_has_error=self.current_run.raise_if_has_error,
     )
     self.on_example_complete(evaluation, item)
     return item

{langfun-0.1.2.dev202503110804 → langfun-0.1.2.dev202503130804}/langfun.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202503110804
+Version: 0.1.2.dev202503130804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors