langfun 0.0.2.dev20240627__py3-none-any.whl → 0.0.2.dev20240629__py3-none-any.whl

langfun/core/language_model.py

@@ -84,20 +84,23 @@ class LMSamplingUsage(pg.Object):
   prompt_tokens: int
   completion_tokens: int
   total_tokens: int
+  num_requests: int = 1
 
   def __add__(self, other: 'LMSamplingUsage') -> 'LMSamplingUsage':
     return LMSamplingUsage(
         prompt_tokens=self.prompt_tokens + other.prompt_tokens,
         completion_tokens=self.completion_tokens + other.completion_tokens,
         total_tokens=self.total_tokens + other.total_tokens,
+        num_requests=self.num_requests + other.num_requests,
     )
 
 
 class UsageNotAvailable(LMSamplingUsage):
   """Usage information not available."""
-  prompt_tokens: pg.typing.Int(0).freeze()        # pytype: disable=invalid-annotation
+  prompt_tokens: pg.typing.Int(0).freeze()       # pytype: disable=invalid-annotation
   completion_tokens: pg.typing.Int(0).freeze()   # pytype: disable=invalid-annotation
   total_tokens: pg.typing.Int(0).freeze()        # pytype: disable=invalid-annotation
+  num_requests: pg.typing.Int(1).freeze()        # pytype: disable=invalid-annotation
 
   def __bool__(self) -> bool:
     return False
@@ -726,7 +729,7 @@ class _UsageTracker:
   def __init__(self, model_ids: set[str] | None):
     self.model_ids = model_ids
     self.usages = {
-        m: LMSamplingUsage(0, 0, 0) for m in model_ids
+        m: LMSamplingUsage(0, 0, 0, 0) for m in model_ids
     } if model_ids else {}
 
   def track(self, model_id: str, usage: LMSamplingUsage):

langfun/core/language_model_test.py

@@ -598,18 +598,18 @@ class LanguageModelTest(unittest.TestCase):
             list(concurrent.concurrent_map(call_lm, ['hi', 'hello']))
 
     self.assertEqual(usages2, {
-        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200, 1),
     })
     self.assertEqual(usages3, {
-        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4),
+        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4, 4),
     })
     self.assertEqual(usages4, {
-        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4),
-        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4, 4),
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200, 1),
     })
     self.assertEqual(usages1, {
-        'model1': lm_lib.LMSamplingUsage(100 * 5, 100 * 5, 200 * 5),
-        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+        'model1': lm_lib.LMSamplingUsage(100 * 5, 100 * 5, 200 * 5, 5),
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200, 1),
     })
 
 

langfun/core/llms/__init__.py

@@ -94,8 +94,10 @@ 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_5_0514
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_0409
 from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_0514
 from langfun.core.llms.vertexai import VertexAIGeminiPro1
 from langfun.core.llms.vertexai import VertexAIGeminiPro1Vision
 from langfun.core.llms.vertexai import VertexAIPalm2

langfun/core/llms/vertexai.py

@@ -24,6 +24,8 @@ import pyglove as pg
 
 
 SUPPORTED_MODELS_AND_SETTINGS = {
+    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=5),
+    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=5),
     'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=5),
     'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=5),
     'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=5),
@@ -330,12 +332,19 @@ _PDF = [
 class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
+  model = 'gemini-1.5-pro-001'
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+
+
+class VertexAIGeminiPro1_5_0514(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro preview model."""
+
   model = 'gemini-1.5-pro-preview-0514'
   supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
 
 
 class VertexAIGeminiPro1_5_0409(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Pro model."""
+  """Vertex AI Gemini 1.5 Pro preview model."""
 
   model = 'gemini-1.5-pro-preview-0409'
   supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
@@ -343,6 +352,12 @@ class VertexAIGeminiPro1_5_0409(VertexAI):  # pylint: disable=invalid-name
 
 class VertexAIGeminiFlash1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
+  model = 'gemini-1.5-flash-001'
+  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+
+
+class VertexAIGeminiFlash1_5_0514(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Flash preview model."""
 
   model = 'gemini-1.5-flash-preview-0514'
   supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES

{langfun-0.0.2.dev20240627.dist-info → langfun-0.0.2.dev20240629.dist-info}/METADATA

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

{langfun-0.0.2.dev20240627.dist-info → langfun-0.0.2.dev20240629.dist-info}/RECORD

@@ -8,8 +8,8 @@ langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
 langfun/core/langfunc.py,sha256=RvIcRjIq0jWYRu1xim-FYe4HSrt97r3GMBO_PuagUmw,11060
 langfun/core/langfunc_test.py,sha256=lyt-UzkD8972cxZwzCkps0_RMLeSsOBrcUFIW-fB6us,8653
-langfun/core/language_model.py,sha256=q41EGhbgoCe68eYWWTAzgt3r6SbdCpINRdXnNGstukQ,23775
-langfun/core/language_model_test.py,sha256=HmA0GACK4-6tCH32TFkfYj9w4CxbrynKqXdKBgiqgwo,21255
+langfun/core/language_model.py,sha256=tHB2otsHx-EM4TaDNgXZ-qvnhJXNJKuK2-OdseQzGdI,23948
+langfun/core/language_model_test.py,sha256=TlNmVUfBfDQZzIiiBqCBTrxgcoyj2qNp3kONvmr2pX4,21273
 langfun/core/logging.py,sha256=FyZRxUy2TTF6tWLhQCRpCvfH55WGUdNgQjUTK_SQLnY,5320
 langfun/core/logging_test.py,sha256=qvm3RObYP3knO2PnXR9evBRl4gH621GnjnwywbGbRfg,1833
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
@@ -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=1J7IATo-8FXUR0SBqk9icztHiM0lWkBFcWUo-vUURgQ,6376
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
-langfun/core/llms/__init__.py,sha256=GDarEQZlDFIxpEW_6Hz969NLmsT7AnRChhnx8uCtDy4,4419
+langfun/core/llms/__init__.py,sha256=3OA4uQczFEKHoIUcWDO3myesdeDISLUWJ8fjp6zjGww,4551
 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,7 +67,7 @@ langfun/core/llms/openai.py,sha256=0z9qIH9FlWj9VWUnhOX321T6JHO-vjY2IozT7OVI4GY,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=wIpckH-rMHUBA1vhauQk4LVrSsPQEsVntz7kLDKwm9g,11359
+langfun/core/llms/vertexai.py,sha256=OH7z10FW5Xer08NZ7MGwSbXB3TcxwhNj3B5sZGnDq3I,11950
 langfun/core/llms/vertexai_test.py,sha256=G18BG36h5KvmX2zutDTLjtYCRjTuP_nWIFm4FMnLnyY,7651
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
@@ -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.0.2.dev20240627.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240627.dist-info/METADATA,sha256=KqNAHkVDoVHKj1Shze7rm6lYJ_MdDWJFkQVAmbg-QI0,3550
-langfun-0.0.2.dev20240627.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
-langfun-0.0.2.dev20240627.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240627.dist-info/RECORD,,
+langfun-0.0.2.dev20240629.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240629.dist-info/METADATA,sha256=a8CBmosl7t2AvpjeT4rU_dt8hi05PNgiLc9302UzzhY,3550
+langfun-0.0.2.dev20240629.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
+langfun-0.0.2.dev20240629.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240629.dist-info/RECORD,,