langfun 0.0.2.dev20240330__py3-none-any.whl → 0.0.2.dev20240511__py3-none-any.whl

langfun/core/llms/fake_test.py

@@ -25,7 +25,24 @@ 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,
+                            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 +51,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):
@@ -53,11 +70,45 @@ 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,
+                            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,
+                            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 +120,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 +136,38 @@ 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,
+                            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,
+                            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 +185,38 @@ 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,
+                            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,
+                            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):

langfun/core/llms/google_genai.py

@@ -34,6 +34,7 @@ class GenAI(lf.LanguageModel):
           'gemini-pro-vision',
           'text-bison-001',
           'chat-bison-001',
+          'gemini-1.5-pro-latest',
       ],
       'Model name.',
   ]
@@ -262,6 +263,13 @@ _GOOGLE_GENAI_MODEL_HUB = _ModelHub()
 #
 
 
+class GeminiPro1_5(GenAI):  # pylint: disable=invalid-name
+  """Gemini Pro latest model."""
+
+  model = 'gemini-1.5-pro-latest'
+  multimodal = True
+
+
 class GeminiPro(GenAI):
   """Gemini Pro model."""
 

langfun/core/llms/google_genai_test.py

@@ -152,10 +152,15 @@ class GenAITest(unittest.TestCase):
     )
 
   def test_model_hub(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+
     model = google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
     self.assertIsNotNone(model)
     self.assertIs(google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
 
+    genai.get_model = orig_get_model
+
   def test_api_key_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
       _ = google_genai.GeminiPro()._api_initialized
@@ -167,7 +172,7 @@ class GenAITest(unittest.TestCase):
 
   def test_call(self):
     with mock.patch(
-        'google.generativeai.generative_models.GenerativeModel.generate_content'
+        'google.generativeai.GenerativeModel.generate_content',
     ) as mock_generate:
       orig_get_model = genai.get_model
       genai.get_model = mock_get_model
@@ -176,7 +181,7 @@ class GenAITest(unittest.TestCase):
       lm = google_genai.GeminiPro(api_key='test_key')
       self.maxDiff = None
       self.assertEqual(
-          lm('hello', temperature=2.0, top_k=20).text,
+          lm('hello', temperature=2.0, top_k=20, max_tokens=1024).text,
           (
               'This is a response to hello with n=1, temperature=2.0, '
               'top_p=None, top_k=20, max_tokens=1024, stop=None.'
@@ -197,7 +202,7 @@ class GenAITest(unittest.TestCase):
         (
             "hello to models/text-bison-001 with {'temperature': 2.0, "
             "'top_k': 20, 'top_p': None, 'candidate_count': 1, "
-            "'max_output_tokens': 1024, 'stop_sequences': None}"
+            "'max_output_tokens': None, 'stop_sequences': None}"
         ),
     )
     genai.get_model = orig_get_model

langfun/core/llms/groq.py

@@ -0,0 +1,260 @@
+# 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 Groq."""
+
+import functools
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+import pyglove as pg
+import requests
+
+
+SUPPORTED_MODELS_AND_SETTINGS = {
+    # Refer https://console.groq.com/docs/models
+    'llama3-8b-8192': pg.Dict(max_tokens=8192, max_concurrency=16),
+    'llama3-70b-8192': pg.Dict(max_tokens=8192, max_concurrency=16),
+    'llama2-70b-4096': pg.Dict(max_tokens=4096, max_concurrency=16),
+    'mixtral-8x7b-32768': pg.Dict(max_tokens=32768, max_concurrency=16),
+    'gemma-7b-it': pg.Dict(max_tokens=8192, max_concurrency=16),
+}
+
+
+class GroqError(Exception):  # pylint: disable=g-bad-exception-name
+  """Base class for Groq errors."""
+
+
+class RateLimitError(GroqError):
+  """Error for rate limit reached."""
+
+
+class OverloadedError(GroqError):
+  """Groq's server is temporarily overloaded."""
+
+
+_CHAT_COMPLETE_API_ENDPOINT = 'https://api.groq.com/openai/v1/chat/completions'
+
+
+@lf.use_init_args(['model'])
+class Groq(lf.LanguageModel):
+  """Groq LLMs through REST APIs (OpenAI compatible).
+
+  See https://platform.openai.com/docs/api-reference/chat
+  """
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+      ),
+      'The name of the model to use.',
+  ]
+
+  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
+      False
+  )
+
+  api_key: Annotated[
+      str | None,
+      (
+          'API key. If None, the key will be read from environment variable '
+          "'GROQ_API_KEY'."
+      ),
+  ] = None
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._api_key = None
+    self.__dict__.pop('_api_initialized', None)
+    self.__dict__.pop('_session', None)
+
+  @functools.cached_property
+  def _api_initialized(self):
+    api_key = self.api_key or os.environ.get('GROQ_API_KEY', None)
+    if not api_key:
+      raise ValueError(
+          'Please specify `api_key` during `__init__` or set environment '
+          'variable `GROQ_API_KEY` with your Groq API key.'
+      )
+    self._api_key = api_key
+    return True
+
+  @functools.cached_property
+  def _session(self) -> requests.Session:
+    assert self._api_initialized
+    s = requests.Session()
+    s.headers.update({
+        'Authorization': f'Bearer {self._api_key}',
+        'Content-Type': 'application/json',
+    })
+    return s
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return self.model
+
+  @property
+  def max_concurrency(self) -> int:
+    return SUPPORTED_MODELS_AND_SETTINGS[self.model].max_concurrency
+
+  def _get_request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    # `logprobs` and `top_logprobs` flags are not supported on Groq yet.
+    args = dict(
+        model=self.model,
+        n=options.n,
+        stream=False,
+    )
+
+    if options.temperature is not None:
+      args['temperature'] = options.temperature
+    if options.max_tokens is not None:
+      args['max_tokens'] = options.max_tokens
+    if options.top_p is not None:
+      args['top_p'] = options.top_p
+    if options.stop:
+      args['stop'] = options.stop
+    return args
+
+  def _content_from_message(self, prompt: lf.Message) -> list[dict[str, Any]]:
+    """Converts an message to Groq's content protocol (list of dicts)."""
+    # Refer: https://platform.openai.com/docs/api-reference/chat/create
+    content = []
+    for chunk in prompt.chunk():
+      if isinstance(chunk, str):
+        item = dict(type='text', text=chunk)
+      elif (
+          self.multimodal
+          and isinstance(chunk, lf_modalities.Image)
+          and chunk.uri
+      ):
+        # NOTE(daiyip): Groq only support image URL.
+        item = dict(type='image_url', image_url=chunk.uri)
+      else:
+        raise ValueError(f'Unsupported modality object: {chunk!r}.')
+      content.append(item)
+    return content
+
+  def _message_from_choice(self, choice: dict[str, Any]) -> lf.Message:
+    """Converts Groq's content protocol to message."""
+    # Refer: https://platform.openai.com/docs/api-reference/chat/create
+    content = choice['message']['content']
+    if isinstance(content, str):
+      return lf.AIMessage(content)
+    return lf.AIMessage.from_chunks(
+        [x['text'] for x in content if x['type'] == 'text']
+    )
+
+  def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
+    """Parses Groq's response."""
+    # Refer: https://platform.openai.com/docs/api-reference/chat/object
+    if response.status_code == 200:
+      output = response.json()
+      samples = [
+          lf.LMSample(self._message_from_choice(choice), score=0.0)
+          for choice in output['choices']
+      ]
+      usage = output['usage']
+      return lf.LMSamplingResult(
+          samples,
+          usage=lf.LMSamplingUsage(
+              prompt_tokens=usage['prompt_tokens'],
+              completion_tokens=usage['completion_tokens'],
+              total_tokens=usage['total_tokens'],
+          ),
+      )
+    else:
+      # https://platform.openai.com/docs/guides/error-codes/api-errors
+      if response.status_code == 429:
+        error_cls = RateLimitError
+      elif response.status_code in (500, 502, 503):
+        error_cls = OverloadedError
+      else:
+        error_cls = GroqError
+      raise error_cls(f'{response.status_code}: {response.content}')
+
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    assert self._api_initialized
+    return self._parallel_execute_with_currency_control(
+        self._sample_single,
+        prompts,
+        retry_on_errors=(RateLimitError, OverloadedError),
+    )
+
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    request = dict()
+    request.update(self._get_request_args(self.sampling_options))
+    request.update(
+        dict(
+            messages=[
+                dict(role='user', content=self._content_from_message(prompt))
+            ]
+        )
+    )
+    try:
+      response = self._session.post(
+          _CHAT_COMPLETE_API_ENDPOINT,
+          json=request,
+          timeout=self.timeout,
+      )
+      return self._parse_response(response)
+    except ConnectionError as e:
+      raise OverloadedError(str(e)) from e
+
+
+class GroqLlama3_8B(Groq):  # pylint: disable=invalid-name
+  """Llama3-8B with 8K context window.
+
+  See: https://huggingface.co/meta-llama/Meta-Llama-3-8B
+  """
+
+  model = 'llama3-8b-8192'
+
+
+class GroqLlama3_70B(Groq):  # pylint: disable=invalid-name
+  """Llama3-70B with 8K context window.
+
+  See: https://huggingface.co/meta-llama/Meta-Llama-3-70B
+  """
+
+  model = 'llama3-70b-8192'
+
+
+class GroqLlama2_70B(Groq):  # pylint: disable=invalid-name
+  """Llama2-70B with 4K context window.
+
+  See: https://huggingface.co/meta-llama/Llama-2-70b
+  """
+
+  model = 'llama2-70b-4096'
+
+
+class GroqMistral_8x7B(Groq):  # pylint: disable=invalid-name
+  """Mixtral 8x7B with 32K context window.
+
+  See: https://huggingface.co/meta-llama/Llama-2-70b
+  """
+
+  model = 'mixtral-8x7b-32768'
+
+
+class GroqGemma7B_IT(Groq):  # pylint: disable=invalid-name
+  """Gemma 7B with 8K context window.
+
+  See: https://huggingface.co/google/gemma-1.1-7b-it
+  """
+
+  model = 'gemma-7b-it'

langfun/core/llms/groq_test.py

@@ -0,0 +1,170 @@
+# 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.
+"""Tests for Groq models."""
+
+import os
+from typing import Any
+import unittest
+from unittest import mock
+from langfun.core import modalities as lf_modalities
+from langfun.core.llms import groq
+import pyglove as pg
+import requests
+
+
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'choices': [{
+          'message': {
+              'content': [{
+                  'type': 'text',
+                  'text': (
+                      f'hello with temperature={json.get("temperature")}, '
+                      f'top_p={json.get("top_p")}, '
+                      f'max_tokens={json.get("max_tokens")}, '
+                      f'stop={json.get("stop")}.'
+                  ),
+              }],
+          }
+      }],
+      'usage': {
+          'prompt_tokens': 2,
+          'completion_tokens': 1,
+          'total_tokens': 3,
+      },
+  }).encode()
+  return response
+
+
+def mock_mm_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  v = json['messages'][0]['content'][0]
+  image = lf_modalities.Image.from_uri(v['image_url'])
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'choices': [
+          {
+              'message': {
+                  'content': [{
+                      'type': 'text',
+                      'text': image.uri,
+                  }],
+              }
+          }
+      ],
+      'usage': {
+          'prompt_tokens': 2,
+          'completion_tokens': 1,
+          'total_tokens': 3,
+      },
+  }).encode()
+  return response
+
+
+def mock_requests_post_error(status_code, error_type, error_message):
+  def _mock_requests(url: str, json: dict[str, Any], **kwargs):
+    del url, json, kwargs
+    response = requests.Response()
+    response.status_code = status_code
+    response._content = pg.to_json_str(
+        {
+            'error': {
+                'type': error_type,
+                'message': error_message,
+            }
+        }
+    ).encode()
+    return response
+
+  return _mock_requests
+
+
+class AuthropicTest(unittest.TestCase):
+
+  def test_basics(self):
+    self.assertEqual(groq.GroqMistral_8x7B().model_id, 'mixtral-8x7b-32768')
+    self.assertEqual(groq.GroqMistral_8x7B().max_concurrency, 16)
+
+  def test_api_key(self):
+    lm = groq.GroqMistral_8x7B()
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      lm('hi')
+
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+
+      lm = groq.GroqMistral_8x7B(api_key='fake key')
+      self.assertRegex(lm('hi').text, 'hello.*')
+
+      os.environ['GROQ_API_KEY'] = 'abc'
+      lm = groq.GroqMistral_8x7B()
+      self.assertRegex(lm('hi').text, 'hello.*')
+      del os.environ['GROQ_API_KEY']
+
+  def test_call(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+      lm = groq.GroqLlama3_70B(api_key='fake_key')
+      response = lm(
+          'hello',
+          temperature=0.0,
+          max_tokens=1024,
+          top_k=0.1,
+          top_p=0.2,
+          stop=['\n'],
+      )
+      self.assertEqual(
+          response.text,
+          (
+              'hello with temperature=0.0, top_p=0.2, '
+              "max_tokens=1024, stop=['\\n']."
+          ),
+      )
+      self.assertIsNotNone(response.usage)
+      self.assertIsNotNone(response.usage.prompt_tokens, 2)
+      self.assertIsNotNone(response.usage.completion_tokens, 1)
+      self.assertIsNotNone(response.usage.total_tokens, 3)
+
+  def test_mm_call(self):
+    with mock.patch('requests.Session.post') as mock_mm_request:
+      mock_mm_request.side_effect = mock_mm_requests_post
+      lm = groq.GroqLlama3_70B(multimodal=True, api_key='fake_key')
+      response = lm(lf_modalities.Image.from_uri('https://fake/image.jpg'))
+      self.assertEqual(response.text, 'https://fake/image.jpg')
+
+  def test_call_errors(self):
+    for status_code, error_type, error_message in [
+        (429, 'rate_limit', 'Rate limit exceeded.'),
+        (503, 'service_unavailable', 'Service unavailable.'),
+        (500, 'bad_request', 'Bad request.'),
+    ]:
+      with mock.patch('requests.Session.post') as mock_mm_request:
+        mock_mm_request.side_effect = mock_requests_post_error(
+            status_code, error_type, error_message
+        )
+        lm = groq.GroqLlama3_70B(api_key='fake_key')
+        with self.assertRaisesRegex(
+            Exception, f'{status_code}:.*{error_type}'
+        ):
+          lm('hello', lm=lm, max_attempts=1)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/llama_cpp.py

@@ -51,10 +51,12 @@ class LlamaCppRemote(lf.LanguageModel):
           data = {
               "prompt": prompt.text,
               "n_predict": self.sampling_options.max_tokens,
-              "temperature": self.sampling_options.temperature,
               "top_k": self.sampling_options.top_k or 50,
               "top_p": self.sampling_options.top_p or 0.95,
           }
+          if self.sampling_options.temperature is not None:
+            data["temperature"] = self.sampling_options.temperature
+
           response = requests.post(
               f"{self.url}/completion",
               json=data,