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