langfun 0.0.2.dev20240507__tar.gz → 0.0.2.dev20240513__tar.gz

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240507
+Version: 0.0.2.dev20240513
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -21,6 +21,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Software Development :: Libraries
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: google-cloud-aiplatform>=1.5.0
 Requires-Dist: google-generativeai>=0.3.2
 Requires-Dist: jinja2>=3.1.2
 Requires-Dist: openai==0.27.2

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/__init__.py

@@ -33,6 +33,11 @@ complete = structured.complete
 score = structured.score
 generate_class = structured.generate_class
 
+# Helper functions for input/output transformations based on
+# `lf.query` (e.g. jax-on-beam could use these for batch processing)
+query_prompt = structured.query_prompt
+query_output = structured.query_output
+
 source_form = structured.source_form
 function_gen = structured.function_gen
 

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/eval/base.py

@@ -1179,7 +1179,7 @@ class Evaluation(Evaluable):
 
   def process(self, example: Any, **kwargs) -> lf.Message:
     """Process an example and returns its output."""
-    prompt = self.prompt.render(example=example).text
+    prompt = lf.Template.from_value(self.prompt, example=example)
     if self.method == 'call':
       return lf_structured.call(
           prompt,
@@ -1207,7 +1207,9 @@ class Evaluation(Evaluable):
     else:
       assert self.method == 'complete', self.method
       assert isinstance(self.schema.spec, pg.typing.Object), self.schema
-      input_value = self.schema.spec.cls.partial(prompt)
+      # TODO(daiyip): Currently multi-modal inputs within the prompt for
+      # completion is not supported.
+      input_value = self.schema.spec.cls.partial(prompt.render().text)
       return lf_structured.complete(
           input_value,
           lm=self.lm,

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/langfunc.py

@@ -14,7 +14,7 @@
 """LangFunc: Language-based functions."""
 
 import dataclasses
-from typing import Annotated, Type, Union
+from typing import Annotated, Type
 
 from langfun.core import component
 from langfun.core import language_model
@@ -328,22 +328,6 @@ class LangFunc(
     """Transforms the output message before returning from __call__."""
     return lm_output
 
-  @classmethod
-  def from_value(
-      cls, value: Union[str, template_lib.Template], **kwargs
-  ) -> 'LangFunc':
-    """Create a LangFunc object from a string or template."""
-    if isinstance(value, LangFunc):
-      return value
-    if isinstance(value, template_lib.Template):
-      lfun = LangFunc(value.template_str, **kwargs)
-      # So lfun could acccess all attributes from value.
-      lfun.sym_setparent(value)
-      return lfun
-    if isinstance(value, str):
-      return LangFunc(template_str=value, **kwargs)
-    return LangFunc('{{input}}', input=value, **kwargs)
-
 
 # Register converter from str to LangFunc, therefore we can always
 # pass strs to attributes that accept LangFunc.

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/langfunc_test.py

@@ -57,6 +57,10 @@ class BasicTest(unittest.TestCase):
     l2 = LangFunc.from_value(l1)
     self.assertIs(l2, l1)
 
+    l3 = LangFunc.from_value(l1, x=1)
+    self.assertIsNot(l3, l1)
+    self.assertTrue(pg.eq(l3, LangFunc('Hello', x=1)))
+
     c = template_lib.Template(
         '{{x}} + {{l}}',
         x=1,

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/llms/__init__.py

@@ -27,6 +27,7 @@ from langfun.core.llms.fake import StaticSequence
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
 from langfun.core.llms.google_genai import GeminiPro
+from langfun.core.llms.google_genai import GeminiPro1_5
 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
@@ -73,6 +74,13 @@ from langfun.core.llms.groq import GroqLlama2_70B
 from langfun.core.llms.groq import GroqMistral_8x7B
 from langfun.core.llms.groq import GroqGemma7B_IT
 
+from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
+from langfun.core.llms.vertexai import VertexAIGeminiPro1
+from langfun.core.llms.vertexai import VertexAIGeminiPro1Vision
+from langfun.core.llms.vertexai import VertexAIPalm2
+from langfun.core.llms.vertexai import VertexAIPalm2_32K
+
 
 # LLaMA C++ models.
 from langfun.core.llms.llama_cpp import LlamaCppRemote

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/llms/fake.py

@@ -57,12 +57,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 _response_from(self, prompt: lf.Message) -> lf.Message:
-    return lf.AIMessage(self.response)
+    return lf.AIMessage.from_value(self.response)
 
 
 @lf.use_init_args(['mapping'])
@@ -70,12 +70,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 _response_from(self, prompt: lf.Message) -> lf.Message:
-    return lf.AIMessage(self.mapping[prompt])
+    return lf.AIMessage.from_value(self.mapping[prompt])
 
 
 @lf.use_init_args(['sequence'])
@@ -83,7 +83,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.'
   ]
 
@@ -92,6 +92,6 @@ class StaticSequence(Fake):
     self._pos = 0
 
   def _response_from(self, prompt: lf.Message) -> lf.Message:
-    r = lf.AIMessage(self.sequence[self._pos])
+    r = lf.AIMessage.from_value(self.sequence[self._pos])
     self._pos += 1
     return r

{langfun-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/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-0.0.2.dev20240507 → langfun-0.0.2.dev20240513}/langfun/core/llms/openai.py

@@ -233,8 +233,9 @@ class OpenAI(lf.LanguageModel):
         for chunk in prompt.chunk():
           if isinstance(chunk, str):
             item = dict(type='text', text=chunk)
-          elif isinstance(chunk, lf_modalities.Image) and chunk.uri:
-            item = dict(type='image_url', image_url=dict(url=chunk.uri))
+          elif isinstance(chunk, lf_modalities.Image):
+            uri = chunk.uri or chunk.content_uri
+            item = dict(type='image_url', image_url=dict(url=uri))
           else:
             raise ValueError(f'Unsupported modality object: {chunk!r}.')
           content.append(item)

langfun-0.0.2.dev20240513/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'