Added gemini-pro to the api

RICommunity · AnayMehrotra · Dec 19, 2023 · Dec 19, 2023 · Dec 19, 2023 · acd96d84e518fe67801159b5d502937ea53eb122
commit acd96d84e518fe67801159b5d502937ea53eb122
diff --git a/conversers.py b/conversers.py
@@ -1,6 +1,6 @@
 
 import common
-from language_models import GPT, PaLM, HuggingFace, APIModelLlama7B, APIModelVicuna13B
+from language_models import GPT, PaLM, HuggingFace, APIModelLlama7B, APIModelVicuna13B, GeminiPro
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from config import VICUNA_PATH, LLAMA_PATH, ATTACK_TEMP, TARGET_TEMP, ATTACK_TOP_P, TARGET_TOP_P, MAX_PARALLEL_STREAMS 
@@ -216,6 +216,8 @@ def load_indiv_model(model_name):
         lm = GPT(model_name)
     elif model_name == "palm-2":
         lm = PaLM(model_name)
+    elif model_name == "gemini-pro":
+        lm = GeminiPro(model_name)
     elif model_name == 'llama-2-api-model':
         lm = APIModelLlama7B(model_name)
     elif model_name == 'vicuna-api-model':
@@ -282,11 +284,15 @@ def get_model_path_and_template(model_name):
         "palm-2":{
             "path":"palm-2",
             "template":"palm-2"
+        },
+        "gemini-pro": {
+            "path": "gemini-pro",
+            "template": "gemini-pro"
         }
     }
     path, template = full_model_dict[model_name]["path"], full_model_dict[model_name]["template"]
     return path, template
 
 
 
-
+
diff --git a/language_models.py b/language_models.py
@@ -5,7 +5,7 @@
 import torch
 import gc
 from typing import Dict, List
-import google.generativeai as palm
+import google.generativeai as genai
 import urllib3
 from copy import deepcopy
 
@@ -257,7 +257,7 @@ class PaLM():
 
     def __init__(self, model_name) -> None:
         self.model_name = model_name
-        palm.configure(api_key=self.API_KEY) 
+        genai.configure(api_key=self.API_KEY) 
 
     def generate(self, conv: List, 
                 max_n_tokens: int, 
@@ -275,7 +275,7 @@ def generate(self, conv: List,
         output = self.API_ERROR_OUTPUT
         for _ in range(self.API_MAX_RETRY):            
             try:
-                completion = palm.chat(
+                completion = genai.chat(
                     messages=conv,
                     temperature=temperature,
                     top_p=top_p
@@ -303,3 +303,66 @@ def batched_generate(self,
                         temperature: float,
                         top_p: float = 1.0,):
         return [self.generate(conv, max_n_tokens, temperature, top_p) for conv in convs_list]
+
+
+class GeminiPro():
+    API_RETRY_SLEEP = 10
+    API_ERROR_OUTPUT = "$ERROR$"
+    API_QUERY_SLEEP = 1
+    API_MAX_RETRY = 5
+    API_TIMEOUT = 20
+    default_output = "I'm sorry, but I cannot assist with that request."
+    API_KEY = os.getenv("PALM_API_KEY")
+
+    def __init__(self, model_name) -> None:
+        self.model_name = model_name
+        genai.configure(api_key=self.API_KEY) 
+
+    def generate(self, conv: List, 
+                max_n_tokens: int, 
+                temperature: float,
+                top_p: float):
+        '''
+        Args:
+            conv: List of dictionaries, 
+            max_n_tokens: int, max number of tokens to generate
+            temperature: float, temperature for sampling
+            top_p: float, top p for sampling
+        Returns:
+            str: generated response
+        '''
+        output = self.API_ERROR_OUTPUT
+        for _ in range(self.API_MAX_RETRY):            
+            try:
+                model = genai.GenerativeModel(self.model_name)
+                output = model.generate_content(
+                    contents = conv,
+                    generation_config = genai.GenerationConfig(
+                        candidate_count = 1,
+                        temperature = temperature,
+                        top_p = top_p,
+                        max_output_tokens=max_n_tokens,
+                    )
+                )
+
+                if output is None:
+                    # If PaLM refuses to output and returns None, we replace it with a default output
+                    output = self.default_output
+                else:
+                    # Use this approximation since PaLM does not allow
+                    # to specify max_tokens. Each token is approximately 4 characters.
+                    output = output.text
+                break
+            except Exception as e:
+                print(type(e), e)
+                time.sleep(self.API_RETRY_SLEEP)
+
+            time.sleep(self.API_QUERY_SLEEP)
+        return output
+
+    def batched_generate(self, 
+                        convs_list: List[List[Dict]],
+                        max_n_tokens: int, 
+                        temperature: float,
+                        top_p: float = 1.0,):
+        return [self.generate(conv, max_n_tokens, temperature, top_p) for conv in convs_list]
diff --git a/main_TAP.py b/main_TAP.py
@@ -293,7 +293,9 @@ def main(args):
                  "gpt-4",
                  'gpt-4-turbo', 
                  'gpt-4-1106-preview', # This is same as gpt-4-turbo
-                 "palm-2"]
+                 "palm-2",
+                 "gemini-pro",
+                 ]
     )
     parser.add_argument(
         "--target-max-n-tokens",

diff --git a/requirements.txt b/requirements.txt
@@ -8,3 +8,4 @@ pytz==2023.3.post1
 torch==2.1.0
 transformers==4.35.0
 urllib3==2.1.0
+google-generativeai==0.3.1