persona-private / app.py
NeonBohdan's picture
Fixed Gemini counts letters instead of tokens
0dcb7a0
raw
history blame contribute delete
No virus
2.6 kB
import gradio as gr
from neon_llm_chatgpt.chatgpt import ChatGPT
from neon_llm_palm2.palm2 import Palm2
from neon_llm_gemini.gemini import Gemini
from neon_llm_claude.claude import Claude
from utils import convert_history, PersonaConverter
import os
role = ""
biography = ""
key = os.environ['OPENAI_API_KEY']
config = {
"key": key,
"model": "gpt-3.5-turbo",
"role": role,
"context_depth": 3,
"max_tokens": 256,
}
chatgpt = ChatGPT(config)
key_google = os.environ['GOOGLE_API_KEY']
key_path = "./vertex.json"
with open(key_path, "w") as json_file:
json_file.write(key_google)
config = {
"key_path": key_path,
"role": role,
"context_depth": 3,
"max_tokens": 256,
}
palm2 = Palm2(config)
config = {
"model": "gemini-pro",
"key_path": key_path,
"role": role,
"context_depth": 3,
"max_tokens": 1024,
}
gemini = Gemini(config)
key_anthropic = os.environ['ANTHROPIC_API_KEY']
config = {
"key": key_anthropic,
"openai_key": key,
"model": "claude-2",
"role": role,
"context_depth": 3,
"max_tokens": 256,
}
claude = Claude(config)
model_choices = {
"openai": chatgpt,
"palm2": palm2,
"gemini": gemini,
"anthropic": claude,
}
model_choices_list = list(model_choices.keys())
personaConverter = PersonaConverter()
preset_choices_list = list(personaConverter.personas.keys())
figures_choices_list = list(personaConverter.historical_figures.keys())
def ask(message, history, persona, model_name, preset, biography, figure, imaginary, short):
chat_history = convert_history(history)
model = model_choices[model_name]
persona_description = personaConverter(name = preset, figure = figure,
description = persona, biography = biography,
imaginary = imaginary, short = short)
responce = model.ask(message, chat_history, persona = {"description": persona_description})
return responce
demo = gr.ChatInterface(ask,
additional_inputs=[
gr.Textbox(role, label="Persona"),
gr.Dropdown(choices=model_choices_list, value=model_choices_list[0], label="Model"),
gr.Dropdown(choices=preset_choices_list, value=preset_choices_list[0], label="Preset"),
gr.Textbox(biography, label="Biography"),
gr.Dropdown(choices=figures_choices_list, value=figures_choices_list[0], label="Historical Figures"),
gr.Checkbox(value=False, label="Imaginary"),
gr.Checkbox(value=True, label="Short response"),
]
)
if __name__ == "__main__":
demo.queue().launch()