add gemini models

Author: EYH0602

main.py

@@ -24,6 +24,7 @@
     GPT_MODELS,
     CLAUDE_MODELS,
     DEEPSEEK_MODELS,
+    GEMINI_MODELS,
     get_ant_model,
     get_ant_ttc_model,
     get_oai_model,
@@ -83,7 +84,12 @@ def main(
     """
     assert (
         model
-        in GPT_MODELS + OLLAMA_MODELS + CLAUDE_MODELS + O1_MODELS + DEEPSEEK_MODELS
+        in GPT_MODELS
+        + OLLAMA_MODELS
+        + CLAUDE_MODELS
+        + O1_MODELS
+        + DEEPSEEK_MODELS
+        + GEMINI_MODELS
     ), f"{model} is not supported."
 
     if output_file is None:
@@ -99,11 +105,11 @@ def main(
     if model in GPT_MODELS:
         assert "OPENAI_API_KEY" in os.environ, "Please set OPEN_API_KEY in environment!"
         client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
-        generate = get_oai_model(client, model, seed, temperature, pure)
+        generate = get_oai_model(client, model, pure)
     elif model in O1_MODELS:
         assert "OPENAI_API_KEY" in os.environ, "Please set OPEN_API_KEY in environment!"
         client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
-        generate = get_o1_model(client, model, seed, temperature, pure)
+        generate = get_o1_model(client, model, pure)
     elif model in CLAUDE_MODELS:
         assert (
             "ANTHROPIC_API_KEY" in os.environ
@@ -120,7 +126,16 @@ def main(
         client = OpenAI(
             api_key=os.environ["DEEPSEEK_API_KEY"], base_url="https://api.deepseek.com"
         )
-        generate = get_oai_model(client, model, seed, temperature, pure)
+        generate = get_oai_model(client, model, pure)
+    elif model in GEMINI_MODELS:
+        assert (
+            "GEMINI_API_KEY" in os.environ
+        ), "Please set GEMINI_API_KEY in environment!"
+        client = OpenAI(
+            api_key=os.environ["GEMINI_API_KEY"],
+            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+        )
+        generate = get_oai_model(client, model)
     else:
         client = OllamaClient(host=f"http://localhost:{port}")
         generate = get_ollama_model(client, model, seed, temperature, pure)

src/experiment.py

@@ -6,6 +6,9 @@
 from anthropic import Anthropic, InternalServerError
 from typing import Callable
 
+from google import genai
+from google.genai import types
+
 from src.common import (
     SEED,
     TEMPERATURE,
@@ -39,12 +42,19 @@
     "deepseek-chat",
 ]
 
+GEMINI_MODELS = [
+    "gemini-2.5-pro-preview-03-25",
+    "gemini-2.0-flash",
+    "gemini-2.0-flash-lite",
+    "gemini-1.5-flash",
+    "gemini-1.5-flash-8b",
+    "gemini-1.5-pro",
+]
+
 
 def get_o1_model(
     client: OpenAI,
     model: str = "o1-preview-2024-09-12",
-    seed: int = SEED,
-    temperature: float = TEMPERATURE,
     pure: bool = False,
 ) -> Callable[[str], str | None]:
     def generate_type_signature(prompt: str) -> str | None:
@@ -64,8 +74,6 @@ def generate_type_signature(prompt: str) -> str | None:
 def get_oai_model(
     client: OpenAI,
     model: str = "gpt-3.5-turbo",
-    seed: int = SEED,
-    temperature: float = TEMPERATURE,
     pure: bool = False,
 ) -> Callable[[str], str | None]:
     def generate_type_signature(prompt: str) -> str | None:
@@ -78,9 +86,6 @@ def generate_type_signature(prompt: str) -> str | None:
                 {"role": "user", "content": prompt},
             ],
             model=model,
-            # Set parameters to ensure reproducibility
-            seed=seed,
-            temperature=temperature,
         )
 
         content = completion.choices[0].message.content
@@ -123,8 +128,6 @@ def generate_type_signature(prompt: str) -> str | None:
 def get_ant_model(
     client: Anthropic,
     model: str = "claude-3-5-sonnet-20240620",
-    seed: int = SEED,
-    temperature: float = TEMPERATURE,
     pure: bool = False,
 ) -> Callable[[str], str | None]:
     def generate_type_signature(prompt: str) -> str | None:
@@ -136,10 +139,6 @@ def generate_type_signature(prompt: str) -> str | None:
                 ],
                 model=model,
                 max_tokens=1024,
-                # ! the following parameters are not supported by Claude API
-                # seed=seed,
-                # temperature=temperature,
-                # top_p=top_p,
             )
         except InternalServerError as e:
             print(e)