langfun 0.1.2.dev202411090804__py3-none-any.whl → 0.1.2.dev202411140804__py3-none-any.whl

langfun/core/console.py

@@ -59,12 +59,20 @@ def under_notebook() -> bool:
   return bool(_notebook)
 
 
-def display(value: Any, clear: bool = False) -> None:  # pylint: disable=redefined-outer-name
+def display(value: Any, clear: bool = False) -> Any:  # pylint: disable=redefined-outer-name
   """Displays object in current notebook cell."""
   if _notebook is not None:
     if clear:
       _notebook.clear_output()
-    _notebook.display(value)
+    return _notebook.display(value)
+  return None
+
+
+def run_script(javascript: str) -> Any:
+  """Runs JavaScript in current notebook cell."""
+  if _notebook is not None:
+    return _notebook.display(_notebook.Javascript(javascript))
+  return
 
 
 def clear() -> None:

langfun/core/console_test.py

@@ -18,6 +18,7 @@ import io
 import unittest
 
 from langfun.core import console
+import pyglove as pg
 
 
 class ConsoleTest(unittest.TestCase):
@@ -32,6 +33,22 @@ class ConsoleTest(unittest.TestCase):
 
   def test_under_notebook(self):
     self.assertFalse(console.under_notebook())
+    console._notebook = True
+    self.assertTrue(console.under_notebook())
+    console._notebook = None
+
+  def test_notebook_interaction(self):
+    console._notebook = pg.Dict(
+        display=lambda x: x, Javascript=lambda x: x, clear_output=lambda: None)
+    self.assertEqual(console.display('hi', clear=True), 'hi')
+    self.assertEqual(
+        console.run_script('console.log("hi")'),
+        'console.log("hi")'
+    )
+    console.clear()
+    console._notebook = None
+    self.assertIsNone(console.display('hi'))
+    self.assertIsNone(console.run_script('console.log("hi")'))
 
 
 if __name__ == '__main__':

langfun/core/eval/__init__.py

@@ -16,6 +16,8 @@
 # pylint: disable=g-importing-member
 # pylint: disable=g-bad-import-order
 
+from langfun.core.eval import v2
+
 from langfun.core.eval.base import register
 from langfun.core.eval.base import registered_names
 from langfun.core.eval.base import get_evaluations

langfun/core/eval/v2/__init__.py

@@ -0,0 +1,38 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""langfun eval framework v2."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-bad-import-order
+from langfun.core.eval.v2.experiment import Experiment
+from langfun.core.eval.v2.experiment import Suite
+from langfun.core.eval.v2.evaluation import Evaluation
+
+from langfun.core.eval.v2.example import Example
+from langfun.core.eval.v2.progress import Progress
+
+from langfun.core.eval.v2.metric_values import MetricValue
+from langfun.core.eval.v2.metric_values import Rate
+from langfun.core.eval.v2.metric_values import Average
+from langfun.core.eval.v2.metrics import Metric
+from langfun.core.eval.v2 import metrics
+
+from langfun.core.eval.v2.experiment import Plugin
+
+from langfun.core.eval.v2.experiment import Runner
+from langfun.core.eval.v2 import runners
+
+
+# pylint: enable=g-bad-import-order
+# pylint: enable=g-importing-member

langfun/core/eval/v2/checkpointing.py

@@ -0,0 +1,135 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Checkpointing evaluation runs."""
+import threading
+
+from langfun.core.eval.v2 import example as example_lib
+from langfun.core.eval.v2 import experiment as experiment_lib
+import pyglove as pg
+
+Example = example_lib.Example
+Experiment = experiment_lib.Experiment
+Runner = experiment_lib.Runner
+
+
+class Checkpointer(experiment_lib.Plugin):
+  """Plugin for checkpointing evaluation runs."""
+
+  checkpoint_filename: str = 'checkpoint.bagz'
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._lock = threading.Lock()
+    self._state_writer = None
+
+  def on_run_start(
+      self,
+      runner: Runner,
+      root: Experiment,
+  ) -> None:
+    self._state_writer = {}
+
+  def on_run_abort(
+      self,
+      runner: Runner,
+      root: Experiment,
+      error: BaseException
+  ) -> None:
+    with self._lock:
+      if self._state_writer is not None:
+        self._state_writer.clear()
+
+  def on_run_complete(
+      self,
+      runner: Runner,
+      root: Experiment,
+  ) -> None:
+    with self._lock:
+      assert self._state_writer is not None and not self._state_writer
+
+  def on_experiment_start(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Creates the checkpoint file."""
+    if not experiment.is_leaf:
+      return
+    # For refresh runs, we don't want to load the previous state.
+    if not runner.current_run.refresh:
+      experiment.load_state(
+          runner.current_run.input_path_for(
+              experiment, self.checkpoint_filename
+          ),
+          raise_if_not_exist=False
+      )
+    state_writer = StateWriter(
+        runner.current_run.output_path_for(
+            experiment, self.checkpoint_filename
+        )
+    )
+    with self._lock:
+      if self._state_writer is not None:
+        self._state_writer[experiment.id] = state_writer
+
+  def on_experiment_complete(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+  ) -> None:
+    """Closes the checkpoint file."""
+    if not experiment.is_leaf:
+      return
+    assert experiment.id in self._state_writer
+    with self._lock:
+      if self._state_writer is not None:
+        del self._state_writer[experiment.id]
+
+  def on_example_complete(
+      self,
+      runner: Runner,
+      experiment: Experiment,
+      example: Example,
+  ) -> None:
+    """Saves the example to the checkpoint file."""
+    assert experiment.id in self._state_writer
+    if not example.has_error:
+      runner.background_run(self._state_writer[experiment.id].add, example)
+
+
+class StateWriter:
+  """Thread safe state writer."""
+
+  def __init__(self, path: str):
+    self._lock = threading.Lock()
+    self._sequence_writer = pg.io.open_sequence(path, 'w')
+
+  def add(self, example: Example):
+    example_blob = pg.to_json_str(
+        example,
+        hide_default_values=True,
+        save_ref_value=True,
+        exclude_input=True
+    )
+    with self._lock:
+      if self._sequence_writer is None:
+        return
+      self._sequence_writer.add(example_blob)
+
+  def __del__(self):
+    # Make sure there is no write in progress.
+    with self._lock:
+      assert self._sequence_writer is not None
+      self._sequence_writer.close()
+      self._sequence_writer = None

langfun/core/eval/v2/checkpointing_test.py

@@ -0,0 +1,89 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import tempfile
+import unittest
+
+from langfun.core.eval.v2 import checkpointing
+from langfun.core.eval.v2 import example as example_lib
+from langfun.core.eval.v2 import runners as runners_lib  # pylint: disable=unused-import
+from langfun.core.eval.v2 import test_helper
+import pyglove as pg
+
+Example = example_lib.Example
+
+
+class StateWriterTest(unittest.TestCase):
+
+  def test_basic(self):
+    file = os.path.join(tempfile.gettempdir(), 'test.jsonl')
+    writer = checkpointing.StateWriter(file)
+    example = Example(id=1, input=pg.Dict(x=1), output=2)
+    writer.add(example)
+    del writer
+    self.assertTrue(pg.io.path_exists(file))
+
+  def test_error_handling(self):
+    file = os.path.join(tempfile.gettempdir(), 'test_error_handling.jsonl')
+    writer = checkpointing.StateWriter(file)
+    writer.add(Example(id=1, input=pg.Dict(x=1), output=2))
+
+    def f():
+      raise ValueError('Intentional error')
+
+    try:
+      writer.add(f())
+    except ValueError:
+      del writer
+
+    self.assertTrue(pg.io.path_exists(file))
+    with pg.io.open_sequence(file, 'r') as f:
+      self.assertEqual(len(list(iter(f))), 1)
+
+
+class CheckpointingTest(unittest.TestCase):
+
+  def test_checkpointing(self):
+    root_dir = os.path.join(tempfile.gettempdir(), 'test_checkpointing')
+    experiment = test_helper.test_experiment()
+    checkpoint_filename = 'checkpoint.jsonl'
+    checkpointer = checkpointing.Checkpointer(checkpoint_filename)
+    run = experiment.run(
+        root_dir, 'new', runner='sequential', plugins=[checkpointer]
+    )
+    self.assertEqual(len(checkpointer._state_writer), 0)
+    num_processed = {}
+    for leaf in experiment.leaf_nodes:
+      ckpt = run.output_path_for(leaf, checkpoint_filename)
+      self.assertTrue(pg.io.path_exists(ckpt))
+      with pg.io.open_sequence(ckpt) as f:
+        self.assertEqual(
+            len(list(iter(f))),
+            leaf.progress.num_completed - leaf.progress.num_failed
+        )
+      if leaf.id not in num_processed:
+        self.assertEqual(leaf.progress.num_skipped, 0)
+        num_processed[leaf.id] = leaf.progress.num_processed
+
+    # Run again, should skip existing.
+    _ = experiment.run(
+        root_dir, 'latest', runner='sequential', plugins=[checkpointer]
+    )
+    self.assertEqual(len(checkpointer._state_writer), 0)
+    for leaf in experiment.leaf_nodes:
+      self.assertEqual(leaf.progress.num_skipped, num_processed[leaf.id])
+
+
+if __name__ == '__main__':
+  unittest.main()