langfun 0.1.1.dev20240822__tar.gz → 0.1.1.dev20240825__tar.gz

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240822
+Version: 0.1.1.dev20240825
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun/core/language_model.py

@@ -712,6 +712,72 @@ class LanguageModel(component.Component):
             color='blue',
         )
 
+  def tokenize(
+      self,
+      prompt: str | message_lib.Message,
+      **kwargs,
+  ) -> list[tuple[str | bytes, int]]:
+    """Tokenizes the given prompt."""
+    prompt = message_lib.UserMessage.from_value(prompt)
+    call_counter = self._call_counter
+    self._call_counter += 1
+
+    with component.context(override_attrs=True, **kwargs):
+      request_start = time.time()
+      tokens = self._tokenize(prompt)
+      elapse = time.time() - request_start
+      self._debug_tokenize(prompt, tokens, call_counter, elapse)
+      return tokens
+
+  def _tokenize(
+      self, prompt: message_lib.Message
+  ) -> list[tuple[str | bytes, int]]:
+    """Subclass to implement."""
+    raise NotImplementedError(
+        f'{self.__class__.__name__} does not support tokenization.'
+    )
+
+  def _debug_tokenize(
+      self,
+      prompt: message_lib.Message,
+      tokens: list[tuple[str | bytes, int]],
+      call_counter: int,
+      elapse: float,
+  ):
+    debug = self.debug
+    if isinstance(debug, bool):
+      debug = LMDebugMode.ALL if debug else LMDebugMode.NONE
+
+    if debug & LMDebugMode.INFO:
+      self._debug_model_info(call_counter, UsageNotAvailable())
+
+    if debug & LMDebugMode.PROMPT:
+      console.write(
+          prompt,
+          title=f'\n[{call_counter}] PROMPT TO TOKENIZE:',
+          color='green',
+      )
+      referred_modalities_lst = [prompt.referred_modalities(),]
+      if referred_modalities_lst:
+        for referred_modalities in referred_modalities_lst:
+          console.write(
+              pg.object_utils.kvlist_str(
+                  [(k, repr(v), None) for k, v in referred_modalities.items()]
+              ),
+              title=f'\n[{call_counter}] MODALITY OBJECTS SENT TO LM:',
+              color='green',
+          )
+
+    if debug & LMDebugMode.RESPONSE:
+      console.write(
+          tokens,
+          title=(
+              f'\n[{call_counter}] {len(tokens)} TOKENS RETURNED '
+              f'(in {elapse:.2f} seconds):'
+          ),
+          color='blue',
+      )
+
   def rate_to_max_concurrency(
       self, requests_per_min: float = 0, tokens_per_min: float = 0
   ) -> int:

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun/core/language_model_test.py

@@ -81,6 +81,13 @@ class MockScoringModel(MockModel):
     ]
 
 
+class MockTokenizeModel(MockModel):
+
+  def _tokenize(
+      self, prompt: message_lib.Message) -> list[tuple[str | bytes, int]]:
+    return [(w, i) for i, w in enumerate(prompt.text.split(' '))]
+
+
 class LMSamplingOptionsTest(unittest.TestCase):
   """Tests for LMSamplingOptions."""
 
@@ -552,6 +559,65 @@ class LanguageModelTest(unittest.TestCase):
     with self.assertRaises(NotImplementedError):
       MockModel().score('hi', ['1', '2'])
 
+  def test_tokenize(self):
+    info_flag = lm_lib.LMDebugMode.INFO
+    prompt_flag = lm_lib.LMDebugMode.PROMPT
+    response_flag = lm_lib.LMDebugMode.RESPONSE
+    debug_prints = {
+        info_flag: 'LM INFO',
+        prompt_flag: 'PROMPT TO TOKENIZE',
+        response_flag: 'TOKENS RETURNED',
+    }
+    debug_modes = [
+        info_flag,
+        prompt_flag,
+        response_flag,
+        info_flag | prompt_flag,
+        info_flag | response_flag,
+        prompt_flag | response_flag,
+        info_flag | prompt_flag | response_flag,
+    ]
+
+    class Image(modality.Modality):
+      def to_bytes(self):
+        return b'fake_image'
+
+    for debug_mode in debug_modes:
+      string_io = io.StringIO()
+      lm = MockTokenizeModel()
+
+      with contextlib.redirect_stdout(string_io):
+        self.assertEqual(
+            lm.tokenize(
+                message_lib.UserMessage('hi <<[[image]]>>', image=Image()),
+                debug=debug_mode),
+            [('hi', 0), ('<<[[image]]>>', 1)],
+        )
+
+      debug_info = string_io.getvalue()
+      expected_included = [
+          debug_prints[f]
+          for f in lm_lib.LMDebugMode
+          if f != lm_lib.LMDebugMode.NONE and f in debug_mode
+      ]
+      expected_excluded = [
+          debug_prints[f]
+          for f in lm_lib.LMDebugMode
+          if f != lm_lib.LMDebugMode.NONE and f not in debug_mode
+      ]
+
+      for expected_include in expected_included:
+        self.assertIn(expected_include, debug_info)
+      for expected_exclude in expected_excluded:
+        self.assertNotIn(expected_exclude, debug_info)
+
+      if debug_mode & lm_lib.LMDebugMode.PROMPT:
+        self.assertIn('[0] MODALITY OBJECTS SENT TO LM', debug_info)
+
+  def test_tokenize_with_unsupported_model(self):
+    with self.assertRaises(NotImplementedError):
+      MockModel().tokenize('hi')
+
   def test_rate_to_max_concurrency_no_rpm_no_tpm(self) -> None:
     lm = MockModel()
     self.assertEqual(

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun/core/llms/fake.py

@@ -25,6 +25,9 @@ class Fake(lf.LanguageModel):
              completions: list[lf.Message]):
     return [lf.LMScoringResult(score=-i * 1.0) for i in range(len(completions))]
 
+  def _tokenize(self, prompt: lf.Message) -> list[tuple[str | bytes, int]]:
+    return [(w, i) for i, w in enumerate(prompt.text.split(' '))]
+
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     results = []
     for prompt in prompts:

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun/core/llms/fake_test.py

@@ -62,6 +62,13 @@ class EchoTest(unittest.TestCase):
         [lf.LMScoringResult(0.0), lf.LMScoringResult(-1.0)],
     )
 
+  def test_tokenize(self):
+    lm = fakelm.Echo()
+    self.assertEqual(
+        lm.tokenize('hi'),
+        [('hi', 0)]
+    )
+
 
 class StaticResponseTest(unittest.TestCase):
 

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun/core/structured/__init__.py

@@ -77,6 +77,8 @@ from langfun.core.structured.completion import complete
 
 from langfun.core.structured.scoring import score
 
+from langfun.core.structured.tokenization import tokenize
+
 # Expose default examples for structured operations so users could refer to
 # them.
 from langfun.core.structured.parsing import default_parse_examples

langfun-0.1.1.dev20240825/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 prompting
+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 = prompting.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-0.1.1.dev20240825/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()

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240822
+Version: 0.1.1.dev20240825
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.1.dev20240822 → langfun-0.1.1.dev20240825}/langfun.egg-info/SOURCES.txt

@@ -116,6 +116,8 @@ langfun/core/structured/schema_generation_test.py
 langfun/core/structured/schema_test.py
 langfun/core/structured/scoring.py
 langfun/core/structured/scoring_test.py
+langfun/core/structured/tokenization.py
+langfun/core/structured/tokenization_test.py
 langfun/core/templates/__init__.py
 langfun/core/templates/completion.py
 langfun/core/templates/completion_test.py