|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import gradio as gr |
|
import torch |
|
|
|
title = "EZChat" |
|
description = "A State-of-the-Art Large-scale Pretrained Response generation model (DialoGPT-medium)" |
|
examples = [["How are you?"]] |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") |
|
tokenizer.padding_side = 'left' |
|
tokenizer.add_special_tokens({'pad_token': '[EOS]'}) |
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium") |
|
|
|
|
|
def predict(input, history=[]): |
|
|
|
new_user_input_ids = tokenizer.encode( |
|
input + tokenizer.eos_token, padding=True, truncation=True, return_tensors="pt" |
|
) |
|
|
|
|
|
bot_input_ids = torch.cat([torch.tensor(history), new_user_input_ids], dim=-1) if history else new_user_input_ids |
|
|
|
|
|
chat_history_ids = model.generate( |
|
bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id |
|
) |
|
|
|
|
|
response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True) |
|
|
|
return response, chat_history_ids.tolist()[0] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
title=title, |
|
description=description, |
|
examples=examples, |
|
inputs=["text", gr.inputs.Slider(0, 4000, default=2000, label='Chat History')], |
|
outputs=["text", "text"], |
|
theme="ParityError/Anime", |
|
) |
|
|
|
iface.launch() |