initial commit v2

app.py

@@ -1,7 +1,40 @@
+from cProfile import label
 import gradio as gr
+import transformers
 
-def greet(name):
-    return "Hello " + name + "!!"
+tokenizer = AutoTokenizer.from_pretrained("wanyu/IteraTeR-PEGASUS-Revision-Generator")
+model = AutoModelForSeq2SeqLM.from_pretrained("wanyu/IteraTeR-PEGASUS-Revision-Generator")
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+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 = "Comparing greedy decoder with beam search CTC decoder, record/ drop your audio!",
+             layout = "horizontal", 
+             theme="huggingface",
+             enable_queue=True)
 iface.launch()