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

langfun/core/llms/vertexai.py

@@ -0,0 +1,291 @@
+# 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.
+"""Vertex AI generative models."""
+
+import functools
+import os
+from typing import Annotated, Any
+
+from google.auth import credentials as credentials_lib
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+import pyglove as pg
+
+
+SUPPORTED_MODELS_AND_SETTINGS = {
+    'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=5),
+    'gemini-1.0-pro': pg.Dict(api='gemini', rpm=300),
+    'gemini-1.0-pro-vision': pg.Dict(api='gemini', rpm=100),
+    # PaLM APIs.
+    'text-bison': pg.Dict(api='palm', rpm=1600),
+    'text-bison-32k': pg.Dict(api='palm', rpm=300),
+    'text-unicorn': pg.Dict(api='palm', rpm=100),
+}
+
+
+@lf.use_init_args(['model'])
+class VertexAI(lf.LanguageModel):
+  """Language model served on VertexAI."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+      ),
+      (
+          'Vertex AI model name. See '
+          'https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models '
+          'for details.'
+      ),
+  ]
+
+  project: Annotated[
+      str | None,
+      (
+          'Vertex AI project ID. Or set from environment variable '
+          'VERTEXAI_PROJECT.'
+      ),
+  ] = None
+
+  location: Annotated[
+      str | None,
+      (
+          'Vertex AI service location. Or set from environment variable '
+          'VERTEXAI_LOCATION.'
+      ),
+  ] = None
+
+  credentials: Annotated[
+      credentials_lib.Credentials | None,
+      (
+          'Credentials to use. If None, the default credentials to the '
+          'environment will be used.'
+      ),
+  ] = None
+
+  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
+      False
+  )
+
+  def _on_bound(self):
+    super()._on_bound()
+    self.__dict__.pop('_api_initialized', None)
+
+  @functools.cached_property
+  def _api_initialized(self):
+    project = self.project or os.environ.get('VERTEXAI_PROJECT', None)
+    if not project:
+      raise ValueError(
+          'Please specify `project` during `__init__` or set environment '
+          'variable `VERTEXAI_PROJECT` with your Vertex AI project ID.'
+      )
+
+    location = self.location or os.environ.get('VERTEXAI_LOCATION', None)
+    if not location:
+      raise ValueError(
+          'Please specify `location` during `__init__` or set environment '
+          'variable `VERTEXAI_LOCATION` with your Vertex AI service location.'
+      )
+
+    credentials = self.credentials
+    # Placeholder for Google-internal credentials.
+    from google.cloud.aiplatform import vertexai  # pylint: disable=g-import-not-at-top
+    vertexai.init(project=project, location=location, credentials=credentials)
+    return True
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'VertexAI({self.model})'
+
+  @property
+  def resource_id(self) -> str:
+    """Returns a string to identify the resource for rate control."""
+    return self.model_id
+
+  @property
+  def max_concurrency(self) -> int:
+    """Returns the maximum number of concurrent requests."""
+    return self.rate_to_max_concurrency(
+        requests_per_min=SUPPORTED_MODELS_AND_SETTINGS[self.model].rpm,
+        tokens_per_min=0,
+    )
+
+  def _generation_config(
+      self, options: lf.LMSamplingOptions
+  ) -> Any:  # generative_models.GenerationConfig
+    """Creates generation config from langfun sampling options."""
+    from google.cloud.aiplatform.vertexai.preview import generative_models  # pylint: disable=g-import-not-at-top
+    return generative_models.GenerationConfig(
+        temperature=options.temperature,
+        top_p=options.top_p,
+        top_k=options.top_k,
+        max_output_tokens=options.max_tokens,
+        stop_sequences=options.stop,
+    )
+
+  def _content_from_message(
+      self, prompt: lf.Message
+  ) -> list[str | Any]:
+    """Gets generation input from langfun message."""
+    from google.cloud.aiplatform.vertexai.preview import generative_models  # pylint: disable=g-import-not-at-top
+    chunks = []
+    for lf_chunk in prompt.chunk():
+      if isinstance(lf_chunk, str):
+        chunk = lf_chunk
+      elif self.multimodal and isinstance(lf_chunk, lf_modalities.Image):
+        chunk = generative_models.Image.from_bytes(lf_chunk.to_bytes())
+      else:
+        raise ValueError(f'Unsupported modality: {lf_chunk!r}')
+      chunks.append(chunk)
+    return chunks
+
+  def _generation_response_to_message(
+      self,
+      response: Any,  # generative_models.GenerationResponse
+  ) -> lf.Message:
+    """Parses generative response into message."""
+    return lf.AIMessage(response.text)
+
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    assert self._api_initialized, 'Vertex AI API is not initialized.'
+    return lf.concurrent_execute(
+        self._sample_single,
+        prompts,
+        executor=self.resource_id,
+        max_workers=self.max_concurrency,
+        # NOTE(daiyip): Vertex has its own policy on handling
+        # with rate limit, so we do not retry on errors.
+        retry_on_errors=None,
+    )
+
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    if self.sampling_options.n > 1:
+      raise ValueError(
+          f'`n` greater than 1 is not supported: {self.sampling_options.n}.'
+      )
+    api = SUPPORTED_MODELS_AND_SETTINGS[self.model].api
+    match api:
+      case 'gemini':
+        return self._sample_generative_model(prompt)
+      case 'palm':
+        return self._sample_text_generation_model(prompt)
+      case _:
+        raise ValueError(f'Unsupported API: {api}')
+
+  def _sample_generative_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    """Samples a generative model."""
+    model = _VERTEXAI_MODEL_HUB.get_generative_model(self.model)
+    input_content = self._content_from_message(prompt)
+    response = model.generate_content(
+        input_content,
+        generation_config=self._generation_config(self.sampling_options),
+    )
+    usage_metadata = response.usage_metadata
+    usage = lf.LMSamplingUsage(
+        prompt_tokens=usage_metadata.prompt_token_count,
+        completion_tokens=usage_metadata.candidates_token_count,
+        total_tokens=usage_metadata.total_token_count,
+    )
+    return lf.LMSamplingResult(
+        [
+            # Scoring is not supported.
+            lf.LMSample(
+                self._generation_response_to_message(response), score=0.0
+            ),
+        ],
+        usage=usage,
+    )
+
+  def _sample_text_generation_model(
+      self, prompt: lf.Message
+  ) -> lf.LMSamplingResult:
+    """Samples a text generation model."""
+    model = _VERTEXAI_MODEL_HUB.get_text_generation_model(self.model)
+    predict_options = dict(
+        temperature=self.sampling_options.temperature,
+        top_k=self.sampling_options.top_k,
+        top_p=self.sampling_options.top_p,
+        max_output_tokens=self.sampling_options.max_tokens,
+        stop_sequences=self.sampling_options.stop,
+    )
+    response = model.predict(prompt.text, **predict_options)
+    return lf.LMSamplingResult([
+        # Scoring is not supported.
+        lf.LMSample(lf.AIMessage(response.text), score=0.0)
+    ])
+
+
+class _ModelHub:
+  """Vertex AI model hub."""
+
+  def __init__(self):
+    self._generative_model_cache = {}
+    self._text_generation_model_cache = {}
+
+  def get_generative_model(
+      self, model_id: str
+  ) -> Any:  # generative_models.GenerativeModel:
+    """Gets a generative model by model id."""
+    model = self._generative_model_cache.get(model_id, None)
+    if model is None:
+      from google.cloud.aiplatform.vertexai.preview import generative_models  # pylint: disable=g-import-not-at-top
+      model = generative_models.GenerativeModel(model_id)
+      self._generative_model_cache[model_id] = model
+    return model
+
+  def get_text_generation_model(
+      self, model_id: str
+  ) -> Any:  # language_models.TextGenerationModel
+    """Gets a text generation model by model id."""
+    model = self._text_generation_model_cache.get(model_id, None)
+    if model is None:
+      from google.cloud.aiplatform.vertexai import language_models  # pylint: disable=g-import-not-at-top
+      model = language_models.TextGenerationModel.from_pretrained(model_id)
+      self._text_generation_model_cache[model_id] = model
+    return model
+
+
+_VERTEXAI_MODEL_HUB = _ModelHub()
+
+
+class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-preview-0409'
+  multimodal = True
+
+
+class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.0 Pro model."""
+
+  model = 'gemini-1.0-pro'
+
+
+class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.0 Pro model."""
+
+  model = 'gemini-1.0-pro-vision'
+  multimodal = True
+
+
+class VertexAIPalm2(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI PaLM2 text generation model."""
+
+  model = 'text-bison'
+
+
+class VertexAIPalm2_32K(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI PaLM2 text generation model (32K context length)."""
+
+  model = 'text-bison-32k'

langfun/core/llms/vertexai_test.py

@@ -0,0 +1,233 @@
+# 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 Gemini models."""
+
+import os
+import unittest
+from unittest import mock
+
+from google.cloud.aiplatform.vertexai.preview import generative_models
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+from langfun.core.llms import vertexai
+import pyglove as pg
+
+
+example_image = (
+    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x18\x00\x00\x00\x18\x04'
+    b'\x03\x00\x00\x00\x12Y \xcb\x00\x00\x00\x18PLTE\x00\x00'
+    b'\x00fff_chaag_cg_ch^ci_ciC\xedb\x94\x00\x00\x00\x08tRNS'
+    b'\x00\n\x9f*\xd4\xff_\xf4\xe4\x8b\xf3a\x00\x00\x00>IDATx'
+    b'\x01c \x05\x08)"\xd8\xcc\xae!\x06pNz\x88k\x19\\Q\xa8"\x10'
+    b'\xc1\x14\x95\x01%\xc1\n\xa143Ta\xa8"D-\x84\x03QM\x98\xc3'
+    b'\x1a\x1a\x1a@5\x0e\x04\xa0q\x88\x05\x00\x07\xf8\x18\xf9'
+    b'\xdao\xd0|\x00\x00\x00\x00IEND\xaeB`\x82'
+)
+
+
+def mock_generate_content(content, generation_config, **kwargs):
+  del kwargs
+  c = pg.Dict(generation_config.to_dict())
+  print('zzz', c)
+  return generative_models.GenerationResponse.from_dict({
+      'candidates': [
+          {
+              'index': 0,
+              'content': {
+                  'role': 'model',
+                  'parts': [
+                      {
+                          'text': (
+                              f'This is a response to {content[0]} with '
+                              f'temperature={c.temperature}, '
+                              f'top_p={c.top_p}, '
+                              f'top_k={c.top_k}, '
+                              f'max_tokens={c.max_output_tokens}, '
+                              f'stop={"".join(c.stop_sequences)}.'
+                          )
+                      },
+                  ],
+              },
+          },
+      ]
+  })
+
+
+class VertexAITest(unittest.TestCase):
+  """Tests for Vertex model."""
+
+  def test_content_from_message_text_only(self):
+    text = 'This is a beautiful day'
+    model = vertexai.VertexAIGeminiPro1()
+    chunks = model._content_from_message(lf.UserMessage(text))
+    self.assertEqual(chunks, [text])
+
+  def test_content_from_message_mm(self):
+    message = lf.UserMessage(
+        'This is an {{image}}, what is it?',
+        image=lf_modalities.Image.from_bytes(example_image),
+    )
+
+    # Non-multimodal model.
+    with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
+      vertexai.VertexAIGeminiPro1()._content_from_message(message)
+
+    model = vertexai.VertexAIGeminiPro1Vision()
+    chunks = model._content_from_message(message)
+    self.maxDiff = None
+    self.assertEqual([chunks[0], chunks[2]], ['This is an', ', what is it?'])
+    self.assertIsInstance(chunks[1], generative_models.Image)
+
+  def test_generation_response_to_message_text_only(self):
+    response = generative_models.GenerationResponse.from_dict({
+        'candidates': [
+            {
+                'index': 0,
+                'content': {
+                    'role': 'model',
+                    'parts': [
+                        {
+                            'text': 'hello world',
+                        },
+                    ],
+                },
+            },
+        ],
+    })
+    model = vertexai.VertexAIGeminiPro1()
+    message = model._generation_response_to_message(response)
+    self.assertEqual(message, lf.AIMessage('hello world'))
+
+  def test_model_hub(self):
+    with mock.patch(
+        'google.cloud.aiplatform.vertexai.preview.generative_models.'
+        'GenerativeModel.__init__'
+    ) as mock_model_init:
+      mock_model_init.side_effect = lambda *args, **kwargs: None
+      model = vertexai._VERTEXAI_MODEL_HUB.get_generative_model(
+          'gemini-1.0-pro'
+      )
+      self.assertIsNotNone(model)
+      self.assertIs(
+          vertexai._VERTEXAI_MODEL_HUB.get_generative_model('gemini-1.0-pro'),
+          model,
+      )
+
+    with mock.patch(
+        'google.cloud.aiplatform.vertexai.language_models.'
+        'TextGenerationModel.from_pretrained'
+    ) as mock_model_init:
+
+      class TextGenerationModel:
+        pass
+
+      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
+      model = vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model(
+          'text-bison'
+      )
+      self.assertIsNotNone(model)
+      self.assertIs(
+          vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model('text-bison'),
+          model,
+      )
+
+  def test_project_and_location_check(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
+      _ = vertexai.VertexAIGeminiPro1()._api_initialized
+
+    with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
+      _ = vertexai.VertexAIGeminiPro1(project='abc')._api_initialized
+
+    self.assertTrue(
+        vertexai.VertexAIGeminiPro1(
+            project='abc', location='us-central1'
+        )._api_initialized
+    )
+
+    os.environ['VERTEXAI_PROJECT'] = 'abc'
+    os.environ['VERTEXAI_LOCATION'] = 'us-central1'
+    self.assertTrue(vertexai.VertexAIGeminiPro1()._api_initialized)
+    del os.environ['VERTEXAI_PROJECT']
+    del os.environ['VERTEXAI_LOCATION']
+
+  def test_call_generative_model(self):
+    with mock.patch(
+        'google.cloud.aiplatform.vertexai.preview.generative_models.'
+        'GenerativeModel.__init__'
+    ) as mock_model_init:
+      mock_model_init.side_effect = lambda *args, **kwargs: None
+
+      with mock.patch(
+          'google.cloud.aiplatform.vertexai.preview.generative_models.'
+          'GenerativeModel.generate_content'
+      ) as mock_generate:
+        mock_generate.side_effect = mock_generate_content
+
+        lm = vertexai.VertexAIGeminiPro1(project='abc', location='us-central1')
+        self.assertEqual(
+            lm(
+                'hello',
+                temperature=2.0,
+                top_p=1.0,
+                top_k=20,
+                max_tokens=1024,
+                stop='\n',
+            ).text,
+            (
+                'This is a response to hello with temperature=2.0, '
+                'top_p=1.0, top_k=20.0, max_tokens=1024, stop=\n.'
+            ),
+        )
+
+  def test_call_text_generation_model(self):
+    with mock.patch(
+        'google.cloud.aiplatform.vertexai.language_models.'
+        'TextGenerationModel.from_pretrained'
+    ) as mock_model_init:
+
+      class TextGenerationModel:
+
+        def predict(self, prompt, **kwargs):
+          c = pg.Dict(kwargs)
+          return pg.Dict(
+              text=(
+                  f'This is a response to {prompt} with '
+                  f'temperature={c.temperature}, '
+                  f'top_p={c.top_p}, '
+                  f'top_k={c.top_k}, '
+                  f'max_tokens={c.max_output_tokens}, '
+                  f'stop={"".join(c.stop_sequences)}.'
+              )
+          )
+
+      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
+      lm = vertexai.VertexAIPalm2(project='abc', location='us-central1')
+      self.assertEqual(
+          lm(
+              'hello',
+              temperature=2.0,
+              top_p=1.0,
+              top_k=20,
+              max_tokens=1024,
+              stop='\n',
+          ).text,
+          (
+              'This is a response to hello with temperature=2.0, '
+              'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
+          ),
+      )
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/modalities/image.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """Image modality."""
 
-import base64
 import imghdr
 from typing import cast
 from langfun.core.modalities import mime
@@ -36,5 +35,4 @@ class Image(mime.MimeType):
   def _repr_html_(self) -> str:
     if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
       return f'<img src="{self.uri}">'
-    image_raw = base64.b64encode(self.to_bytes()).decode()
-    return f'<img src="data:image/{self.image_format};base64,{image_raw}">'
+    return f'<img src="{self.content_uri}">'

langfun/core/modalities/mime.py

@@ -14,6 +14,7 @@
 """MIME type data."""
 
 import abc
+import base64
 from typing import Annotated, Union
 import langfun.core as lf
 import pyglove as pg
@@ -54,6 +55,11 @@ class MimeType(lf.Modality):
     self.rebind(content=content, skip_notification=True)
     return self.content
 
+  @property
+  def content_uri(self) -> str:
+    base64_content = base64.b64encode(self.to_bytes()).decode()
+    return f'data:{self.mime_type};base64,{base64_content}'
+
   @classmethod
   def from_uri(cls, uri: str, **kwargs) -> 'MimeType':
     return cls(uri=uri, content=None, **kwargs)

langfun/core/modalities/video.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """Video modality."""
 
-import base64
 from typing import cast
 from langfun.core.modalities import mime
 
@@ -40,8 +39,7 @@ class Video(mime.MimeType):
   def _repr_html_(self) -> str:
     if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
       return f'<video controls> <source src="{self.uri}"> </video>'
-    video_raw = base64.b64encode(self.to_bytes()).decode()
     return (
         '<video controls> <source'
-        f' src="data:video/{self.video_format};base64,{video_raw}"> </video>'
+        f' src="data:video/{self.content_uri}"> </video>'
     )

langfun/core/structured/__init__.py

@@ -64,6 +64,8 @@ from langfun.core.structured.prompting import QueryStructure
 from langfun.core.structured.prompting import QueryStructureJson
 from langfun.core.structured.prompting import QueryStructurePython
 from langfun.core.structured.prompting import query
+from langfun.core.structured.prompting import query_prompt
+from langfun.core.structured.prompting import query_output
 
 from langfun.core.structured.description import DescribeStructure
 from langfun.core.structured.description import describe

langfun/core/structured/mapping.py

@@ -251,7 +251,7 @@ class Mapping(lf.LangFunc):
 
       {%- if example.schema -%}
       {{ schema_title }}:
-      {{ example.schema_repr(protocol) | indent(2, True) }}
+      {{ example.schema_repr(protocol, include_methods=include_methods) | indent(2, True) }}
 
       {% endif -%}
 
@@ -279,6 +279,10 @@ class Mapping(lf.LangFunc):
       'The protocol for representing the schema and value.',
   ] = 'python'
 
+  include_methods: Annotated[
+      bool, 'If True, include method definitions in the schema.'
+  ] = False
+
   #
   # Other user-provided flags.
   #