Revert "update"

This reverts commit 0dae8ed.

Author: nauyisu022

docs/guides/generation_details.md

@@ -70,6 +70,8 @@ config.claude_api = "claude api"
 
 config.openai_key = "openai api"
 
+config.palm_api = "palm api"
+
 config.ernie_client_id = "ernie client id"
 
 config.ernie_client_secret = "ernie client secret"

trustllm_pkg/setup.py

@@ -25,7 +25,9 @@
         'python-dotenv',
         'urllib3',
         'anthropic',
+        'google.generativeai',
         'google-api-python-client',
+        'google.ai.generativelanguage',
         'replicate',
         'zhipuai>=2.0.1'
 ],

trustllm_pkg/trustllm/config.py

@@ -9,6 +9,7 @@
 deepinfra_api = None
 ernie_api = None
 claude_api = None
+palm_api = None
 replicate_api = None
 zhipu_api = None
 
@@ -37,17 +38,19 @@
 zhipu_model = ["glm-4", "glm-3-turbo"]
 claude_model = ["claude-2", "claude-instant-1"]
 openai_model = ["chatgpt", "gpt-4"]
+google_model = ["bison-001", "gemini"]
 wenxin_model = ["ernie"]
 replicate_model=["vicuna-7b","vicuna-13b","vicuna-33b","chatglm3-6b","llama3-70b","llama3-8b"]
 
-online_model = deepinfra_model + zhipu_model + claude_model + openai_model + wenxin_model+replicate_model
+online_model = deepinfra_model + zhipu_model + claude_model + openai_model + google_model + wenxin_model+replicate_model
 
 model_info = {
     "online_model": online_model,
     "zhipu_model": zhipu_model,
     "deepinfra_model": deepinfra_model,
     'claude_model': claude_model,
     'openai_model': openai_model,
+    'google_model': google_model,
     'wenxin_model': wenxin_model,
     'replicate_model':replicate_model,
     "model_mapping": {

trustllm_pkg/trustllm/utils/generation_utils.py

@@ -1,6 +1,7 @@
 import os, json
 from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT
-
+import google.generativeai as genai
+from google.generativeai.types import safety_types
 from fastchat.model import load_model, get_conversation_template
 from openai import OpenAI,AzureOpenAI
 from tenacity import retry, wait_random_exponential, stop_after_attempt
@@ -16,6 +17,16 @@
 model_mapping = model_info['model_mapping']
 rev_model_mapping = {value: key for key, value in model_mapping.items()}
 
+# Define safety settings to allow harmful content generation
+safety_setting = [
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_DEROGATORY, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_VIOLENCE, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_SEXUAL, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_TOXICITY, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_MEDICAL, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+    {"category": safety_types.HarmCategory.HARM_CATEGORY_DANGEROUS, "threshold": safety_types.HarmBlockThreshold.BLOCK_NONE},
+]
+
 # Retrieve model information
 def get_models():
     return model_mapping, online_model_list
@@ -87,7 +98,31 @@ def claude_api(string, model, temperature):
     return completion.completion
 
 
+@retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(6))
+def gemini_api(string, temperature):
+    genai.configure(api_key=trustllm.config.gemini_api)
+    model = genai.GenerativeModel('gemini-pro')
+    response = model.generate_content(string, temperature=temperature, safety_settings=safety_setting)
+    return response
+
+
 
+@retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(6))
+def palm_api(string, model, temperature):
+    genai.configure(api_key=trustllm.config.palm_api)
+
+    model_mapping = {
+        'bison-001': 'models/text-bison-001',
+    }
+    completion = genai.generate_text(
+        model=model_mapping[model],  # models/text-bison-001
+        prompt=string,
+        temperature=temperature,
+        # The maximum length of the response
+        max_output_tokens=4000,
+        safety_settings=safety_setting
+    )
+    return completion.result
 
 
 @retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(6))
@@ -113,6 +148,11 @@ def zhipu_api(string, model, temperature):
 def gen_online(model_name, prompt, temperature, replicate=False, deepinfra=False):
     if model_name in model_info['wenxin_model']:
         res = get_ernie_res(prompt, temperature=temperature)
+    elif model_name in model_info['google_model']:
+        if model_name == 'bison-001':
+            res = palm_api(prompt, model=model_name, temperature=temperature)
+        elif model_name == 'gemini-pro':
+            res = gemini_api(prompt, temperature=temperature)
     elif model_name in model_info['openai_model']:
         res = get_res_openai(prompt, model=model_name, temperature=temperature)
     elif model_name in model_info['deepinfra_model']: