From 70db89685790a9f4a82070a315865985c3f59437 Mon Sep 17 00:00:00 2001 From: Firoj Alam Date: Sat, 21 Jun 2025 00:44:30 +0300 Subject: [PATCH 1/3] updated gemini code --- llmebench/models/Gemini.py | 78 +++++++++++++++++++++----------------- 1 file changed, 43 insertions(+), 35 deletions(-) diff --git a/llmebench/models/Gemini.py b/llmebench/models/Gemini.py index 934c675e..e4caac61 100644 --- a/llmebench/models/Gemini.py +++ b/llmebench/models/Gemini.py @@ -6,6 +6,7 @@ import vertexai import vertexai.preview.generative_models as generative_models +from google.oauth2 import service_account from vertexai.generative_models import FinishReason, GenerativeModel, Part from llmebench.models.model_base import ModelBase @@ -53,50 +54,71 @@ class GeminiModel(ModelBase): def __init__( self, project_id=None, - api_key=None, model_name=None, + location=None, + credentials_path=None, # path to JSON file + credentials_info=None, # dict or JSON string timeout=20, temperature=0, + tolerance=1e-7, top_p=0.95, max_tokens=2000, **kwargs, ): - # API parameters - # self.api_url = api_url or os.getenv("AZURE_DEPLOYMENT_API_URL") - self.api_key = api_key or os.getenv("GOOGLE_API_KEY") self.project_id = project_id or os.getenv("GOOGLE_PROJECT_ID") self.model_name = model_name or os.getenv("MODEL") - if self.api_key is None: + self.location = location or os.getenv("VERTEX_LOCATION") or "us-central1" + self.credentials = None + + # 1. Prefer explicit credentials_info (dict or JSON string) + if credentials_info: + if isinstance(credentials_info, str): + credentials_info = json.loads(credentials_info) + self.credentials = service_account.Credentials.from_service_account_info( + credentials_info + ) + # 2. Else, load from path (arg or env) + elif credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS"): + path = credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS") + with open(path, "r") as f: + info = json.load(f) + self.credentials = service_account.Credentials.from_service_account_info( + info + ) + # 3. Else, None: will fall back to ADC (Application Default Credentials) + + if not self.project_id: raise Exception( - "API Key must be provided as model config or environment variable (`GOOGLE_API_KEY`)" + "PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)" ) - if self.project_id is None: + if not self.model_name: + raise Exception("MODEL must be set (argument or `MODEL` in .env)") + if not self.location: raise Exception( - "PROJECT_ID must be provided as model config or environment variable (`GOOGLE_PROJECT_ID`)" + "LOCATION must be set (argument or `VERTEX_LOCATION` in .env)" ) - self.api_timeout = timeout + + vertexai.init( + project=self.project_id, + location=self.location, + credentials=self.credentials, + ) + + self.tolerance = tolerance + self.temperature = max(temperature, tolerance) + self.top_p = top_p + self.max_tokens = max_tokens + self.safety_settings = { generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH, generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH, generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH, generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH, } - # Parameters - tolerance = 1e-7 - self.temperature = temperature - if self.temperature < tolerance: - # Currently, the model inference fails if temperature - # is exactly 0, so we nudge it slightly to work around - # the issue - self.temperature += tolerance - self.top_p = top_p - self.max_tokens = max_tokens super(GeminiModel, self).__init__( retry_exceptions=(TimeoutError, GeminiFailure), **kwargs ) - vertexai.init(project=self.project_id, location="us-central1") - # self.client = GenerativeModel(self.model_name) def summarize_response(self, response): """Returns the "outputs" key's value, if available""" @@ -127,20 +149,6 @@ def prompt(self, processed_input): This method raises this exception if the server responded with a non-ok response """ - # headers = { - # "Content-Type": "application/json", - # "Authorization": "Bearer " + self.api_key, - # } - # body = { - # "input_data": { - # "input_string": processed_input, - # "parameters": { - # "max_tokens": self.max_tokens, - # "temperature": self.temperature, - # "top_p": self.top_p, - # }, - # } - # } generation_config = { "max_output_tokens": self.max_tokens, "temperature": self.temperature, From f9c81dc9775fbb21844bc0b3aedefe121d852d13 Mon Sep 17 00:00:00 2001 From: Firoj Alam Date: Sat, 21 Jun 2025 01:39:03 +0300 Subject: [PATCH 2/3] updated gemini tests --- llmebench/models/Gemini.py | 29 +++++++++++------------------ tests/models/test_Gemini.py | 16 ++++++++-------- 2 files changed, 19 insertions(+), 26 deletions(-) diff --git a/llmebench/models/Gemini.py b/llmebench/models/Gemini.py index e4caac61..44f31ff0 100644 --- a/llmebench/models/Gemini.py +++ b/llmebench/models/Gemini.py @@ -6,9 +6,8 @@ import vertexai import vertexai.preview.generative_models as generative_models -from google.oauth2 import service_account from vertexai.generative_models import FinishReason, GenerativeModel, Part - +from google.oauth2 import service_account from llmebench.models.model_base import ModelBase @@ -56,11 +55,11 @@ def __init__( project_id=None, model_name=None, location=None, - credentials_path=None, # path to JSON file - credentials_info=None, # dict or JSON string + credentials_path=None, # path to JSON file + credentials_info=None, # dict or JSON string timeout=20, temperature=0, - tolerance=1e-7, + tolerance = 1e-7, top_p=0.95, max_tokens=2000, **kwargs, @@ -74,34 +73,28 @@ def __init__( if credentials_info: if isinstance(credentials_info, str): credentials_info = json.loads(credentials_info) - self.credentials = service_account.Credentials.from_service_account_info( - credentials_info - ) + self.credentials = service_account.Credentials.from_service_account_info(credentials_info) # 2. Else, load from path (arg or env) elif credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS"): path = credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS") with open(path, "r") as f: info = json.load(f) - self.credentials = service_account.Credentials.from_service_account_info( - info - ) + self.credentials = service_account.Credentials.from_service_account_info(info) + elif os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is not None: + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") # 3. Else, None: will fall back to ADC (Application Default Credentials) if not self.project_id: - raise Exception( - "PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)" - ) + raise Exception("PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)") if not self.model_name: raise Exception("MODEL must be set (argument or `MODEL` in .env)") if not self.location: - raise Exception( - "LOCATION must be set (argument or `VERTEX_LOCATION` in .env)" - ) + raise Exception("LOCATION must be set (argument or `VERTEX_LOCATION` in .env)") vertexai.init( project=self.project_id, location=self.location, - credentials=self.credentials, + credentials=self.credentials ) self.tolerance = tolerance diff --git a/tests/models/test_Gemini.py b/tests/models/test_Gemini.py index 7d189962..d7fc9165 100644 --- a/tests/models/test_Gemini.py +++ b/tests/models/test_Gemini.py @@ -53,18 +53,18 @@ class TestGeminiDepModelConfig(unittest.TestCase): def test_gemini_deployed_model_config(self): "Test if model config parameters passed as arguments are used" model = GeminiModel( - project_id="test_project_id", api_key="secret-key", model_name="gemini-test" + project_id="test_project_id", model_name="gemini-test", location="us-central1" ) self.assertEqual(model.project_id, "test_project_id") - self.assertEqual(model.api_key, "secret-key") + self.assertEqual(model.location, "us-central1") self.assertEqual(model.model_name, "gemini-test") @patch.dict( "os.environ", { "GOOGLE_PROJECT_ID": "test_project_id", - "GOOGLE_API_KEY": "secret-key", + "LOCATION": "us-central1", "MODEL": "gemini-test", }, ) @@ -73,23 +73,23 @@ def test_gemini_deployed_model_config_env_var(self): model = GeminiModel() self.assertEqual(model.project_id, "test_project_id") - self.assertEqual(model.api_key, "secret-key") + self.assertEqual(model.location, "us-central1") self.assertEqual(model.model_name, "gemini-test") @patch.dict( "os.environ", { "GOOGLE_PROJECT_ID": "test_project_id", - "GOOGLE_API_KEY": "secret-env-key", + "LOCATION": "us-central1", "MODEL": "gemini-test", }, ) def test_gemini_deployed_model_config_priority(self): "Test if model config parameters passed directly get priority" model = GeminiModel( - project_id="test_project_id", api_key="secret-key", model_name="gemini_test" + project_id="test_project_id", model_name="gemini_test", location="us-central1" ) - self.assertEqual(model.project_id, "test_project_id") - self.assertEqual(model.api_key, "secret-key") + self.assertEqual(model.project_id, "test_project_id") + self.assertEqual(model.location, "us-central1") self.assertEqual(model.model_name, "gemini_test") From ed02089580f03c08b9b066caa3d81b8aa10ee74d Mon Sep 17 00:00:00 2001 From: Firoj Alam Date: Sat, 21 Jun 2025 01:40:01 +0300 Subject: [PATCH 3/3] updated gemini tests --- llmebench/models/Gemini.py | 31 +++++++++++++++++++++---------- tests/models/test_Gemini.py | 10 +++++++--- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/llmebench/models/Gemini.py b/llmebench/models/Gemini.py index 44f31ff0..9b8dd9d5 100644 --- a/llmebench/models/Gemini.py +++ b/llmebench/models/Gemini.py @@ -6,8 +6,9 @@ import vertexai import vertexai.preview.generative_models as generative_models -from vertexai.generative_models import FinishReason, GenerativeModel, Part from google.oauth2 import service_account +from vertexai.generative_models import FinishReason, GenerativeModel, Part + from llmebench.models.model_base import ModelBase @@ -55,11 +56,11 @@ def __init__( project_id=None, model_name=None, location=None, - credentials_path=None, # path to JSON file - credentials_info=None, # dict or JSON string + credentials_path=None, # path to JSON file + credentials_info=None, # dict or JSON string timeout=20, temperature=0, - tolerance = 1e-7, + tolerance=1e-7, top_p=0.95, max_tokens=2000, **kwargs, @@ -73,28 +74,38 @@ def __init__( if credentials_info: if isinstance(credentials_info, str): credentials_info = json.loads(credentials_info) - self.credentials = service_account.Credentials.from_service_account_info(credentials_info) + self.credentials = service_account.Credentials.from_service_account_info( + credentials_info + ) # 2. Else, load from path (arg or env) elif credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS"): path = credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS") with open(path, "r") as f: info = json.load(f) - self.credentials = service_account.Credentials.from_service_account_info(info) + self.credentials = service_account.Credentials.from_service_account_info( + info + ) elif os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is not None: - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.getenv( + "GOOGLE_APPLICATION_CREDENTIALS" + ) # 3. Else, None: will fall back to ADC (Application Default Credentials) if not self.project_id: - raise Exception("PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)") + raise Exception( + "PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)" + ) if not self.model_name: raise Exception("MODEL must be set (argument or `MODEL` in .env)") if not self.location: - raise Exception("LOCATION must be set (argument or `VERTEX_LOCATION` in .env)") + raise Exception( + "LOCATION must be set (argument or `VERTEX_LOCATION` in .env)" + ) vertexai.init( project=self.project_id, location=self.location, - credentials=self.credentials + credentials=self.credentials, ) self.tolerance = tolerance diff --git a/tests/models/test_Gemini.py b/tests/models/test_Gemini.py index d7fc9165..55c399e6 100644 --- a/tests/models/test_Gemini.py +++ b/tests/models/test_Gemini.py @@ -53,7 +53,9 @@ class TestGeminiDepModelConfig(unittest.TestCase): def test_gemini_deployed_model_config(self): "Test if model config parameters passed as arguments are used" model = GeminiModel( - project_id="test_project_id", model_name="gemini-test", location="us-central1" + project_id="test_project_id", + model_name="gemini-test", + location="us-central1", ) self.assertEqual(model.project_id, "test_project_id") @@ -87,9 +89,11 @@ def test_gemini_deployed_model_config_env_var(self): def test_gemini_deployed_model_config_priority(self): "Test if model config parameters passed directly get priority" model = GeminiModel( - project_id="test_project_id", model_name="gemini_test", location="us-central1" + project_id="test_project_id", + model_name="gemini_test", + location="us-central1", ) - self.assertEqual(model.project_id, "test_project_id") + self.assertEqual(model.project_id, "test_project_id") self.assertEqual(model.location, "us-central1") self.assertEqual(model.model_name, "gemini_test")