Spaces:
Paused
Paused
File size: 2,460 Bytes
c89a08c |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
from openai import OpenAI
from groq import Groq
class LLM:
def __init__(self, api_key, modelo, temperatura=0.5, quantidade_tokens=2048, mensagem=""):
self.temperatura = temperatura
self.quantidade_tokens = quantidade_tokens
self.mensagem = mensagem
self.modelo = modelo
self.api_key = api_key
self.openai_client = None
self.groq_client = None
def _init_openai_client(self):
if not self.openai_client:
self.openai_client = OpenAI(api_key=self.api_key)
def _init_groq_client(self):
if not self.groq_client:
self.groq_client = Groq(api_key=self.api_key)
def gerar_resposta_openai(self):
self._init_openai_client()
messages = []
if self.mensagem.strip():
messages.append({"role": "system", "content": self.mensagem})
messages.append({"role": "user", "content": self.prompt})
response = self.openai_client.chat.completions.create(
model=self.modelo,
messages=messages,
temperature=self.temperatura,
max_tokens=self.quantidade_tokens,
top_p=1,
frequency_penalty=0,
presence_penalty=0
)
return response.choices[0].message.content.strip()
def gerar_resposta_groq(self):
self._init_groq_client()
messages = []
if self.mensagem.strip():
messages.append({"role": "system", "content": self.mensagem})
messages.append({"role": "user", "content": self.prompt})
completion = self.groq_client.chat.completions.create(
model=self.modelo,
messages=messages,
temperature=self.temperatura,
max_tokens=self.quantidade_tokens,
top_p=0.65,
stream=False, # Mudamos para False para manter consistência com OpenAI
stop=None
)
return completion.choices[0].message.content.strip()
def gerar_resposta(self):
"""Método principal para gerar respostas baseado no modelo selecionado"""
if "gpt" in self.modelo: # Modelos OpenAI
return self.gerar_resposta_openai()
elif "llama" in self.modelo or "mixtral" in self.modelo: # Modelos Groq
return self.gerar_resposta_groq()
else:
raise ValueError(f"Modelo não suportado: {self.modelo}")
|