Spaces:
Runtime error
Runtime error
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() |