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()