joaomorossini's picture
First commit for new HF Space. Adapted from https://huggingface.co/spaces/tseronni/startup_genome
f11dfb5
raw
history blame
3.85 kB
import openai
class MessageChatCompletion:
def __init__(self,
model: str = 'gpt-3.5-turbo-0125',
message: str = '',
api_key: str = '',
temperature: float = 0.10,
top_p: float = 0.95,
n: int = 1,
stream: bool = False,
stop: str = "\n",
max_tokens: int = 4096,
presence_penalty: float = 0.0,
frequency_penalty: float = 0.0,
logit_bias: int = None,
user: str = ''):
self.api_key = api_key
openai.api_key = self.api_key
if model in ["gpt-4", "gpt-4-turbo-preview", "gpt-3.5-turbo", "gpt-3.5-turbo-0125"]:
self.endpoint = "https://api.openai.com/v1/chat/completions"
else:
self.endpoint = "https://api.openai.com/v1/completions"
self.headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
}
self.prompt = {
"model": model,
"messages": [],
"temperature": temperature,
"top_p": top_p,
"n": n,
"stream": stream,
"stop": stop,
"presence_penalty": presence_penalty,
"frequency_penalty": frequency_penalty
}
if max_tokens is not None:
self.prompt["max_tokens"] = max_tokens
if logit_bias is not None:
self.prompt["logit_bias"] = logit_bias
if user != '':
self.prompt["user"] = user
if message != '':
self.new_user_message(content=message)
self.response = ''
self.error = False
def new_message(self, role: str = 'user', content: str = '', name: str = ''):
new_message = {"role": role, "content": f"{content}"}
if name != '':
new_message['name'] = name
self.prompt['messages'].append(new_message)
def new_user_message(self, content: str = '', name: str = ''):
self.new_message(role='user', content=content, name=name)
def new_system_message(self, content: str = '', name: str = ''):
self.new_message(role='system', content=content, name=name)
def new_assistant_message(self, content: str = '', name: str = ''):
self.new_message(role='assistant', content=content, name=name)
def get_last_message(self):
return self.prompt['messages'][-1]['content']
def send_message(self):
try:
self.error = False
response = openai.chat.completions.create(
model=self.prompt['model'],
messages=self.prompt['messages'],
frequency_penalty=self.prompt['frequency_penalty'],
temperature=self.prompt['temperature'],
max_tokens=self.prompt['max_tokens'],
top_p=self.prompt['top_p'],
presence_penalty=self.prompt['presence_penalty'],
stream=self.prompt['stream']
)
full_response = response.choices[0].message.content
# if stream = True
# full_response = ""
# for chunk in response:
# chunk_message = chunk.choices[0].delta.content
# if chunk_message != '':
# full_response += chunk_message
self.new_system_message(content=full_response)
return self.response
except Exception as e:
self.error = True
if self.api_key == '' or self.api_key is None:
self.new_system_message(content="API key is missing")
else:
self.new_system_message(content=f"Unable to generate ChatCompletion response\nException: {e}")
return e