langfun 0.1.1.dev20240808__py3-none-any.whl → 0.1.1.dev20240811__py3-none-any.whl

langfun/core/llms/__init__.py

@@ -105,6 +105,7 @@ 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
+from langfun.core.llms.vertexai import VertexAICustom
 
 
 # LLaMA C++ models.

langfun/core/llms/vertexai.py

@@ -18,6 +18,7 @@ import os
 from typing import Annotated, Any
 
 from google.auth import credentials as credentials_lib
+from google.cloud.aiplatform import aiplatform
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 import pyglove as pg
@@ -35,6 +36,9 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     '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),
+    # Endpoint
+    # TODO(chengrun): Set a more appropriate rpm for endpoint.
+    'custom': pg.Dict(api='endpoint', rpm=20),
 }
 
 
@@ -53,6 +57,11 @@ class VertexAI(lf.LanguageModel):
       ),
   ]
 
+  endpoint_name: pg.typing.Annotated[
+      str | None,
+      'Vertex Endpoint name or ID.',
+  ]
+
   project: Annotated[
       str | None,
       (
@@ -177,6 +186,13 @@ class VertexAI(lf.LanguageModel):
     """Parses generative response into message."""
     return lf.AIMessage(response.text)
 
+  def _generation_endpoint_response_to_message(
+      self,
+      response: Any,  # google.cloud.aiplatform.aiplatform.models.Prediction
+  ) -> lf.Message:
+    """Parses Endpoint response into message."""
+    return lf.AIMessage(response.predictions[0])
+
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized, 'Vertex AI API is not initialized.'
     # TODO(yifenglu): It seems this exception is due to the instability of the
@@ -212,6 +228,8 @@ class VertexAI(lf.LanguageModel):
         return self._sample_generative_model(prompt)
       case 'palm':
         return self._sample_text_generation_model(prompt)
+      case 'endpoint':
+        return self._sample_endpoint_model(prompt)
       case _:
         raise ValueError(f'Unsupported API: {api}')
 
@@ -257,6 +275,34 @@ class VertexAI(lf.LanguageModel):
         lf.LMSample(lf.AIMessage(response.text), score=0.0)
     ])
 
+  def _sample_endpoint_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    """Samples a text generation model."""
+    model = aiplatform.Endpoint(self.endpoint_name)
+    # TODO(chengrun): Add support for stop_sequences.
+    predict_options = dict(
+        temperature=self.sampling_options.temperature
+        if self.sampling_options.temperature is not None
+        else 1.0,
+        top_k=self.sampling_options.top_k
+        if self.sampling_options.top_k is not None
+        else 32,
+        top_p=self.sampling_options.top_p
+        if self.sampling_options.top_p is not None
+        else 1,
+        max_tokens=self.sampling_options.max_tokens
+        if self.sampling_options.max_tokens is not None
+        else 8192,
+    )
+    instances = [{'prompt': prompt.text, **predict_options}]
+    response = model.predict(instances=instances)
+
+    return lf.LMSamplingResult([
+        # Scoring is not supported.
+        lf.LMSample(
+            self._generation_endpoint_response_to_message(response), score=0.0
+        )
+    ])
+
 
 class _ModelHub:
   """Vertex AI model hub."""
@@ -387,3 +433,9 @@ class VertexAIPalm2_32K(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI PaLM2 text generation model (32K context length)."""
 
   model = 'text-bison-32k'
+
+
+class VertexAICustom(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Custom model (Endpoint)."""
+
+  model = 'custom'

langfun/core/llms/vertexai_test.py

@@ -17,6 +17,7 @@ import os
 import unittest
 from unittest import mock
 
+from google.cloud.aiplatform.aiplatform import models as aiplatform_models
 from vertexai import generative_models
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
@@ -63,6 +64,20 @@ def mock_generate_content(content, generation_config, **kwargs):
   })
 
 
+def mock_endpoint_predict(instances, **kwargs):
+  del kwargs
+  assert len(instances) == 1
+  return aiplatform_models.Prediction(
+      predictions=[
+          f"This is a response to {instances[0]['prompt']} with"
+          f" temperature={instances[0]['temperature']},"
+          f" top_p={instances[0]['top_p']}, top_k={instances[0]['top_k']},"
+          f" max_tokens={instances[0]['max_tokens']}."
+      ],
+      deployed_model_id='',
+  )
+
+
 class VertexAITest(unittest.TestCase):
   """Tests for Vertex model."""
 
@@ -227,6 +242,30 @@ class VertexAITest(unittest.TestCase):
           ),
       )
 
+  def test_call_endpoint_model(self):
+    with mock.patch(
+        'google.cloud.aiplatform.aiplatform.Endpoint.predict'
+    ) as mock_model_predict:
+
+      mock_model_predict.side_effect = mock_endpoint_predict
+      lm = vertexai.VertexAI(
+          'custom',
+          endpoint_name='123',
+          project='abc',
+          location='us-central1',
+      )
+      self.assertEqual(
+          lm(
+              'hello',
+              temperature=2.0,
+              top_p=1.0,
+              top_k=20,
+              max_tokens=50,
+          ),
+          'This is a response to hello with temperature=2.0, top_p=1.0,'
+          ' top_k=20, max_tokens=50.',
+      )
+
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.1.dev20240808.dist-info → langfun-0.1.1.dev20240811.dist-info}/METADATA

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

{langfun-0.1.1.dev20240808.dist-info → langfun-0.1.1.dev20240811.dist-info}/RECORD

@@ -52,7 +52,7 @@ langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0
 langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrCG1L6w,4775
 langfun/core/eval/scoring.py,sha256=AlCwEVrU6nvURDB1aPxA2XBUmOjWxuNJDXJoS4-6VbU,6386
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
-langfun/core/llms/__init__.py,sha256=YGILcGi2QTxDG0v-0Gd4uAj1HL_zRhtllOM9EURxzDg,4712
+langfun/core/llms/__init__.py,sha256=ggPvkKq8l-B8rN3ZD6d7r3_d5DYAd5fC7FgC6ZI2Wzo,4766
 langfun/core/llms/anthropic.py,sha256=Gon3fOi31RhZFgNd0ijyTnKnUdp9hrWrCoSXyO4UaLw,7316
 langfun/core/llms/anthropic_test.py,sha256=T-swuMkfnlgs8Fpif4rtXs579exGk0TsbLMirXDZCkg,5533
 langfun/core/llms/fake.py,sha256=Dd7-6ka9pFf3fcWZyczamjOqQ91MOI-m7We3Oc9Ffmo,2927
@@ -67,8 +67,8 @@ langfun/core/llms/openai.py,sha256=jILxfFb3vBuyf1u_2-LVfs_wekPF2RVuNFzNVg25pEA,1
 langfun/core/llms/openai_test.py,sha256=3muDTnW7UBOSHq694Fi2bofqhe8Pkj0Tl8IShoLCTOM,15525
 langfun/core/llms/rest.py,sha256=laopuq-zD8V-3Y6eFDngftHEbE66VlUkCD2-rvvRaLU,3388
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
-langfun/core/llms/vertexai.py,sha256=pSvEWZ2jqS21TEEoztn3VWaqT09abkMcj07Oe2BJR1s,12017
-langfun/core/llms/vertexai_test.py,sha256=G18BG36h5KvmX2zutDTLjtYCRjTuP_nWIFm4FMnLnyY,7651
+langfun/core/llms/vertexai.py,sha256=-YoEUlG19CWIhJb8S6puPqdX9SoiT5NNAItefwdCfsk,13781
+langfun/core/llms/vertexai_test.py,sha256=N3k4N9_bVjC6_Qtg4WO9jYNv8M9xmv5UdODvIKG2upo,8835
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
 langfun/core/llms/cache/in_memory.py,sha256=YfFyJEhLs73cUiB0ZfhMxYpdE8Iuxxw-dvMFwGHTSHw,4742
@@ -117,8 +117,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=rBW2Qr8yi-aWYwoTwRR-n1peKyMX9QXPZXURjLgoiRs,2264
-langfun-0.1.1.dev20240808.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.1.dev20240808.dist-info/METADATA,sha256=fhgQonI-ve4MRW400VgoZI4e3cA5GMpQZmjtdo1tOpc,5234
-langfun-0.1.1.dev20240808.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-langfun-0.1.1.dev20240808.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.1.dev20240808.dist-info/RECORD,,
+langfun-0.1.1.dev20240811.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.1.dev20240811.dist-info/METADATA,sha256=zo0-slFX-N_vCKGKOHamespDwn3SbGpM0TsIkdwB0Fk,5234
+langfun-0.1.1.dev20240811.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+langfun-0.1.1.dev20240811.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.1.dev20240811.dist-info/RECORD,,