import random
import gradio as gr
from transformers import AutoTokenizer, pipeline, T5ForConditionalGeneration

model_name = "versae/byt5-base-finetuned-modernisa"  # "versae/modernisa-pre"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def modernisa(lines=None, file_obj=None):
    if file_obj:
        with open(file_obj.name) as file:
            lines = file.read()
    lines = [line for line in lines.strip().split('\n') if line.strip()]
    text2text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
    outputs = []
    for line in lines:
        outputs += text2text_generator([line], max_length=150)
    return "<br/>".join(output["generated_text"] for output in outputs)


samples = [
"""Otra vez, Don Iuan, me dad,
y otras mil vezes los braços.
Otra, y otras mil sean lazos
de nuestra antigua amistad.
Como venis?
Yo me siento
tan alegre, tan vfano,
tan venturoso, tan vano,
que no podrà el pensamiento
encareceros jamàs
las venturas que posseo,
porque el pensamiento creo"""
]

gr.Interface(
    fn=modernisa,
    inputs=[
        gr.inputs.Textbox(lines=12, label="Enter Old Spanish", default=random.choice(samples), optional=True),
        gr.inputs.File(file_count="single", label="Or upload a plain text file (.txt)" type="file", optional=True),
    ],
    outputs=gr.outputs.HTML(label="Modern Spanish"),
).launch(inline=False)