Spaces:
Runtime error
Runtime error
File size: 1,554 Bytes
3060e5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import os
import google.generativeai as genai
class GeminiClient:
def __init__(self, system_message=None):
self._system_message = system_message
self._connect_client()
def _connect_client(self):
if not os.getenv("GOOGLE_PALM_KEY"):
raise Exception("Please set your Google MakerSuite API key")
api_key = os.getenv("GOOGLE_PALM_KEY")
genai.configure(api_key=api_key)
safety_settings = [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_ONLY_HIGH",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_ONLY_HIGH",
},
]
defaults = {
"temperature": 0.7,
"top_k": 40,
"top_p": 0.95,
"max_output_tokens": 1024,
}
self._model = genai.GenerativeModel(
model_name="gemini-pro",
generation_config=defaults,
safety_settings=safety_settings,
)
def generate_text(self, prompt: str) -> str:
full_prompt = self._system_message + prompt
try:
response = self._model.generate_content(full_prompt).text
except Exception as e:
print(f"Error: {e}")
response = ""
return response
|