langfun 0.0.2.dev20240329__tar.gz → 0.0.2.dev20240401__tar.gz

{langfun-0.0.2.dev20240329 → langfun-0.0.2.dev20240401}/PKG-INFO

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

{langfun-0.0.2.dev20240329 → langfun-0.0.2.dev20240401}/langfun/core/llms/__init__.py

@@ -25,9 +25,11 @@ from langfun.core.llms.fake import StaticResponse
 from langfun.core.llms.fake import StaticSequence
 
 # Gemini models.
-from langfun.core.llms.gemini import Gemini
-from langfun.core.llms.gemini import GeminiPro
-from langfun.core.llms.gemini import GeminiProVision
+from langfun.core.llms.google_genai import GenAI
+from langfun.core.llms.google_genai import GeminiPro
+from langfun.core.llms.google_genai import GeminiProVision
+from langfun.core.llms.google_genai import Palm2
+from langfun.core.llms.google_genai import Palm2_IT
 
 # OpenAI models.
 from langfun.core.llms.openai import OpenAI

langfun-0.0.2.dev20240329/langfun/core/llms/gemini.py → langfun-0.0.2.dev20240401/langfun/core/llms/google_genai.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Gemini models exposed through Google Generative AI APIs."""
 
+import abc
 import functools
 import os
 from typing import Annotated, Any, Literal
@@ -20,14 +21,20 @@ from typing import Annotated, Any, Literal
 import google.generativeai as genai
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+import pyglove as pg
 
 
 @lf.use_init_args(['model'])
-class Gemini(lf.LanguageModel):
-  """Language model served on VertexAI."""
+class GenAI(lf.LanguageModel):
+  """Language models provided by Google GenAI."""
 
   model: Annotated[
-      Literal['gemini-pro', 'gemini-pro-vision', ''],
+      Literal[
+          'gemini-pro',
+          'gemini-pro-vision',
+          'text-bison-001',
+          'chat-bison-001',
+      ],
       'Model name.',
   ]
 
@@ -35,7 +42,8 @@ class Gemini(lf.LanguageModel):
       str | None,
       (
           'API key. If None, the key will be read from environment variable '
-          "'GOOGLE_API_KEY'."
+          "'GOOGLE_API_KEY'. "
+          'Get an API key at https://ai.google.dev/tutorials/setup'
       ),
   ] = None
 
@@ -43,6 +51,9 @@ class Gemini(lf.LanguageModel):
       False
   )
 
+  # Set the default max concurrency to 8 workers.
+  max_concurrency = 8
+
   def _on_bound(self):
     super()._on_bound()
     self.__dict__.pop('_api_initialized', None)
@@ -67,7 +78,11 @@ class Gemini(lf.LanguageModel):
     return [
         m.name.lstrip('models/')
         for m in genai.list_models()
-        if 'generateContent' in m.supported_generation_methods
+        if (
+            'generateContent' in m.supported_generation_methods
+            or 'generateText' in m.supported_generation_methods
+            or 'generateMessage' in m.supported_generation_methods
+        )
     ]
 
   @property
@@ -80,11 +95,6 @@ class Gemini(lf.LanguageModel):
     """Returns a string to identify the resource for rate control."""
     return self.model_id
 
-  @property
-  def max_concurrency(self) -> int:
-    """Max concurrent requests."""
-    return 8
-
   def _generation_config(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     """Creates generation config from langfun sampling options."""
     return genai.GenerationConfig(
@@ -117,7 +127,7 @@ class Gemini(lf.LanguageModel):
     return chunks
 
   def _response_to_result(
-      self, response: genai.types.GenerateContentResponse
+      self, response: genai.types.GenerateContentResponse | pg.Dict
   ) -> lf.LMSamplingResult:
     """Parses generative response into message."""
     samples = []
@@ -149,17 +159,97 @@ class Gemini(lf.LanguageModel):
     return self._response_to_result(response)
 
 
+class _LegacyGenerativeModel(pg.Object):
+  """Base for legacy GenAI generative model."""
+
+  model: str
+
+  def generate_content(
+      self,
+      input_content: list[str | genai.types.BlobDict],
+      generation_config: genai.GenerationConfig,
+  ) -> pg.Dict:
+    """Generate content."""
+    segments = []
+    for s in input_content:
+      if not isinstance(s, str):
+        raise ValueError(f'Unsupported modality: {s!r}')
+      segments.append(s)
+    return self.generate(' '.join(segments), generation_config)
+
+  @abc.abstractmethod
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig) -> pg.Dict:
+    """Generate response based on prompt."""
+
+
+class _LegacyCompletionModel(_LegacyGenerativeModel):
+  """Legacy GenAI completion model."""
+
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig
+  ) -> pg.Dict:
+    completion: genai.types.Completion = genai.generate_text(
+        model=f'models/{self.model}',
+        prompt=prompt,
+        temperature=generation_config.temperature,
+        top_k=generation_config.top_k,
+        top_p=generation_config.top_p,
+        candidate_count=generation_config.candidate_count,
+        max_output_tokens=generation_config.max_output_tokens,
+        stop_sequences=generation_config.stop_sequences,
+    )
+    return pg.Dict(
+        candidates=[
+            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['output'])]))
+            for c in completion.candidates
+        ]
+    )
+
+
+class _LegacyChatModel(_LegacyGenerativeModel):
+  """Legacy GenAI chat model."""
+
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig
+  ) -> pg.Dict:
+    response: genai.types.ChatResponse = genai.chat(
+        model=f'models/{self.model}',
+        messages=prompt,
+        temperature=generation_config.temperature,
+        top_k=generation_config.top_k,
+        top_p=generation_config.top_p,
+        candidate_count=generation_config.candidate_count,
+    )
+    return pg.Dict(
+        candidates=[
+            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['content'])]))
+            for c in response.candidates
+        ]
+    )
+
+
 class _ModelHub:
   """Google Generative AI model hub."""
 
   def __init__(self):
     self._model_cache = {}
 
-  def get(self, model_name: str) -> genai.GenerativeModel:
+  def get(
+      self, model_name: str
+  ) -> genai.GenerativeModel | _LegacyGenerativeModel:
     """Gets a generative model by model id."""
     model = self._model_cache.get(model_name, None)
     if model is None:
-      model = genai.GenerativeModel(model_name)
+      model_info = genai.get_model(f'models/{model_name}')
+      if 'generateContent' in model_info.supported_generation_methods:
+        model = genai.GenerativeModel(model_name)
+      elif 'generateText' in model_info.supported_generation_methods:
+        model = _LegacyCompletionModel(model_name)
+      elif 'generateMessage' in model_info.supported_generation_methods:
+        model = _LegacyChatModel(model_name)
+      else:
+        raise ValueError(f'Unsupported model: {model_name!r}')
       self._model_cache[model_name] = model
     return model
 
@@ -172,14 +262,26 @@ _GOOGLE_GENAI_MODEL_HUB = _ModelHub()
 #
 
 
-class GeminiPro(Gemini):
+class GeminiPro(GenAI):
   """Gemini Pro model."""
 
   model = 'gemini-pro'
 
 
-class GeminiProVision(Gemini):
+class GeminiProVision(GenAI):
   """Gemini Pro vision model."""
 
   model = 'gemini-pro-vision'
   multimodal = True
+
+
+class Palm2(GenAI):
+  """PaLM2 model."""
+
+  model = 'text-bison-001'
+
+
+class Palm2_IT(GenAI):  # pylint: disable=invalid-name
+  """PaLM2 instruction-tuned model."""
+
+  model = 'chat-bison-001'

langfun-0.0.2.dev20240329/langfun/core/llms/gemini_test.py → langfun-0.0.2.dev20240401/langfun/core/llms/google_genai_test.py

@@ -20,7 +20,7 @@ from unittest import mock
 from google import generativeai as genai
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
-from langfun.core.llms import gemini
+from langfun.core.llms import google_genai
 import pyglove as pg
 
 
@@ -36,6 +36,29 @@ example_image = (
 )
 
 
+def mock_get_model(model_name, *args, **kwargs):
+  del args, kwargs
+  if 'gemini' in model_name:
+    method = 'generateContent'
+  elif 'chat' in model_name:
+    method = 'generateMessage'
+  else:
+    method = 'generateText'
+  return pg.Dict(supported_generation_methods=[method])
+
+
+def mock_generate_text(*, model, prompt, **kwargs):
+  return pg.Dict(
+      candidates=[pg.Dict(output=f'{prompt} to {model} with {kwargs}')]
+  )
+
+
+def mock_chat(*, model, messages, **kwargs):
+  return pg.Dict(
+      candidates=[pg.Dict(content=f'{messages} to {model} with {kwargs}')]
+  )
+
+
 def mock_generate_content(content, generation_config, **kwargs):
   del kwargs
   c = generation_config
@@ -68,12 +91,12 @@ def mock_generate_content(content, generation_config, **kwargs):
   )
 
 
-class GeminiTest(unittest.TestCase):
-  """Tests for Evergreen language model."""
+class GenAITest(unittest.TestCase):
+  """Tests for Google GenAI model."""
 
   def test_content_from_message_text_only(self):
     text = 'This is a beautiful day'
-    model = gemini.GeminiPro()
+    model = google_genai.GeminiPro()
     chunks = model._content_from_message(lf.UserMessage(text))
     self.assertEqual(chunks, [text])
 
@@ -85,9 +108,9 @@ class GeminiTest(unittest.TestCase):
 
     # Non-multimodal model.
     with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
-      gemini.GeminiPro()._content_from_message(message)
+      google_genai.GeminiPro()._content_from_message(message)
 
-    model = gemini.GeminiProVision()
+    model = google_genai.GeminiProVision()
     chunks = model._content_from_message(message)
     self.maxDiff = None
     self.assertEqual(
@@ -118,7 +141,7 @@ class GeminiTest(unittest.TestCase):
             ],
         ),
     )
-    model = gemini.GeminiProVision()
+    model = google_genai.GeminiProVision()
     result = model._response_to_result(response)
     self.assertEqual(
         result,
@@ -129,26 +152,28 @@ class GeminiTest(unittest.TestCase):
     )
 
   def test_model_hub(self):
-    model = gemini._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
+    model = google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
     self.assertIsNotNone(model)
-    self.assertIs(gemini._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
+    self.assertIs(google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
 
   def test_api_key_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      _ = gemini.GeminiPro()._api_initialized
+      _ = google_genai.GeminiPro()._api_initialized
 
-    self.assertTrue(gemini.GeminiPro(api_key='abc')._api_initialized)
+    self.assertTrue(google_genai.GeminiPro(api_key='abc')._api_initialized)
     os.environ['GOOGLE_API_KEY'] = 'abc'
-    self.assertTrue(gemini.GeminiPro()._api_initialized)
+    self.assertTrue(google_genai.GeminiPro()._api_initialized)
     del os.environ['GOOGLE_API_KEY']
 
   def test_call(self):
     with mock.patch(
         'google.generativeai.generative_models.GenerativeModel.generate_content'
     ) as mock_generate:
+      orig_get_model = genai.get_model
+      genai.get_model = mock_get_model
       mock_generate.side_effect = mock_generate_content
 
-      lm = gemini.GeminiPro(api_key='test_key')
+      lm = google_genai.GeminiPro(api_key='test_key')
       self.maxDiff = None
       self.assertEqual(
           lm('hello', temperature=2.0, top_k=20).text,
@@ -157,6 +182,44 @@ class GeminiTest(unittest.TestCase):
               'top_p=None, top_k=20, max_tokens=1024, stop=None.'
           ),
       )
+      genai.get_model = orig_get_model
+
+  def test_call_with_legacy_completion_model(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+    orig_generate_text = genai.generate_text
+    genai.generate_text = mock_generate_text
+
+    lm = google_genai.Palm2(api_key='test_key')
+    self.maxDiff = None
+    self.assertEqual(
+        lm('hello', temperature=2.0, top_k=20).text,
+        (
+            "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}"
+        ),
+    )
+    genai.get_model = orig_get_model
+    genai.generate_text = orig_generate_text
+
+  def test_call_with_legacy_chat_model(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+    orig_chat = genai.chat
+    genai.chat = mock_chat
+
+    lm = google_genai.Palm2_IT(api_key='test_key')
+    self.maxDiff = None
+    self.assertEqual(
+        lm('hello', temperature=2.0, top_k=20).text,
+        (
+            "hello to models/chat-bison-001 with {'temperature': 2.0, "
+            "'top_k': 20, 'top_p': None, 'candidate_count': 1}"
+        ),
+    )
+    genai.get_model = orig_get_model
+    genai.chat = orig_chat
 
 
 if __name__ == '__main__':

{langfun-0.0.2.dev20240329 → langfun-0.0.2.dev20240401}/langfun.egg-info/PKG-INFO

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

{langfun-0.0.2.dev20240329 → langfun-0.0.2.dev20240401}/langfun.egg-info/SOURCES.txt

@@ -57,8 +57,8 @@ langfun/core/eval/scoring_test.py
 langfun/core/llms/__init__.py
 langfun/core/llms/fake.py
 langfun/core/llms/fake_test.py
-langfun/core/llms/gemini.py
-langfun/core/llms/gemini_test.py
+langfun/core/llms/google_genai.py
+langfun/core/llms/google_genai_test.py
 langfun/core/llms/llama_cpp.py
 langfun/core/llms/llama_cpp_test.py
 langfun/core/llms/openai.py