langfun 0.0.2.dev20240422__py3-none-any.whl → 0.0.2.dev20240425__py3-none-any.whl

langfun/__init__.py

@@ -55,6 +55,7 @@ Video = modalities.Video
 PDF = modalities.PDF
 
 # Error types.
+MappingError = structured.MappingError
 SchemaError = structured.SchemaError
 JsonError = structured.JsonError
 CodeError = coding.CodeError

langfun/core/component.py

@@ -210,6 +210,12 @@ def get_contextual_override(var_name: str) -> ContextualOverride | None:
   return _get_scoped_value(_global_tls, _CONTEXT_OVERRIDES, var_name)
 
 
+def all_contextual_values() -> dict[str, Any]:
+  """Returns all contextual values provided from `lf.context` in scope."""
+  overrides = getattr(_global_tls, _CONTEXT_OVERRIDES, {})
+  return {k: v.value for k, v in overrides.items()}
+
+
 @contextlib.contextmanager
 def _contextual_scope(
     tls: threading.local, tls_key, **variables

langfun/core/component_test.py

@@ -84,6 +84,7 @@ class ComponentContextTest(unittest.TestCase):
           lf.get_contextual_override('y'),
           lf.ContextualOverride(3, cascade=False, override_attrs=False),
       )
+      self.assertEqual(lf.all_contextual_values(), dict(x=3, y=3, z=3))
 
       # Member attributes take precedence over `lf.context`.
       self.assertEqual(a1.x, 1)

langfun/core/eval/__init__.py

@@ -16,6 +16,8 @@
 # pylint: disable=g-importing-member
 # pylint: disable=g-bad-import-order
 
+from langfun.core.eval.base import app_run
+
 from langfun.core.eval.base import Evaluable
 from langfun.core.eval.base import Evaluation
 from langfun.core.eval.base import Suite

langfun/core/eval/base.py

@@ -26,6 +26,8 @@ import threading
 import time
 from typing import Annotated, Any, Callable, Iterator, Literal, Optional, Sequence, Type, Union
 
+from absl import app
+from absl import flags
 import langfun.core as lf
 import langfun.core.coding as lf_coding
 from langfun.core.llms.cache import in_memory
@@ -546,6 +548,16 @@ class Evaluable(lf.Component):
         )
         s.write(pg.format(m.result))
         s.write('</div>')
+      if 'usage' in m.metadata:
+        s.write(
+            '<div style="background-color: #EEEEEE; color: black; '
+            'white-space: pre-wrap; padding: 10px; border: 0px solid; '
+            'margin: 10px">'
+            f'prompt: {m.usage.prompt_tokens} tokens, '
+            f'response: {m.usage.completion_tokens} tokens, '
+            f'total: {m.usage.total_tokens} tokens'
+            '</div>'
+        )
       s.write('</div>')
 
   @classmethod
@@ -810,6 +822,30 @@ class Evaluation(Evaluable):
       return 0.0
     return self.num_failures / self.num_completed
 
+  @property
+  def has_usage(self) -> bool:
+    """Returns True if token usage is enabled."""
+    return self._num_usages > 0
+
+  @property
+  def average_prompt_tokens(self) -> int:
+    """Returns the average prompt tokens."""
+    if not self.has_usage:
+      return 0
+    return self._total_prompt_tokens // self._num_usages
+
+  @property
+  def average_completion_tokens(self) -> int:
+    """Returns the average completion tokens."""
+    if not self.has_usage:
+      return 0
+    return self._total_completion_tokens // self._num_usages
+
+  @property
+  def average_total_tokens(self) -> int:
+    """Returns the average total tokens."""
+    return self.average_prompt_tokens + self.average_completion_tokens
+
   @functools.cached_property
   def schema(self) -> lf_structured.Schema | None:
     """Schema."""
@@ -856,7 +892,11 @@ class Evaluation(Evaluable):
             'Encountered: {annotation!r}.'
         )
       self._maybe_adjust_schema_for_completion(annotation)
-    return lf_structured.Schema.from_value(annotation)
+    schema = lf_structured.Schema.from_value(annotation)
+    # NOTE(daiyip): add references to the dependent classes of the returned type
+    # to prevent unused subclasses get garbage collected by Python.
+    setattr(schema, '__dependencies__', schema.class_dependencies())
+    return schema
 
   def _maybe_adjust_schema_for_completion(self, cls):
     if (self.completion_prompt_field is None
@@ -933,6 +973,10 @@ class Evaluation(Evaluable):
     self._failures = []
     self._num_completed = 0
 
+    self._total_prompt_tokens = 0
+    self._total_completion_tokens = 0
+    self._num_usages = 0
+
   @property
   def failures_link(self) -> str | None:
     """Returns the link to the failures page."""
@@ -952,7 +996,7 @@ class Evaluation(Evaluable):
     example = example or self.examples[0]
 
     # We make a copy to avoid pollute the state of current object.
-    copy = self.clone()
+    copy: Evaluation = self.clone()
     copy.__dict__['examples'] = [example]
 
     # We set the symbolic parent of the cloned to access contextual information
@@ -982,9 +1026,9 @@ class Evaluation(Evaluable):
           color='blue',
       )
 
-    # Audit the result.
-    copy.audit(example, output, output_message)
+    copy.audit(example, output_message, None, dryrun=True)
     result = copy.summarize()
+
     if verbose:
       lf.console.write('')
       lf.console.write(
@@ -1031,11 +1075,12 @@ class Evaluation(Evaluable):
             status_fn=self._status,
         ):
           if error is not None:
-            self._failures.append((example, str(error)))
-          else:
-            output = message.text if self.schema is None else message.result
-            self.audit(example, output, message)
-          self._num_completed += 1
+            message = (
+                error.lm_response
+                if isinstance(error, lf_structured.MappingError)
+                else None
+            )
+          self.audit(example, message, error)
       finally:
         # Save cache upon completion or interruption.
         if self.dir and self.cache:
@@ -1138,6 +1183,19 @@ class Evaluation(Evaluable):
       )
     else:
       cache_stats = dict(use_cache=False)
+
+    if self.has_usage:
+      usage = pg.Dict(
+          total_prompt_tokens=self._total_prompt_tokens,
+          total_completion_tokens=self._total_completion_tokens,
+          num_usages=self._num_usages,
+          average_prompt_tokens=self.average_prompt_tokens,
+          average_completion_tokens=self.average_completion_tokens,
+          average_total_tokens=self.average_total_tokens,
+      )
+    else:
+      usage = None
+
     result = pg.Dict(
         experiment_setup=pg.Dict(
             id=self.id,
@@ -1153,6 +1211,7 @@ class Evaluation(Evaluable):
             failures=self.num_failures,
             failure_rate=self.failure_rate,
         ),
+        usage=usage,
     )
     return result
 
@@ -1174,9 +1233,28 @@ class Evaluation(Evaluable):
           '</td></tr><tr><td>'
       )
       self._render_metric(s)
+
+      # Summarize average usage.
+      if self.result.usage is not None:
+        self._render_usage(s)
+
     s.write('</td></tr></table></div>')
     return s.getvalue()
 
+  def _render_usage(self, s: io.StringIO) -> None:
+    """Renders usage in HTML."""
+    usage = self.result.usage
+    total = usage.total_prompt_tokens + usage.total_completion_tokens
+    s.write(
+        '&nbsp;<a title="'
+        f'# of usages: {usage.num_usages}&#013;'
+        f'total prompt: {usage.total_prompt_tokens}&#013;'
+        f'total response: {usage.total_completion_tokens}&#013;'
+        f'avg prompt: {usage.average_prompt_tokens}&#013;'
+        f'avg response: {usage.average_completion_tokens}'
+        f'" style="color:gray">({total} tokens)</a>'
+    )
+
   def _render_metric(self, s: io.StringIO) -> None:
     """Renders metrics in HTML."""
     assert self.result is not None
@@ -1191,17 +1269,48 @@ class Evaluation(Evaluable):
         )
     )
 
-  def audit(self, example: Any, output: Any, message: lf.Message) -> None:
+  def audit(
+      self,
+      example: Any,
+      message: lf.Message | None,
+      error: Exception | None = None,
+      dryrun: bool = False,
+  ) -> None:
     """Audits the example against the output. Subclasses should override.
 
     Args:
       example: The input object.
-      output: The output from LM. For `lf.call`, if `schema_fn` is not provided,
-        it will be the raw LM response string. Otherwise it will be the
-        structured output from the LM.
       message: The entire message returned by the LM, which could be used to
-        trace the LM input, response and parsed structure.
+        trace the LM input, response and parsed structure. If error is raised
+        before LLM could return a response, None will be its value.
+      error: The exception during processing the example.
+      dryrun: Whether or not audition takes place during dryrun.
     """
+    if error is not None:
+      self._failures.append((example, str(error)))
+      if isinstance(error, lf_structured.MappingError):
+        message = error.lm_response
+    else:
+      assert message is not None
+      output = message.text if self.schema is None else message.result
+      self.audit_processed(example, output, message, dryrun=dryrun)
+
+    # Audit usage.
+    if message is not None:
+      self.audit_usage(message, dryrun=dryrun)
+    self._num_completed += 1
+
+  def audit_usage(self, message: lf.Message, dryrun: bool = False) -> None:
+    for m in message.trace():
+      if 'usage' in m.metadata:
+        self._total_prompt_tokens += m.usage.prompt_tokens
+        self._total_completion_tokens += m.usage.completion_tokens
+        self._num_usages += 1
+
+  def audit_processed(
+      self, example: Any, output: Any, message: lf.Message, dryrun: bool = False
+  ) -> None:
+    """Audits a successfully processed example. Subclass should override."""
 
   def save(
       self, definition: bool = True, result: bool = True, report: bool = True
@@ -1245,8 +1354,10 @@ class Evaluation(Evaluable):
         '<td>Prompt</td>'
         '<td>Schema</td>'
         '<td>Additional Args</td>'
-        '<td>Failures</td>'
     )
+    if self.result.usage is not None:
+      s.write('<td>Usage</td>')
+    s.write('<td>Failures</td>')
 
   def _render_result_row(self, s: io.StringIO) -> None:
     s.write(
@@ -1271,6 +1382,12 @@ class Evaluation(Evaluable):
         '<td style="color:purple" '
         f'{_html_repr(self.additional_args, compact=False)}</td>'
     )
+    # Usage.
+    if self.result.usage is not None:
+      s.write('<td>')
+      self._render_usage(s)
+      s.write('</td>')
+
     # Failures.
     s.write(
         '<td><span style="color:orange">%s</span>%s</td>'
@@ -1369,8 +1486,8 @@ class Summary(pg.Object):
           Type[lf.LanguageModel],
           tuple[lf.LanguageModel | Type[lf.LanguageModel], ...],
       ] = lf.LanguageModel,
-      method: Union[str, tuple[str], None] = None,
-      schema_fn: Union[pg.Functor, tuple[pg.Functor], None] = None,
+      method: Union[str, tuple[str, ...], None] = None,
+      schema_fn: Union[pg.Functor, tuple[pg.Functor, ...], None] = None,
       completed: bool | None = None,
       pivot_field: str | None = None,
   ) -> 'Summary':
@@ -1564,6 +1681,7 @@ class Summary(pg.Object):
       for entry in self.select(task=task).evaluations:
         results.append(
             pg.Dict(
+                id=entry.id,
                 experiment=entry,
                 dir=entry.dir,
                 metrics=entry.result.metrics if entry.result else None,
@@ -1789,3 +1907,43 @@ def monitor_async(
       scan_interval=scan_interval,
       refresh_when_stop=refresh_when_stop,
   )
+
+
+def app_run(target: Evaluable):
+  """Runs the target evaluation as an absl app.
+
+  Args:
+    target: An Langfun evaluable object.
+  """
+  flags.DEFINE_string(
+      'root_dir', None, 'Root directory for running the evaluation.'
+  )
+
+  flags.DEFINE_bool(
+      'dryrun', False, 'If True, dryrun the experiment instead of running it.'
+  )
+
+  flags.DEFINE_bool(
+      'debug', False, 'If True, output prompt and response to the console.'
+  )
+
+  flags.DEFINE_bool(
+      'rerun',
+      False,
+      'If True, rerun the experiment even a cached result is found.',
+  )
+
+  FLAGS = flags.FLAGS  # pylint: disable=invalid-name
+
+  def _main(argv):
+    if len(argv) > 1:
+      raise app.UsageError('Too many command-line arguments.')
+
+    if FLAGS.root_dir:
+      target.rebind(root_dir=FLAGS.root_dir, raise_on_no_change=False)
+    if FLAGS.dryrun:
+      target.dryrun(debug=FLAGS.debug)
+    else:
+      target.run(debug=FLAGS.debug, rerun=FLAGS.rerun)
+
+  app.run(_main)

langfun/core/eval/base_test.py

@@ -101,7 +101,7 @@ class EvaluationTest(unittest.TestCase):
     self.assertEqual(s.dir, os.path.join(s.root_dir, s.id))
     self.assertEqual(s.hash, s.clone().hash)
     # Test persistent hash.
-    self.assertEqual(s.hash, '436dc80c')
+    self.assertEqual(s.hash, 'ae86c703')
     self.assertEqual(
         s.hash, s.clone(override={'max_workers': 2, 'lm.timeout': 20}).hash
     )
@@ -210,7 +210,7 @@ class EvaluationTest(unittest.TestCase):
         s.result,
         dict(
             experiment_setup=dict(
-                id='Evaluation@f1aa5126',
+                id='Evaluation@0fade07d',
                 dir=s.dir,
                 model='StaticSequence',
                 prompt_template='{{example.question}}',
@@ -221,6 +221,14 @@ class EvaluationTest(unittest.TestCase):
                 use_cache=True, num_queries=2, num_hits=0, num_updates=2
             ),
             metrics=dict(total=2, failures=1, failure_rate=0.5),
+            usage=dict(
+                total_prompt_tokens=774,
+                total_completion_tokens=25,
+                num_usages=2,
+                average_prompt_tokens=387,
+                average_completion_tokens=12,
+                average_total_tokens=399,
+            ),
         ),
     )
     self.assertTrue(
@@ -285,8 +293,11 @@ class EvaluationTest(unittest.TestCase):
     s = eval_set(
         'run_filter_test', pg.oneof(['call', 'query']),
         schema_fn=answer_schema(), lm=lm)
+    result = s.run(
+        filter=lambda x: x.method == 'query', dryrun=True, summary=False
+    )
     self.assertEqual(
-        s.run(filter=lambda x: x.method == 'query', dryrun=True, summary=False),
+        result,
         {
             s.children[0].id: None,
             s.children[1].id: dict(
@@ -302,7 +313,8 @@ class EvaluationTest(unittest.TestCase):
                     use_cache=True, num_queries=2, num_hits=0, num_updates=2
                 ),
                 metrics=dict(total=2, failures=0, failure_rate=0.0),
-            )
+                usage=s.children[1].result.usage,
+            ),
         },
     )
 
@@ -336,7 +348,6 @@ class EvaluationTest(unittest.TestCase):
 
     summary = s.run(verbose=True)
     self.assertEqual(len(summary.evaluations), 2)
-
     self.assertEqual(
         s.result,
         {
@@ -353,6 +364,7 @@ class EvaluationTest(unittest.TestCase):
                     use_cache=True, num_queries=2, num_hits=0, num_updates=2
                 ),
                 metrics=dict(total=2, failures=1, failure_rate=0.5),
+                usage=s.children[0].result.usage,
             ),
             s.children[1].id: dict(
                 experiment_setup=dict(
@@ -367,6 +379,7 @@ class EvaluationTest(unittest.TestCase):
                     use_cache=True, num_queries=2, num_hits=0, num_updates=2
                 ),
                 metrics=dict(total=2, failures=1, failure_rate=0.5),
+                usage=s.children[1].result.usage,
             ),
         },
     )
@@ -459,7 +472,7 @@ class SuiteTest(unittest.TestCase):
         lm=lm
     )
     # Test for persistent hash.
-    self.assertEqual(s.hash, 'bbfdc7a8')
+    self.assertEqual(s.hash, '26e6cc25')
     s.run()
     expected = {
         s.children[0].id: dict(
@@ -475,6 +488,7 @@ class SuiteTest(unittest.TestCase):
                 use_cache=True, num_queries=2, num_hits=0, num_updates=2
             ),
             metrics=dict(total=2, failures=1, failure_rate=0.5),
+            usage=s.children[0].result.usage,
         ),
         s.children[1].id: {
             s.children[1]
@@ -492,6 +506,7 @@ class SuiteTest(unittest.TestCase):
                     use_cache=True, num_queries=4, num_hits=1, num_updates=3
                 ),
                 metrics=dict(total=2, failures=2, failure_rate=1.0),
+                usage=s.children[1].children[0].result.usage,
             ),
             s.children[1]
             .children[2]
@@ -511,6 +526,7 @@ class SuiteTest(unittest.TestCase):
                     num_updates=2,
                 ),
                 metrics=dict(total=2, failures=1, failure_rate=0.5),
+                usage=s.children[1].children[2].result.usage,
             ),
         },
     }
@@ -682,5 +698,17 @@ class SummaryTest(unittest.TestCase):
     self.assertTrue(pg.io.path_exists(summary_file))
 
 
+class AppRunTest(unittest.TestCase):
+
+  def test_app_run(self):
+    lm = fake.StaticSequence(['two', 'Solution(final_answer=2)'])
+    try:
+      base.app_run(
+          eval_set('app_run_test', 'query', schema_fn=answer_schema(), lm=lm)
+      )
+    except SystemExit:
+      pass
+
+
 if __name__ == '__main__':
   unittest.main()

langfun/core/eval/matching.py

@@ -86,9 +86,26 @@ class Matching(base.Evaluation):
     self._matches = []
     self._mismatches = []
 
-  def audit(self, example: Any, output: Any, message: lf.Message) -> None:
+  def audit_processed(
+      self, example: Any, output: Any, message: lf.Message, dryrun: bool = False
+  ) -> None:
     groundtruth = self.groundtruth(example)
     answer = self.answer(output, example)
+
+    if dryrun:
+      lf.console.write('')
+      lf.console.write(
+          str(groundtruth),
+          title='GROUDTRUTH',
+          color='green',
+      )
+      lf.console.write('')
+      lf.console.write(
+          str(answer),
+          title='ANSWER',
+          color='blue',
+      )
+
     if self.match(answer, groundtruth):
       self._matches.append((example, output, message))
     else:

langfun/core/eval/matching_test.py

@@ -103,7 +103,7 @@ class MatchingTest(unittest.TestCase):
         s.result,
         dict(
             experiment_setup=dict(
-                id='MyTask@acd56a61',
+                id='MyTask@739a174b',
                 dir=s.dir,
                 model='StaticSequence',
                 prompt_template='{{example.question}}',
@@ -125,6 +125,7 @@ class MatchingTest(unittest.TestCase):
                 num_mismatches=1,
                 mismatch_rate=0.25,
             ),
+            usage=s.result.usage,
         ),
     )
     self.assertTrue(

langfun/core/eval/scoring.py

@@ -61,8 +61,18 @@ class Scoring(base.Evaluation):
     super()._reset()
     self._scored = []
 
-  def audit(self, example: Any, output: Any, message: lf.Message) -> None:
+  def audit_processed(
+      self, example: Any, output: Any, message: lf.Message, dryrun: bool = False
+  ) -> None:
     score = self.score(example, output)
+
+    if dryrun:
+      lf.console.write('')
+      lf.console.write(
+          str(score),
+          title='SCORE',
+          color='blue',
+      )
     self._scored.append((example, output, score, message))
 
   @abc.abstractmethod

langfun/core/eval/scoring_test.py

@@ -81,7 +81,7 @@ class ScoringTest(unittest.TestCase):
         s.result,
         dict(
             experiment_setup=dict(
-                id='ConstraintFollowing@a44d8b89',
+                id='ConstraintFollowing@5c88a5eb',
                 dir=s.dir,
                 model='StaticSequence',
                 prompt_template='{{example}}',
@@ -102,6 +102,7 @@ class ScoringTest(unittest.TestCase):
                 score_rate=1.0,
                 avg_score=0.5,
             ),
+            usage=s.result.usage,
         ),
     )
     self.assertTrue(

langfun/core/language_model.py

@@ -24,6 +24,9 @@ from langfun.core import console
 from langfun.core import message as message_lib
 import pyglove as pg
 
+TOKENS_PER_REQUEST = 250  # Estimated num tokens for a single request
+DEFAULT_MAX_CONCURRENCY = 1  # Use this as max concurrency if no RPM or TPM data
+
 
 class LMSample(pg.Object):
   """Response candidate."""
@@ -604,3 +607,14 @@ class LanguageModel(component.Component):
             f'score: {r.score}',
             color='blue',
         )
+
+  def rate_to_max_concurrency(
+      self, requests_per_min: float = 0, tokens_per_min: float = 0
+  ) -> int:
+    """Converts a rate to a max concurrency."""
+    if tokens_per_min > 0:
+      return max(int(tokens_per_min / TOKENS_PER_REQUEST / 60), 1)
+    elif requests_per_min > 0:
+      return max(int(requests_per_min / 60), 1)  # Max concurrency can't be zero
+    else:
+      return DEFAULT_MAX_CONCURRENCY  # Default of 1

langfun/core/language_model_test.py

@@ -394,6 +394,38 @@ class LanguageModelTest(unittest.TestCase):
     with self.assertRaises(NotImplementedError):
       MockModel().score('hi', ['1', '2'])
 
+  def test_rate_to_max_concurrency_no_rpm_no_tpm(self) -> None:
+    lm = MockModel()
+    self.assertEqual(
+        lm_lib.DEFAULT_MAX_CONCURRENCY,
+        lm.rate_to_max_concurrency(requests_per_min=0, tokens_per_min=0),
+    )
+    self.assertEqual(
+        lm_lib.DEFAULT_MAX_CONCURRENCY,
+        lm.rate_to_max_concurrency(requests_per_min=-1, tokens_per_min=-1),
+    )
+
+  def test_rate_to_max_concurrency_only_rpm_specified_uses_rpm(self) -> None:
+    lm = MockModel()
+    test_rpm = 1e4
+    self.assertEqual(
+        lm.rate_to_max_concurrency(requests_per_min=test_rpm),
+        int(test_rpm / 60)
+    )
+
+  def test_rate_to_max_concurrency_tpm_specified_uses_tpm(self) -> None:
+    lm = MockModel()
+    test_tpm = 1e7
+    self.assertEqual(
+        lm.rate_to_max_concurrency(requests_per_min=1, tokens_per_min=test_tpm),
+        int(test_tpm / lm_lib.TOKENS_PER_REQUEST / 60),
+    )
+
+  def test_rate_to_max_concurrency_small_rate_returns_one(self) -> None:
+    lm = MockModel()
+    self.assertEqual(lm.rate_to_max_concurrency(requests_per_min=1), 1)
+    self.assertEqual(lm.rate_to_max_concurrency(tokens_per_min=1), 1)
+
 
 if __name__ == '__main__':
   unittest.main()