DialoGPT / app.py
hf-dongpyo's picture
Update app.py
875c396
raw
history blame contribute delete
No virus
1.98 kB
from transformers import AutoModelForCausalLM, AutoTokenizer, BlenderbotForConditionalGeneration
import torch
chat_tkn = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium')
mdl = AutoModelForCausalLM.from_pretrained('microsoft/DialoGPT-medium')
# chat_tkn = AutoTokenizer.from_pretrained('facebook/blenderbot-400M-distill')
# mdl = BlenderbotForConditionalGeneration.from_pretrained('facebook/blenderbot-400M-distill')
def converse(user_input, chat_history = []):
user_input_ids = chat_tkn(user_input + chat_tkn.eos_token, return_tensors = 'pt').input_ids
# keep history in the tensor
bot_input_ids = torch.cat([torch.LongTensor(chat_history), user_input_ids], dim = -1)
# get response
chat_history = mdl.generate(
bot_input_ids,
max_length = 1000,
pad_token_id = chat_tkn.eos_token_id
).tolist()
print(chat_history)
response = chat_tkn.decode(chat_history[0]).split("<|endoftext|>")
print("staring to print response")
print(response)
# html for display
html = "<div class='mybot'>"
for x, mesg in enumerate(response):
if x % 2 != 0:
mesg = "Alicia: " + mesg
clazz = "alicia"
else:
clazz = "user"
print("value of x")
print(x)
print("message")
print(mesg)
html += "<div class='mesg {}'> {}</div>".format(clazz, mesg)
html += "</div>"
print(html)
return html, chat_history
import gradio as grad
css = """
.mychat {display:flex;flex-direction:column}
.mesg {padding:5px;margin-bottom:5px;border-radius:5px;width:75%}
.mesg.user {background-color:lightblue;color:white}
.mesg.alicia {background-color:orange;color:white,align-self:self-end}
.footer {display:none !important}
"""
text = grad.inputs.Textbox(placeholder = "Lets chat")
grad.Interface(
fn = converse,
theme = "default",
inputs = [text, "state"],
outputs = ["html", "state"],
css = css
).launch()