nugentc's picture
add chat agent
245a96a
raw
history blame
2.24 kB
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration, T5Tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
grammar_tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')
grammar_model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector')
import torch
import gradio as gr
def chat(message, history, bot_input_ids):
history = history or []
bot_input_ids = bot_input_ids or []
new_user_input_ids = tokenizer.encode(message+tokenizer.eos_token, return_tensors='pt')
# append the new user input tokens to the chat history
bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids
# generated a response while limiting the total chat history to 1000 tokens,
chat_history_ids = model.generate(bot_input_ids, max_length=5000, pad_token_id=tokenizer.eos_token_id)
print("The text is ", [text])
# pretty print last ouput tokens from bot
reponse = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
history.append((message, response))
return history, bot_input_ids, feedback(message)
def feedback(text):
num_return_sequences=1
batch = grammar_tokenizer([text],truncation=True,padding='max_length',max_length=64, return_tensors="pt")
corrections= grammar_model.generate(**batch,max_length=64,num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)
print("The corrections are: ", corrections)
if len(corrections) == 0:
feedback = f'Looks good! Keep up the good work'
else:
suggestion = grammar_tokenizer.batch_decode(corrections[0], skip_special_tokens=True)
suggestion = [sug for sug in suggestion if '<' not in sug]
feedback = f'\'{" ".join(suggestion)}\' might be a little better'
return feedback
iface = gr.Interface(
chat,
["text", "state", "state"],
["chatbot", "state", "state", "text"],
allow_screenshot=False,
allow_flagging="never",
)
iface.launch()