import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

choices = ['Formal to Informal', 'Informal to Formal']

tokenizer_1 = AutoTokenizer.from_pretrained('prithivida/formal_to_informal_styletransfer')
model_1 = AutoModelForSeq2SeqLM.from_pretrained('prithivida/formal_to_informal_styletransfer')

tokenizer_2 = AutoTokenizer.from_pretrained('prithivida/informal_to_formal_styletransfer')
model_2 = AutoModelForSeq2SeqLM.from_pretrained('prithivida/informal_to_formal_styletransfer')

def model_selection(choices, text):
    if choices == "Formal to Informal":
        inputs = tokenizer_1.encode(text, return_tensors = 'pt')
        outputs = model_1.generate(inputs)
        clean = tokenizer_1.decode(outputs[0])
        
    else:
        inputs = tokenizer_2.encode(text, return_tensors = 'pt')
        outputs = model_2.generate(inputs)
        clean = tokenizer_2.decode(outputs[0])
        
    replace = clean.replace('<pad>', '').replace('</s>', '')
    return replace    

input_1 = gr.inputs.Radio(choices = choices, label='Choose a model.')
input_2 = gr.inputs.Textbox(placeholder='Enter your text here...', label = 'Input')

article = "<p style='text-align: center'><a href='https://github.com/PrithivirajDamodaran/Styleformer'>Styleformer GitHub</a></p>"
examples = [["Formal to Informal", "Your mother is so old that her last name is asaurus."], ["Informal to Formal", "Yo what's up with the weather?"]]

iface = gr.Interface(
    model_selection,
    [input_1, input_2], "text", theme = 'huggingface', article=article, examples = examples)

if __name__ == "__main__":
    iface.launch()