text-iterater / app.py
Vaibhav Srivastav
adding examples
4ffb297
from cProfile import label
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("wanyu/IteraTeR-PEGASUS-Revision-Generator")
model = AutoModelForSeq2SeqLM.from_pretrained("wanyu/IteraTeR-PEGASUS-Revision-Generator")
def prep_input(text):
text = text.strip()
clarity_input = "<clarity> " + text
fluency_input = "<fluency> " + text
coherence_input = "<coherence> " + text
style_input = "<style> " + text
return [clarity_input, fluency_input, coherence_input, style_input]
def get_model_output(text):
model_input = tokenizer(text, return_tensors='pt')
model_outputs = model.generate(**model_input, num_beams=8, max_length=1024)
pred = tokenizer.batch_decode(model_outputs, skip_special_tokens=True)[0]
return pred
def return_predictions(text):
all_predictions = []
prepped_input = prep_input(text)
for input in prepped_input:
all_predictions.append(get_model_output(input))
return all_predictions[0], all_predictions[1], all_predictions[2], all_predictions[3]
iface = gr.Interface(fn=return_predictions,
inputs=gr.inputs.Textbox(label="Sentence/ Paragraph"),
outputs = [gr.outputs.Textbox(label="Clarity"),
gr.outputs.Textbox(label="Fluency"),
gr.outputs.Textbox(label="Coherence"),
gr.outputs.Textbox(label="Style")],
title="IteraTeR: Understanding Iterative Revision from Human-Written text",
description = "The model (pegasus-large) generates a revised sentence based on a given intention!",
layout = "horizontal",
examples = ["The changes made the paper better than before.",
"She went to the markt",
"She works hard. She is successful.",
"Everything was rotten."],
theme="huggingface",
enable_queue=True)
iface.launch()