Spaces:
Sleeping
Sleeping
File size: 1,846 Bytes
e9559d9 a2ba326 e9559d9 a2ba326 e9559d9 |
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 53 54 55 |
import os
import google.generativeai as genai
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()
generation_config = {
"temperature": 0,
"top_k": 1,
"max_output_tokens": 4000,
}
class LLM:
def __init__(self, model_name) -> None:
self.model_name = model_name
self.model = self.create_model(model_name)
def create_model(self, model_name):
match model_name:
case "gemini-pro-vision":
genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
return genai.GenerativeModel(model_name)
case "gemini-pro":
genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
return genai.GenerativeModel(
model_name,generation_config=generation_config)
case "OpenAI":
return OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
case _:
print("Not Implemented")
def __call__(self, prompt, image=None):
if self.model_name == 'gemini-pro-vision':
response = self.model.generate_content(
[image, prompt]
)
elif self.model_name == "gemini-pro":
response = self.model.generate_content(
prompt)
return response.text
elif self.model_name == 'OpenAI':
res = self.model.chat.completions.create(
model="gpt-3.5-turbo-1106",
# response_format={"type": "json_object"},
messages=[
# {"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": f"{prompt}"},
],
# seed=10,
temperature=0
)
return res.choices[0].message.content |