nugentc's picture
swap out grammar model
71f2227
raw
history blame
No virus
1.98 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):
history = history or []
if message.startswith("How many"):
response = random.randint(1, 10)
elif message.startswith("How"):
response = random.choice(["Great", "Good", "Okay", "Bad"])
elif message.startswith("Where"):
response = random.choice(["Here", "There", "Somewhere"])
else:
response = "I don't know"
history.append((message, response))
return history, feedback(message)
def feedback(text):
# tokenized_phrases = grammar_tokenizer([text], return_tensors='pt', padding=True)
# corrections = grammar_model.generate(**tokenized_phrases)
# corrections = grammar_tokenizer.batch_decode(corrections, skip_special_tokens=True)
batch = grammar_tokenizer([input_text],truncation=True,padding='max_length',max_length=64, return_tensors="pt").to(torch_device)
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 = tokenizer.batch_decode(corrections[0], skip_special_tokens=True)
feedback = f'\'{suggestion}\' might be a little better'
return f'FEEDBACK: {feedback}'
iface = gr.Interface(
chat,
["text", "state"],
["chatbot", "text"],
allow_screenshot=False,
allow_flagging="never",
)
iface.launch()