langfun 0.0.2.dev20240330__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/structured/schema_test.py

@@ -18,6 +18,7 @@ import inspect
 import typing
 import unittest
 
+import langfun.core as lf
 from langfun.core.llms import fake
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
@@ -192,9 +193,9 @@ class SchemaTest(unittest.TestCase):
     self.assertEqual(schema.parse('{"result": 1}'), 1)
     schema = schema_lib.Schema(dict[str, int])
     self.assertEqual(
-        schema.parse(
-            '{"result": {"_type": "Unknown", "x": 1}}}', force_dict=True),
-        dict(x=1))
+        schema.parse('{"result": {"x": 1}}}'),
+        dict(x=1)
+    )
     with self.assertRaisesRegex(
         schema_lib.SchemaError, 'Expect .* but encountered .*'):
       schema.parse('{"result": "def"}')
@@ -260,15 +261,18 @@ class ClassDependenciesTest(unittest.TestCase):
     class A(pg.Object):
       foo: tuple[Foo, int]
 
+    class B(pg.Object):
+      pass
+
     class X(pg.Object):
-      k: int
+      k: dict[str, B]
 
-    class B(A):
+    class C(A):
       bar: Bar
       foo2: Foo | X
 
     a = A(foo=(Foo(1), 0))
-    self.assertEqual(schema_lib.class_dependencies(a), [Foo, A, Bar, X, B])
+    self.assertEqual(schema_lib.class_dependencies(a), [Foo, A, Bar, B, X, C])
 
     self.assertEqual(schema_lib.class_dependencies(1), [])
 
@@ -280,9 +284,10 @@ class SchemaPythonReprTest(unittest.TestCase):
       value_spec: pg.typing.ValueSpec,
       expected_annotation: str,
       strict: bool = False,
+      **kwargs,
   ) -> None:
     self.assertEqual(
-        schema_lib.annotation(value_spec, strict=strict),
+        schema_lib.annotation(value_spec, strict=strict, **kwargs),
         expected_annotation,
     )
 
@@ -358,11 +363,27 @@ class SchemaPythonReprTest(unittest.TestCase):
     self.assert_annotation(
         pg.typing.Object(Activity).noneable(), 'Activity | None'
     )
+    self.assert_annotation(
+        pg.typing.Object(Activity).noneable(), 'Activity | None',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.Object(Activity).noneable(), 'Any | None',
+        allowed_dependencies=set(),
+    )
 
     # List.
     self.assert_annotation(
         pg.typing.List(pg.typing.Object(Activity)), 'list[Activity]'
     )
+    self.assert_annotation(
+        pg.typing.List(pg.typing.Object(Activity)), 'list[Activity]',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.List(pg.typing.Object(Activity)), 'list[Any]',
+        allowed_dependencies=set(),
+    )
     self.assert_annotation(
         pg.typing.List(pg.typing.Object(Activity)).noneable(),
         'list[Activity] | None',
@@ -374,16 +395,35 @@ class SchemaPythonReprTest(unittest.TestCase):
 
     # Tuple.
     self.assert_annotation(
-        pg.typing.Tuple([pg.typing.Int(), pg.typing.Str()]), 'tuple[int, str]'
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Activity, str]'
+    )
+    self.assert_annotation(
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Activity, str]',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Any, str]',
+        allowed_dependencies=set(),
     )
     self.assert_annotation(
-        pg.typing.Tuple([pg.typing.Int(), pg.typing.Str()]).noneable(),
-        'tuple[int, str] | None',
+        pg.typing.Tuple([Activity, pg.typing.Str()]).noneable(),
+        'tuple[Activity, str] | None',
     )
 
     # Dict.
     self.assert_annotation(
-        pg.typing.Dict({'x': int, 'y': str}), '{\'x\': int, \'y\': str}'
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Activity, \'y\': str}'
+    )
+    self.assert_annotation(
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Activity, \'y\': str}',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Any, \'y\': str}',
+        allowed_dependencies=set(),
     )
     self.assert_annotation(
         pg.typing.Dict({'x': int, 'y': str}),
@@ -395,6 +435,15 @@ class SchemaPythonReprTest(unittest.TestCase):
         'dict[str, Any]',
         strict=False,
     )
+
+    class DictValue(pg.Object):
+      pass
+
+    self.assert_annotation(
+        pg.typing.Dict([(pg.typing.StrKey(), DictValue)]),
+        'dict[str, DictValue]',
+        strict=False,
+    )
     self.assert_annotation(
         pg.typing.Dict(),
         'dict[str, Any]',
@@ -408,6 +457,13 @@ class SchemaPythonReprTest(unittest.TestCase):
         ).noneable(),
         'Union[Activity, Itinerary, None]',
     )
+    self.assert_annotation(
+        pg.typing.Union(
+            [pg.typing.Object(Activity), pg.typing.Object(Itinerary)]
+        ).noneable(),
+        'Union[Activity, Any, None]',
+        allowed_dependencies=set([Activity]),
+    )
 
     # Any.
     self.assert_annotation(pg.typing.Any(), 'Any')
@@ -415,13 +471,13 @@ class SchemaPythonReprTest(unittest.TestCase):
 
   def test_class_definition(self):
     self.assertEqual(
-        schema_lib.class_definition(Activity),
+        schema_lib.class_definition(Activity, allowed_dependencies=set()),
         'class Activity:\n  description: str\n',
     )
     self.assertEqual(
         schema_lib.class_definition(Itinerary),
         inspect.cleandoc("""
-            class Itinerary:
+            class Itinerary(Object):
               \"\"\"A travel itinerary for a day.\"\"\"
               day: int(min=1)
               type: Literal['daytime', 'nighttime']
@@ -431,7 +487,9 @@ class SchemaPythonReprTest(unittest.TestCase):
             """) + '\n',
     )
     self.assertEqual(
-        schema_lib.class_definition(PlaceOfInterest),
+        schema_lib.class_definition(
+            PlaceOfInterest, allowed_dependencies=set()
+        ),
         inspect.cleandoc("""
             class PlaceOfInterest:
               \"\"\"The name of a place of interest.
@@ -447,11 +505,11 @@ class SchemaPythonReprTest(unittest.TestCase):
       pass
 
     self.assertEqual(
-        schema_lib.class_definition(A),
+        schema_lib.class_definition(A, allowed_dependencies=set()),
         'class A:\n  pass\n',
     )
     self.assertEqual(
-        schema_lib.class_definition(A, include_pg_object_as_base=True),
+        schema_lib.class_definition(A),
         'class A(Object):\n  pass\n',
     )
 
@@ -459,9 +517,27 @@ class SchemaPythonReprTest(unittest.TestCase):
       x: str
       __kwargs__: typing.Any
 
-    with self.assertRaisesRegex(
-        TypeError, 'Variable-length keyword arguments is not supported'):
-      schema_lib.class_definition(C)
+    self.assertEqual(
+        schema_lib.class_definition(C), 'class C(Object):\n  x: str\n'
+    )
+
+    class D(pg.Object):
+      x: str
+      @schema_lib.include_method_in_prompt
+      def __call__(self, y: int) -> int:
+        return len(self.x) + y
+
+    self.assertEqual(
+        schema_lib.class_definition(D),
+        inspect.cleandoc(
+            """
+            class D(Object):
+              x: str
+
+              def __call__(self, y: int) -> int:
+                return len(self.x) + y
+            """) + '\n'
+    )
 
   def test_repr(self):
     class Foo(pg.Object):
@@ -479,10 +555,21 @@ class SchemaPythonReprTest(unittest.TestCase):
     class A(pg.Object):
       foo: Foo
 
+      @schema_lib.include_method_in_prompt
+      def foo_value(self) -> int:
+        return self.foo.x
+
+      def baz_value(self) -> str:
+        return 'baz'
+
     class B(A):
       bar: Bar
       foo2: Foo
 
+      @schema_lib.include_method_in_prompt
+      def bar_value(self) -> str:
+        return self.bar.y
+
     schema = schema_lib.Schema([B])
     self.assertEqual(
         schema_lib.SchemaPythonRepr().class_definitions(schema),
@@ -490,9 +577,6 @@ class SchemaPythonReprTest(unittest.TestCase):
             class Foo:
               x: int
 
-            class A:
-              foo: Foo
-
             class Bar:
               """Class Bar."""
               y: str
@@ -501,10 +585,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             ''') + '\n',
     )
 
@@ -521,9 +611,6 @@ class SchemaPythonReprTest(unittest.TestCase):
             class Foo:
               x: int
 
-            class A:
-              foo: Foo
-
             class Bar:
               """Class Bar."""
               y: str
@@ -532,10 +619,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             ```
             '''),
     )
@@ -543,16 +636,12 @@ class SchemaPythonReprTest(unittest.TestCase):
         schema_lib.SchemaPythonRepr().repr(
             schema,
             include_result_definition=False,
-            include_pg_object_as_base=True,
             markdown=False,
         ),
         inspect.cleandoc('''
-            class Foo(Object):
+            class Foo:
               x: int
 
-            class A(Object):
-              foo: Foo
-
             class Bar:
               """Class Bar."""
               y: str
@@ -561,10 +650,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             '''),
     )
 
@@ -600,6 +695,10 @@ class ValuePythonReprTest(unittest.TestCase):
         schema_lib.ValuePythonRepr().repr(1, schema_lib.Schema(int)),
         '```python\n1\n```'
     )
+    self.assertEqual(
+        schema_lib.ValuePythonRepr().repr(lf.Template('hi, {{a}}', a='foo')),
+        'hi, foo'
+    )
     self.assertEqual(
         schema_lib.ValuePythonRepr().repr(
             A([Foo(1), Foo(2)], 'bar'), schema_lib.Schema(A), markdown=False,
@@ -612,7 +711,7 @@ class ValuePythonReprTest(unittest.TestCase):
             ```python
             class Foo(Object):
               x: int
-            
+
             class A(Object):
               foo: list[Foo]
               y: str | None

langfun/core/structured/scoring.py

@@ -17,13 +17,13 @@ from typing import Any, Type, Union
 
 import langfun.core as lf
 from langfun.core.structured import mapping
-from langfun.core.structured import prompting
+from langfun.core.structured import querying
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
 def score(
-    prompt: Union[str, pg.Symbolic],
+    prompt: Union[str, pg.Symbolic] | list[str | pg.Symbolic],
     completions: list[str | pg.Symbolic],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
@@ -32,9 +32,58 @@ def score(
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
     protocol: schema_lib.SchemaProtocol = 'python',
+    return_scoring_results: bool = False,
     **kwargs,
-) -> list[float]:
-  """Scores the outputs based on the prompt."""
+) -> list[float] | list[lf.LMScoringResult]:
+  """Scores the outputs based on the prompt.
+
+  Examples:
+    ```
+    # Example 1: Scoring text output based on the user prompt.
+    scores = lf.score('{{x}} + {{y}} =', ['1', '2', '3'], lm=lm, x=1, y=2)
+    assert len(scores) == 3
+
+    # Example 2: Scoring int output based on the formulated OOP prompt.
+    scores = lf.score('1 + 1 =', [1, 2, 3], lm=lm)
+    assert len(scores) == 3
+
+    class Answer(pg.Object):
+      result: int
+
+    # Example 3: Scoring object output based on the formulated OOP prompt.
+    scores = lf.score('1 + 1 =', [Answer(1), Answer(2), Answer(3)], lm=lm)
+    assert len(scores) == 3
+
+    # Example 4: Scoring object field value based on the formulated OOP prompt
+    # and the generated tokens before the first `pg.oneof`.
+    scores = lf.score('1 + 1 =', [Answer(pg.oneof([1, 2, 3]))], lm=lm)
+    assert len(scores) == 3
+
+    # Example 5: Scoring multiple prompt/completion pairs.
+    scores = lf.score(
+        ['1 + 1=', '2 + 3='],
+        ['2', '4'],
+        lm=lm
+    )
+    assert len(scores) == 2
+    ```
+
+  Args:
+    prompt: The prompt(s) based on which each completion will be scored.
+    completions: A list of strings or symbolic objects as the output.
+    schema: The schema as the output type. If None, it will be inferred from
+      the completions.
+    lm: The language model used for scoring.
+    examples: Fewshot exemplars used together with the prompt in getting the
+      completions.
+    protocol: The protocol for formulating the prompt based on objects.
+    return_scoring_results: If True, returns a list of `lf.LMScoringResult`,
+      otherwise returns a list of floats as the scores of each completion.
+    **kwargs: Keyword arguments that are referred by the prompt.
+
+  Returns:
+    A list of floats or `lf.LMScoringResult` as the score of each completion.
+  """
   if not completions:
     raise ValueError('`completions` must not be empty.')
 
@@ -48,28 +97,85 @@ def score(
             f'{[type(c) for c in completions]}.'
         )
 
-  input_message = prompting.query(
-      prompt,
-      schema,
-      examples=examples,
-      protocol=protocol,
-      skip_lm=True,
-      returns_message=True,
-      **kwargs,
-  )
+  if isinstance(prompt, list):
+    prompts = []
+    for p in prompt:
+      prompts.append(
+          querying.query_prompt(
+              p,
+              schema,
+              examples=examples,
+              protocol=protocol,
+              **kwargs,
+          )
+      )
+    input_message = prompts
+  else:
+    input_message = querying.query_prompt(
+        prompt,
+        schema,
+        examples=examples,
+        protocol=protocol,
+        **kwargs,
+    )
   if lm is None:
     lm_override = lf.get_contextual_override('lm')
     if lm_override is None:
       raise ValueError('`lm` must be specified or provided from `lf.context`.')
     lm = lm_override.value
 
+  completion_reprs = []
+  for c in completions:
+    if isinstance(c, mapping.MappingError):
+      completion_reprs.append(c.lm_response)
+    else:
+      rep = mapping.MappingExample.value_repr(
+          c, protocol=protocol, compact=False, verbose=False
+      )
+
+      # NOTE(daiyip): supporting scenario of scoring object field with
+      # `pg.oneof`.
+      oneof_pos = rep.find('OneOf(')
+      if oneof_pos == -1:
+        completion_reprs.append(rep)
+      else:
+        assert protocol == 'python', protocol
+        if isinstance(input_message, list):
+          raise ValueError(
+              'Scoring on object fields using `pg.oneof` must share the '
+              f'same prompt. Encountered: {prompt}'
+          )
+        input_message.text += '\n' + rep[:oneof_pos]
+        oneof = _get_first_oneof(c)
+        for v in oneof.candidates:
+          completion_reprs.append(
+              pg.format(
+                  v,
+                  python_format=True,
+                  compact=False,
+                  verbose=False,
+                  root_indent=oneof.sym_path.depth
+              )
+          )
+
   results = lm.score(
       input_message,
-      [
-          mapping.MappingExample.value_repr(
-              c, protocol=protocol, compact=False, verbose=False
-          )
-          for c in completions
-      ],
+      completion_reprs,
   )
+  if return_scoring_results:
+    return results
   return [r.score for r in results]
+
+
+def _get_first_oneof(value: Any) -> pg.hyper.OneOf:
+  """Gets the first pg.oneof from a symbolic object."""
+  oneofs = []
+  def select_oneofs(k, v, p):
+    del k, p
+    if isinstance(v, pg.hyper.OneOf):
+      oneofs.append(v)
+      return pg.TraverseAction.CONTINUE
+    return pg.TraverseAction.ENTER
+  pg.traverse(value, select_oneofs)
+  assert oneofs
+  return oneofs[0]

langfun/core/structured/scoring_test.py

@@ -16,6 +16,11 @@ import unittest
 import langfun.core as lf
 from langfun.core.llms import fake
 from langfun.core.structured import scoring
+import pyglove as pg
+
+
+class Answer(pg.Object):
+  result: int
 
 
 class ScoringTest(unittest.TestCase):
@@ -32,9 +37,34 @@ class ScoringTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, '`lm` must be specified'):
       scoring.score('hi', [1, 2])
 
+    with self.assertRaisesRegex(
+        ValueError,
+        'Scoring on object fields using `pg.oneof` must share the same prompt',
+    ):
+      scoring.score(
+          ['1 + 1=', '2 + 3='],
+          [Answer(pg.oneof([1, 2, 3]))],
+          lm=fake.Echo(),
+      )
+
   def test_score(self):
     self.assertEqual(scoring.score('hi', [1, 2], lm=fake.Echo()), [0.0, -1.0])
 
+  def test_score_on_field_values(self):
+    self.assertEqual(
+        scoring.score(
+            '1 + 1=',
+            [Answer(pg.oneof([1, 2, 3]))], lm=fake.Echo()
+        ),
+        [0.0, -1.0, -2.0]
+    )
+
+  def test_score_returning_scoring_results(self):
+    self.assertEqual(scoring.score(
+        'hi', [1, 2], lm=fake.Echo(), return_scoring_results=True),
+                     [lf.LMScoringResult(score=0.0, gradients=None),
+                      lf.LMScoringResult(score=-1.0, gradients=None)])
+
   def test_scope_with_lm_from_the_context(self):
     with lf.context(lm=fake.Echo()):
       self.assertEqual(scoring.score('hi', [1, 2]), [0.0, -1.0])

langfun/core/structured/tokenization.py

@@ -0,0 +1,64 @@
+# Copyright 2023 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.
+"""Tokenize the prompt for `lf.query`."""
+
+from typing import Any, Type, Union
+
+import langfun.core as lf
+from langfun.core.structured import mapping
+from langfun.core.structured import querying
+from langfun.core.structured import schema as schema_lib
+import pyglove as pg
+
+
+def tokenize(
+    prompt: Union[str, pg.Symbolic] | list[str | pg.Symbolic],
+    schema: Union[
+        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
+    ] = None,
+    *,
+    lm: lf.LanguageModel | None = None,
+    examples: list[mapping.MappingExample] | None = None,
+    protocol: schema_lib.SchemaProtocol = 'python',
+    **kwargs,
+) -> list[tuple[str | bytes, int]]:
+  """Tokenize the prompt for `lf.query`.
+
+  Args:
+    prompt: The prompt(s) based on which each completion will be scored.
+    schema: The schema as the output type. If None, it will be inferred from
+      the completions.
+    lm: The language model used for scoring.
+    examples: Fewshot exemplars used together with the prompt in getting the
+      completions.
+    protocol: The protocol for formulating the prompt based on objects.
+    **kwargs: Keyword arguments that are referred by the prompt.
+
+  Returns:
+    A list of (text, token_id) tuples.
+  """
+  input_message = querying.query_prompt(
+      prompt,
+      schema,
+      examples=examples,
+      protocol=protocol,
+      **kwargs,
+  )
+  if lm is None:
+    lm_override = lf.get_contextual_override('lm')
+    if lm_override is None:
+      raise ValueError('`lm` must be specified or provided from `lf.context`.')
+    lm = lm_override.value
+
+  return lm.tokenize(input_message)

langfun/core/structured/tokenization_test.py

@@ -0,0 +1,48 @@
+# Copyright 2023 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 unittest
+import langfun.core as lf
+from langfun.core.llms import fake
+from langfun.core.structured import tokenization
+import pyglove as pg
+
+
+class Answer(pg.Object):
+  result: int
+
+
+class TokenizationTest(unittest.TestCase):
+
+  def test_bad_call(self):
+
+    with self.assertRaisesRegex(ValueError, '`lm` must be specified'):
+      tokenization.tokenize('hi')
+
+  def test_tokenize(self):
+    self.assertEqual(
+        tokenization.tokenize('hi', lm=fake.Echo()),
+        [('hi', 0)]
+    )
+
+  def test_tokenize_with_lm_from_the_context(self):
+    with lf.context(lm=fake.Echo()):
+      self.assertEqual(
+          tokenization.tokenize('hi'),
+          [('hi', 0)]
+      )
+
+
+if __name__ == '__main__':
+  unittest.main()