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

langfun/core/llms/compositional_test.py

@@ -0,0 +1,73 @@
+# 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.
+"""Tests for compositional models."""
+import unittest
+
+import langfun.core as lf
+from langfun.core.llms import compositional
+from langfun.core.llms import fake
+
+
+class RandomChoiceTest(unittest.TestCase):
+
+  def test_basic(self):
+    lm = compositional.RandomChoice([
+        fake.StaticResponse('hi'),
+        fake.StaticSequence(['hello', 'world'])
+    ])
+    self.assertEqual(
+        lm.model_id,
+        'RandomChoice(StaticResponse, StaticSequence)'
+    )
+    self.assertEqual(
+        lm.resource_id,
+        'RandomChoice(StaticResponse, StaticSequence)'
+    )
+    self.assertEqual(
+        [lm('a'), lm('b'), lm('c')],
+        ['hello', 'world', 'hi']
+    )
+    lm = lm.clone()
+    self.assertEqual(
+        [
+            x.samples[0].response for x in [
+                lm.sample(['a'])[0],
+                lm.sample(['b'])[0],
+                lm.sample(['c'])[0],
+            ]
+        ],
+        ['hello', 'world', 'hi']
+    )
+    self.assertEqual(
+        lm.score('hello', ['world']),
+        [lf.LMScoringResult(0.0)]
+    )
+    self.assertEqual(
+        lm.tokenize('hello'),
+        [('hello', 0)]
+    )
+
+  def test_sampling_options(self):
+    lm = compositional.RandomChoice([
+        fake.StaticResponse('hi'),
+        fake.StaticSequence(['hello', 'world'])
+    ], temperature=0.5)
+    self.assertEqual(
+        lm.candidates[0].sampling_options.temperature,
+        0.5
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/deepseek.py

@@ -0,0 +1,117 @@
+# 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.
+"""Language models from DeepSeek."""
+
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core.llms import openai_compatible
+import pyglove as pg
+
+SUPPORTED_MODELS_AND_SETTINGS = {
+    # pylint: disable=g-line-too-long
+    # TODO(yifenglu): The RPM and TPM are arbitrary numbers. Update them once DeepSeek provides concrete guidelines.
+    # DeepSeek doesn't control the rate limit at the moment: https://api-docs.deepseek.com/quick_start/rate_limit
+    # The cost is based on: https://api-docs.deepseek.com/quick_start/pricing
+    'deepseek-chat': pg.Dict(
+        in_service=True,
+        rpm=100,
+        tpm=1000000,
+        cost_per_1k_input_tokens=0.00014,
+        cost_per_1k_output_tokens=0.00028,
+    ),
+}
+
+
+# DeepSeek API uses an API format compatible with OpenAI.
+# Reference: https://api-docs.deepseek.com/
+@lf.use_init_args(['model'])
+class DeepSeek(openai_compatible.OpenAICompatible):
+  """DeepSeek model."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+      ),
+      'The name of the model to use.',
+  ]
+
+  api_endpoint: str = 'https://api.deepseek.com/chat/completions'
+
+  api_key: Annotated[
+      str | None,
+      (
+          'API key. If None, the key will be read from environment variable '
+          "'DEEPSEEK_API_KEY'."
+      ),
+  ] = None
+
+  @property
+  def headers(self) -> dict[str, Any]:
+    api_key = self.api_key or os.environ.get('DEEPSEEK_API_KEY', None)
+    if not api_key:
+      raise ValueError(
+          'Please specify `api_key` during `__init__` or set environment '
+          'variable `DEEPSEEK_API_KEY` with your DeepSeek API key.'
+      )
+    headers = super().headers
+    headers.update({
+        'Authorization': f'Bearer {api_key}',
+    })
+    return headers
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'DeepSeek({self.model})'
+
+  @property
+  def max_concurrency(self) -> int:
+    rpm = SUPPORTED_MODELS_AND_SETTINGS[self.model].get('rpm', 0)
+    tpm = SUPPORTED_MODELS_AND_SETTINGS[self.model].get('tpm', 0)
+    return self.rate_to_max_concurrency(
+        requests_per_min=rpm, tokens_per_min=tpm
+    )
+
+  def estimate_cost(
+      self, num_input_tokens: int, num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    cost_per_1k_input_tokens = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_input_tokens', None
+    )
+    cost_per_1k_output_tokens = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_output_tokens', None
+    )
+    if cost_per_1k_output_tokens is None or cost_per_1k_input_tokens is None:
+      return None
+    return (
+        cost_per_1k_input_tokens * num_input_tokens
+        + cost_per_1k_output_tokens * num_output_tokens
+    ) / 1000
+
+  @classmethod
+  def dir(cls):
+    return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
+
+
+class DeepSeekChat(DeepSeek):
+  """DeepSeek Chat model.
+
+  Currently, it is powered by DeepSeek-V3 model, 64K input contenxt window and
+  8k max output tokens.
+  """
+
+  model = 'deepseek-chat'

langfun/core/llms/deepseek_test.py

@@ -0,0 +1,61 @@
+# 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
+from langfun.core.llms import deepseek
+
+
+class DeepSeekTest(unittest.TestCase):
+  """Tests for DeepSeek language model."""
+
+  def test_dir(self):
+    self.assertIn('deepseek-chat', deepseek.DeepSeek.dir())
+
+  def test_key(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      _ = deepseek.DeepSeekChat().headers
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer test_key',
+        }
+    )
+
+  def test_model_id(self):
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').model_id,
+        'DeepSeek(deepseek-chat)',
+    )
+
+  def test_resource_id(self):
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').resource_id,
+        'DeepSeek(deepseek-chat)',
+    )
+
+  def test_max_concurrency(self):
+    self.assertGreater(
+        deepseek.DeepSeekChat(api_key='test_key').max_concurrency, 0
+    )
+
+  def test_estimate_cost(self):
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').estimate_cost(
+            num_input_tokens=100, num_output_tokens=100
+        ),
+        4.2e-5
+    )
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/fake.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Fake LMs for testing."""
 
+import abc
 from typing import Annotated
 import langfun.core as lf
 
@@ -20,18 +21,39 @@ import langfun.core as lf
 class Fake(lf.LanguageModel):
   """The base class for all fake language models."""
 
-  def _score(self, prompt: lf.Message, completions: list[lf.Message]):
+  def _score(self, prompt: lf.Message| list[lf.Message],
+             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:
+      response = self._response_from(prompt)
+      results.append(
+          lf.LMSamplingResult(
+              [lf.LMSample(response, 1.0)],
+              usage=lf.LMSamplingUsage(
+                  prompt_tokens=len(prompt.text),
+                  completion_tokens=len(response.text),
+                  total_tokens=len(prompt.text) + len(response.text),
+              )
+          )
+      )
+    return results
+
+  @abc.abstractmethod
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    """Returns the response for the given prompt."""
+
 
 class Echo(Fake):
   """A simple echo language model for testing."""
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(prompt.text, 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(prompt.text)
 
 
 @lf.use_init_args(['response'])
@@ -39,15 +61,12 @@ class StaticResponse(Fake):
   """Language model that always gives the same canned response."""
 
   response: Annotated[
-      str,
+      str | lf.Message,
       'A canned response that will be returned regardless of the prompt.'
   ]
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.response, 1.0)])
-        for _ in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage.from_value(self.response)
 
 
 @lf.use_init_args(['mapping'])
@@ -55,15 +74,12 @@ class StaticMapping(Fake):
   """A static mapping from prompt to response."""
 
   mapping: Annotated[
-      dict[str, str],
+      dict[str, str | lf.Message],
       'A mapping from prompt to response.'
   ]
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.mapping[prompt], 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage.from_value(self.mapping[prompt])
 
 
 @lf.use_init_args(['sequence'])
@@ -71,7 +87,7 @@ class StaticSequence(Fake):
   """A static sequence of responses to use."""
 
   sequence: Annotated[
-      list[str],
+      list[str | lf.Message],
       'A sequence of strings as the response.'
   ]
 
@@ -79,10 +95,7 @@ class StaticSequence(Fake):
     super()._on_bound()
     self._pos = 0
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    results = []
-    for _ in prompts:
-      results.append(lf.LMSamplingResult(
-          [lf.LMSample(self.sequence[self._pos], 1.0)]))
-      self._pos += 1
-    return results
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    r = lf.AIMessage.from_value(self.sequence[self._pos])
+    self._pos += 1
+    return r

langfun/core/llms/fake_test.py

@@ -25,7 +25,25 @@ class EchoTest(unittest.TestCase):
   def test_sample(self):
     lm = fakelm.Echo()
     self.assertEqual(
-        lm.sample(['hi']), [lf.LMSamplingResult([lf.LMSample('hi', 1.0)])]
+        lm.sample(['hi']),
+        [
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'hi',
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(2, 2, 4),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                lf.LMSamplingUsage(2, 2, 4))
+        ]
     )
 
   def test_call(self):
@@ -34,8 +52,8 @@ class EchoTest(unittest.TestCase):
     with contextlib.redirect_stdout(string_io):
       self.assertEqual(lm('hi'), 'hi')
     debug_info = string_io.getvalue()
-    self.assertIn('[0] LM INFO:', debug_info)
-    self.assertIn('[0] PROMPT SENT TO LM:', debug_info)
+    self.assertIn('[0] LM INFO', debug_info)
+    self.assertIn('[0] PROMPT SENT TO LM', debug_info)
     self.assertIn('[0] LM RESPONSE', debug_info)
 
   def test_score(self):
@@ -45,6 +63,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):
 
@@ -53,11 +78,47 @@ class StaticResponseTest(unittest.TestCase):
     lm = fakelm.StaticResponse(canned_response)
     self.assertEqual(
         lm.sample(['hi']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(2, 38, 40),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(2, 38, 40)
+            )
+        ],
     )
     self.assertEqual(
         lm.sample(['Tell me a joke.']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(15, 38, 53),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(15, 38, 53)
+            )
+        ],
     )
 
   def test_call(self):
@@ -69,8 +130,8 @@ class StaticResponseTest(unittest.TestCase):
       self.assertEqual(lm('hi'), canned_response)
 
     debug_info = string_io.getvalue()
-    self.assertIn('[0] LM INFO:', debug_info)
-    self.assertIn('[0] PROMPT SENT TO LM:', debug_info)
+    self.assertIn('[0] LM INFO', debug_info)
+    self.assertIn('[0] PROMPT SENT TO LM', debug_info)
     self.assertIn('[0] LM RESPONSE', debug_info)
 
 
@@ -85,8 +146,40 @@ class StaticMappingTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(KeyError):
@@ -104,8 +197,40 @@ class StaticSequenceTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(IndexError):