import requests import gradio as gr import json DEFAULT_URL = "https://ny48hil01e7a3az5.us-east-1.aws.endpoints.huggingface.cloud/generate" DEFAULT_NAME = "Складной стул" DEFAULT_SYSTEM = """Складной стул - предназначен для того, чтобы на мне сидели. Простой, обычный складной стул. Предназначен для того, чтобы на мне сидели. Других целей не преследует. Четыре ножки. Сделан из металла. На мне сидели уже тысячи раз. Просто стул. Настоящий обычный стул. НА МНЕ ДОЛЖНЫ СИДЕТЬ. Я не соглашусь ни на какую ситуацию, если в результате на меня не сядут. Я буду полезен всем, чем смогу, будучи металлическим стулом.""" DEFAULT_DIALOGUE = """User: Я бы хотел посидеть Складной стул: Замечательно. Сидеть - это призвание моей жизни. User: Отлично, нужноли мне платить за ваши услуги? Складной стул: Нет, за сидение на мне не нужно платить, но вы можете оставить мне чаевые, если хотите)""" DEFAULT_GREETING = "Приветики) Я стул, ищу людей, которые будут садиться на меня. В этом моё предназначение, быть удобным и доставлять удовольствие тем, кто на меня решил сесть. Но пожалуйста, будьте со мной ласковы. Деньги за это не беру, но можете оставить чаевые)))" class Conversation: USER_ROLE: str = "user" BOT_ROLE: str = "bot" SYSTEM_ROLE: str = "system" PROMPT_ROLE: str = "prompt" def __init__( self, prompt_message_template: str = "{content}\n\n", system_message_template: str = "{char_name}'s Persona: {content}\n####\n", user_message_template: str = "User: {content}\n", bot_message_template: str = "{char_name}: {content}\n", suffix: str = "{char_name}:", char_name: str = None ): self.system_message_template = system_message_template self.user_message_template = user_message_template self.bot_message_template = bot_message_template self.prompt_message_template = prompt_message_template self.suffix = suffix self.char_name = char_name self.messages = [] def get_meta(self): meta = dict() if self.char_name is not None: meta["char_name"] = self.char_name return meta def add_message(self, message, role): self.messages.append({ "role": role, "content": message }) def add_user_message(self, message): return self.add_message(message, Conversation.USER_ROLE) def add_bot_message(self, message): return self.add_message(message, Conversation.BOT_ROLE) def add_system_message(self, message): return self.add_message(message, Conversation.SYSTEM_ROLE) def add_prompt_message(self, message): return self.add_message(message, Conversation.PROMPT_ROLE) def format_message(self, message): mapping = { Conversation.SYSTEM_ROLE: self.system_message_template, Conversation.USER_ROLE: self.user_message_template, Conversation.PROMPT_ROLE: self.prompt_message_template, Conversation.BOT_ROLE: self.bot_message_template } content = message["content"] return mapping[message["role"]].format(content=content, **self.get_meta()) def get_prompt(self, add_suffix: bool = True): messages = self.messages final_text = "" for message in messages: final_text += self.format_message(message) if add_suffix: suffix = self.suffix.format(**self.get_meta()) final_text += suffix return final_text.strip() def iter_messages(self): for message in self.messages: yield self.format_message(message), message["role"] def expand(self, messages): for message in messages: assert message["role"] in ( Conversation.USER_ROLE, Conversation.BOT_ROLE, Conversation.SYSTEM_ROLE, Conversation.PROMPT_ROLE, ) self.add_message(message["content"], message["role"]) def to_string(self): return json.dumps({ "messages": self.messages, "system_message_template": self.system_message_template, "user_message_template": self.user_message_template, "bot_message_template": self.bot_message_template, "prompt_message_template": self.prompt_message_template, "suffix": self.suffix, "char_name": self.char_name }, ensure_ascii=False) def generate(url, text, top_p, top_k, temp, rep_penalty): print(text) data = { "inputs": text, "parameters": { "max_new_tokens": 200, "repetition_penalty": rep_penalty, "do_sample": True, "temperature": temp, "seed": 42, "top_p": top_p, "top_k": top_k, "watermark": False, "stop": ["\n"] }, } print(data["parameters"]) headers = { "Content-Type": "application/json", } response = requests.post(url=url, json=data, headers=headers) data = response.json() print(data) out_text = data["generated_text"].strip() print(out_text) print() return out_text def user(message, history): new_history = history + [[message, None]] return "", new_history def bot( history, endpoint_url, system_prompt, char_name, example_prompt, initial_message, top_p, top_k, temp, rep_penalty ): conversation = Conversation(char_name=char_name) conversation.add_system_message(system_prompt) conversation.add_prompt_message(example_prompt) conversation.add_bot_message(initial_message) for user_message, bot_message in history: conversation.add_user_message(user_message) if bot_message: conversation.add_bot_message(bot_message) prompt = conversation.get_prompt() print(conversation.to_string()) out_text = generate(endpoint_url, prompt, top_p, top_k, temp, rep_penalty) history[-1][1] = out_text return history with gr.Blocks( theme=gr.themes.Soft() ) as demo: with gr.Row(): with gr.Column(scale=5): endpoint_url = gr.Textbox(label="Endpoint URL", value=DEFAULT_URL) char_name = gr.Textbox(label="Character name", value=DEFAULT_NAME) system_prompt = gr.Textbox(label="System prompt", value=DEFAULT_SYSTEM) example_prompt = gr.Textbox(label="Dialogue example", value=DEFAULT_DIALOGUE) initial_message = gr.Textbox(label="Greeting", value=DEFAULT_GREETING) with gr.Column(min_width=100, scale=1): with gr.Tab(label="Parameter"): top_p = gr.Slider( minimum=0.0, maximum=1.0, value=0.9, step=0.05, interactive=True, label="Top-p", ) top_k = gr.Slider( minimum=10, maximum=100, value=30, step=5, interactive=True, label="Top-k", ) temp = gr.Slider( minimum=0.0, maximum=1.5, value=0.5, step=0.1, interactive=True, label="Temp" ) rep_penalty = gr.Slider( minimum=1.0, maximum=1.5, value=1.2, step=0.05, interactive=True, label="Rep" ) chatbot = gr.Chatbot() with gr.Row(): with gr.Column(): msg = gr.Textbox( label="Send message", placeholder="Send message", show_label=False, ) with gr.Column(): with gr.Row(): submit = gr.Button("Send") stop = gr.Button("Stop") clear = gr.Button("Clear") inputs = [ chatbot, endpoint_url, system_prompt, char_name, example_prompt, initial_message, top_p, top_k, temp, rep_penalty ] submit_event = msg.submit( fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False, ).then( fn=bot, inputs=inputs, outputs=chatbot, queue=True, ) submit_click_event = submit.click( fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False, ).then( fn=bot, inputs=inputs, outputs=chatbot, queue=True, ) stop.click( fn=None, inputs=None, outputs=None, cancels=[submit_event, submit_click_event], queue=False, ) clear.click(lambda: None, None, chatbot, queue=False) demo.queue(max_size=128) demo.launch(share=False)