ConorDY/feedback-chatbot

1from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration, T5Tokenizer2tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")3model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")4grammar_tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')5grammar_model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector')6import torch7import gradio as gr8 9 10 11def chat(message, history=[]):12    new_user_input_ids = tokenizer.encode(message+tokenizer.eos_token, return_tensors='pt')13    if  len(history) > 0:14        last_set_of_ids = history[len(history)-1][2]15        bot_input_ids = torch.cat([last_set_of_ids, new_user_input_ids], dim=-1) 16    else:17        bot_input_ids = new_user_input_ids18    chat_history_ids = model.generate(bot_input_ids, max_length=5000, pad_token_id=tokenizer.eos_token_id)19    response_ids = chat_history_ids[:, bot_input_ids.shape[-1]:][0]20    response = tokenizer.decode(response_ids, skip_special_tokens=True)21    history.append((message, response, chat_history_ids))22    return history, history, feedback(message)23 24 25def feedback(text):26    num_return_sequences=127    batch =  grammar_tokenizer([text],truncation=True,padding='max_length',max_length=64, return_tensors="pt")28    corrections = grammar_model.generate(**batch,max_length=64,num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)29    corrected_text = grammar_tokenizer.decode(corrections[0], clean_up_tokenization_spaces=True, skip_special_tokens=True)30    print("The corrected text is: ", corrected_text)31    print("The orig text is: ", text)32    if corrected_text.rstrip('.') == text.rstrip('.'):33    # if corrected_text == text:34        feedback = f'Looks good! Keep up the good work'35    else:36        feedback = f'\'{corrected_text}\' might be a little better'37    return feedback38 39 40title = "A chatbot that provides grammar feedback"41description = "A quick proof of concept using Gradio"42article = "<p style='text-align: center'><a href='https://docs.google.com/presentation/d/11fiO91MKZVgNoQJh5pn3Tw8-inHe6XbWYB2r1f701WI/edit?usp=sharing'> A conversational agent for Language learning</a> | <a href='https://github.com/ConorNugent/gradio-chatbot-demo'>Github Repo</a></p>"43examples = [44            ["Have you read the play what I wrote?"],45            ["Were do you live?"],46]47 48iface = gr.Interface(49    chat,50    [gr.Textbox(label="Send messages here"), "state"],51    [gr.Chatbot(label='Conversation'), "state", gr.Textbox(52            label="Feedback",53            lines=154        )],55    allow_screenshot=False,56    allow_flagging="never",57    title=title, 58    description=description, 59    article=article, 60    examples=examples)61iface.launch()62