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

langfun/core/llms/groq.py

@@ -0,0 +1,276 @@
+# 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 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 = {
+    # Refer https://console.groq.com/docs/models
+    # Price in US dollars at https://groq.com/pricing/ as of 2024-10-10.
+    'llama-3.2-3b-preview': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=64,
+        cost_per_1k_input_tokens=0.00006,
+        cost_per_1k_output_tokens=0.00006,
+    ),
+    'llama-3.2-1b-preview': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=64,
+        cost_per_1k_input_tokens=0.00004,
+        cost_per_1k_output_tokens=0.00004,
+    ),
+    'llama-3.1-70b-versatile': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=16,
+        cost_per_1k_input_tokens=0.00059,
+        cost_per_1k_output_tokens=0.00079,
+    ),
+    'llama-3.1-8b-instant': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=32,
+        cost_per_1k_input_tokens=0.00005,
+        cost_per_1k_output_tokens=0.00008,
+    ),
+    'llama3-70b-8192': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=16,
+        cost_per_1k_input_tokens=0.00059,
+        cost_per_1k_output_tokens=0.00079,
+    ),
+    'llama3-8b-8192': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=32,
+        cost_per_1k_input_tokens=0.00005,
+        cost_per_1k_output_tokens=0.00008,
+    ),
+    'llama2-70b-4096': pg.Dict(
+        max_tokens=4096,
+        max_concurrency=16,
+    ),
+    'mixtral-8x7b-32768': pg.Dict(
+        max_tokens=32768,
+        max_concurrency=16,
+        cost_per_1k_input_tokens=0.00024,
+        cost_per_1k_output_tokens=0.00024,
+    ),
+    'gemma2-9b-it': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=32,
+        cost_per_1k_input_tokens=0.0002,
+        cost_per_1k_output_tokens=0.0002,
+    ),
+    'gemma-7b-it': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=32,
+        cost_per_1k_input_tokens=0.00007,
+        cost_per_1k_output_tokens=0.00007,
+    ),
+    'whisper-large-v3': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=16,
+    ),
+    'whisper-large-v3-turbo': pg.Dict(
+        max_tokens=8192,
+        max_concurrency=16,
+    )
+}
+
+
+@lf.use_init_args(['model'])
+class Groq(openai_compatible.OpenAICompatible):
+  """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.',
+  ]
+
+  api_key: Annotated[
+      str | None,
+      (
+          'API key. If None, the key will be read from environment variable '
+          "'GROQ_API_KEY'."
+      ),
+  ] = None
+
+  api_endpoint: str = 'https://api.groq.com/openai/v1/chat/completions'
+
+  @property
+  def headers(self) -> dict[str, Any]:
+    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.'
+      )
+    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 self.model
+
+  @property
+  def max_concurrency(self) -> int:
+    return SUPPORTED_MODELS_AND_SETTINGS[self.model].max_concurrency
+
+  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_input_tokens is None or cost_per_1k_output_tokens is None:
+      return None
+    return (
+        cost_per_1k_input_tokens * num_input_tokens
+        + cost_per_1k_output_tokens * num_output_tokens
+    ) / 1000
+
+  def _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 = super()._request_args(options)
+    args.pop('logprobs', None)
+    args.pop('top_logprobs', None)
+    return args
+
+
+class GroqLlama3_2_3B(Groq):  # pylint: disable=invalid-name
+  """Llama3.2-3B with 8K context window.
+
+  See: https://huggingface.co/meta-llama/Llama-3.2-3B
+  """
+
+  model = 'llama-3.2-3b-preview'
+
+
+class GroqLlama3_2_1B(Groq):  # pylint: disable=invalid-name
+  """Llama3.2-1B with 8K context window.
+
+  See: https://huggingface.co/meta-llama/Llama-3.2-1B
+  """
+
+  model = 'llama-3.2-3b-preview'
+
+
+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_1_70B(Groq):  # pylint: disable=invalid-name
+  """Llama3.1-70B with 8K context window.
+
+  See: https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/MODEL_CARD.md   # pylint: disable=line-too-long
+  """
+
+  model = 'llama-3.1-70b-versatile'
+
+
+class GroqLlama3_1_8B(Groq):  # pylint: disable=invalid-name
+  """Llama3.1-8B with 8K context window.
+
+  See: https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/MODEL_CARD.md   # pylint: disable=line-too-long
+  """
+
+  model = 'llama-3.1-8b-instant'
+
+
+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 GroqGemma2_9B_IT(Groq):  # pylint: disable=invalid-name
+  """Gemma2 9B with 8K context window.
+
+  See: https://huggingface.co/google/gemma-2-9b-it
+  """
+
+  model = 'gemma2-9b-it'
+
+
+class GroqGemma_7B_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'
+
+
+class GroqWhisper_Large_v3(Groq):  # pylint: disable=invalid-name
+  """Whisper Large V3 with 8K context window.
+
+  See: https://huggingface.co/openai/whisper-large-v3
+  """
+
+  model = 'whisper-large-v3'
+
+
+class GroqWhisper_Large_v3Turbo(Groq):  # pylint: disable=invalid-name
+  """Whisper Large V3 Turbo with 8K context window.
+
+  See: https://huggingface.co/openai/whisper-large-v3-turbo
+  """
+
+  model = 'whisper-large-v3-turbo'

langfun/core/llms/groq_test.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.
+import os
+import unittest
+import langfun.core as lf
+from langfun.core.llms import groq
+
+
+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)
+    self.assertEqual(groq.GroqMistral_8x7B().estimate_cost(100, 100), 4.8e-5)
+
+  def test_request_args(self):
+    args = groq.GroqMistral_8x7B()._request_args(
+        lf.LMSamplingOptions(
+            temperature=1.0, stop=['\n'], n=1, random_seed=123,
+            logprobs=True, top_logprobs=True
+        )
+    )
+    self.assertNotIn('logprobs', args)
+    self.assertNotIn('top_logprobs', args)
+
+  def test_api_key(self):
+    lm = groq.GroqMistral_8x7B()
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      _ = lm.headers
+
+    lm = groq.GroqMistral_8x7B(api_key='fake key')
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer fake key',
+        }
+    )
+
+    os.environ['GROQ_API_KEY'] = 'abc'
+    lm = groq.GroqMistral_8x7B()
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer abc',
+        }
+    )
+    del os.environ['GROQ_API_KEY']
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/llama_cpp.py

@@ -14,59 +14,34 @@
 """Language models from llama.cpp."""
 
 from typing import Annotated
+from langfun.core.llms import openai_compatible
+import pyglove as pg
 
-import langfun.core as lf
-import requests
 
-
-@lf.use_init_args(["url"])
-class LlamaCppRemote(lf.LanguageModel):
+@pg.use_init_args(['url', 'model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class LlamaCppRemote(openai_compatible.OpenAICompatible):
   """The remote LLaMA C++ model.
 
   The Remote LLaMA C++ models can be launched via
   https://github.com/ggerganov/llama.cpp/tree/master/examples/server
   """
-
   url: Annotated[
       str,
-      "The name of the model to use.",
-  ] = ""
+      'The URL of the LLaMA C++ server.',
+  ]
 
-  name: Annotated[
+  model: Annotated[
       str,
-      "The abbreviation for the LLaMA CPP-based model name.",
-  ] = ""
+      'The name of the model to use.',
+  ] = ''
+
+  @property
+  def api_endpoint(self) -> str:
+    return self.url + '/completion'
 
   @property
   def model_id(self) -> str:
     """Returns a string to identify the model."""
-    return f"LLaMAC++({self.name})"
-
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    def _complete_fn(cur_prompts):
-      results = []
-      for prompt in cur_prompts:
-        result = lf.LMSamplingResult()
-        for _ in range(self.sampling_options.n or 1):
-          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,
-          }
-          response = requests.post(
-              f"{self.url}/completion",
-              json=data,
-              headers={"Content-Type": "application/json"},
-              timeout=self.timeout,
-          )
-          decoded_response = response.json()
-          response = decoded_response["content"]
-          result.samples.append(lf.LMSample(response, score=0.0))
-        results.append(result)
-      return results
+    return f'LLaMAC++({self.model or ""})'
 
-    return self._parallel_execute_with_currency_control(
-        _complete_fn, [prompts]
-    )[0]

langfun/core/llms/llama_cpp_test.py

@@ -11,44 +11,18 @@
 # 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 llama cpp models."""
-
-import typing
 import unittest
-from unittest import mock
-
-import langfun.core as lf
 from langfun.core.llms import llama_cpp
 
 
-def mock_requests_post(url: str, json: typing.Dict[str, typing.Any], **kwargs):
-  del kwargs
-
-  class TEMP:
-
-    def json(self):
-      return {"content": json["prompt"] + "\n" + url}
-
-  return TEMP()
-
-
 class LlamaCppRemoteTest(unittest.TestCase):
   """Tests for the LlamaCppRemote model."""
 
-  def test_call_completion(self):
-    with mock.patch("requests.post") as mock_request:
-      mock_request.side_effect = mock_requests_post
-      lm = llama_cpp.LlamaCppRemote(url="http://127.0.0.1:8080")
-      response = lm("hello", sampling_options=lf.LMSamplingOptions(n=1))
-      self.assertEqual(
-          response.text,
-          "hello\nhttp://127.0.0.1:8080/completion",
-      )
-
-  def test_name(self):
-    lm = llama_cpp.LlamaCppRemote()
+  def test_basics(self):
+    lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
+    self.assertEqual(lm.api_endpoint, "http://127.0.0.1:8080/completion")
     self.assertEqual(lm.model_id, "LLaMAC++()")
-    lm = llama_cpp.LlamaCppRemote(url="xxx", name="x")
+    lm = llama_cpp.LlamaCppRemote("xxx", model="x")
     self.assertEqual(lm.model_id, "LLaMAC++(x)")