Aira-Demo / app.py
nicholasKluge's picture
Create app.py
6ddba8c
raw
history blame
1.25 kB
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import gradio as gr
tokenizer = GPT2Tokenizer.from_pretrained('nicholasKluge/Aira-Instruct-124M')
model = GPT2LMHeadModel.from_pretrained('nicholasKluge/Aira-Instruct-124M')
import gradio as gr
title = "AIRA Demo 🤓"
with gr.Blocks() as demo:
chatbot = gr.Chatbot()
msg = gr.Textbox()
clear = gr.Button("Clear Conversation")
def respond(message, chat_history):
inputs = tokenizer(tokenizer.bos_token + message + tokenizer.eos_token, return_tensors="pt")
response = model.generate(**inputs,
bos_token_id=tokenizer.bos_token_id,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
do_sample=True,
early_stopping=True,
top_k=50,
max_length=200,
top_p=0.95,
temperature=0.7,
num_return_sequences=1)
chat_history.append((f"👤 {message}", f"""🤖 {tokenizer.decode(response[0], skip_special_tokens=True).replace(message, "")}"""))
return "", chat_history
msg.submit(respond, [msg, chatbot], [msg, chatbot])
clear.click(lambda: None, None, chatbot, queue=False)
demo.launch()