|
import gradio as gr |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
model_name = "distilgpt2" |
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
|
def respond(message, history, system_message, max_tokens, temperature, top_p): |
|
messages = [{"role": "system", "content": system_message}] |
|
|
|
for val in history: |
|
if val[0]: |
|
messages.append({"role": "user", "content": val[0]}) |
|
if val[1]: |
|
messages.append({"role": "assistant", "content": val[1]}) |
|
|
|
messages.append({"role": "user", "content": message}) |
|
|
|
input_text = "\n".join([msg["content"] for msg in messages]) |
|
|
|
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True) |
|
|
|
outputs = model.generate( |
|
inputs["input_ids"], |
|
max_length=max_tokens, |
|
temperature=temperature, |
|
top_p=top_p, |
|
do_sample=True, |
|
) |
|
|
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
response += "\nСделано больницей EMS штата Alta!" |
|
return response |
|
|
|
|
|
demo = gr.Interface( |
|
fn=respond, |
|
inputs=[ |
|
gr.Textbox(value="Здравствуйте. Отвечай кратко...", label="System message"), |
|
gr.Slider(minimum=1, maximum=2048, value=512, label="Max Tokens"), |
|
gr.Slider(minimum=0.1, maximum=4.0, value=0.7, label="Temperature"), |
|
gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p"), |
|
], |
|
outputs="text", |
|
) |
|
|
|
demo.launch() |
|
|