import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

pretrained_model = "datificate/gpt2-small-spanish"
tokenizer = AutoTokenizer.from_pretrained(pretrained_model, use_fast=True)
sonnets_tuned_model = 'hackathon-pln-es/gpt2-small-spanish-disco-poetry'
sonnets_pipe = pipeline('text2text-generation', model=sonnets_tuned_model, tokenizer=tokenizer)

def make_new_sonnet(prompt, max_lenght):
    ouputs = sonnets_pipe(prompt, max_length=max_lenght, 
                  num_beams=5, 
                  early_stopping=True, 
                  repetition_penalty=20.0, 
                  num_return_sequences=1)
    return ouputs[0]['generated_text']


description = """
<p>Este experimento pretende poner a prueba la capacidad de un modelo de inteligencia artificial para generar texto a partir de patrones observados en la hechicería del lenguaje poético en español.
Algo fascinante de lo que Borges habla en una conferencia de su Arte Poética: <i>"Uso la palabra «modelo» porque las metáforas que voy a citar, aunque parezcan muy distintas a la imaginación, para un lógico serían casi idénticas. Así que podríamos hablar de ellas como ecuaciones."</i>
"""
article = """
<p>
<h3>¿Por qué sonetos?</h3>
Este space usa un modelo gpt-2 entrenado con un corpus de sonetos del siglo 15 al 20 en español (DISCO), el cual que fue portado al formato de Datasets de Huggin Face. La elección de los sonetos obedece a una razón fundamental: su estilo y métrica específica, buscando entre otras cosas, tener consistencia en la rima y coherencia del texto generado.
</p>
<p>
<h3>¿Y para qué la poesía en la sociedad?</h3>
La poesía, <i>esa cosa liviana, alada y sagrada</i> como expresión artística ha acompañado la humanidad desde tiempos inmemoriales. Desde la oralidad, hasta la escritura. Lo que eso puede significar sobrepasa el alcance de este experimento. Pero podriamos decir, que esa constante obedece a unas necesidades humanas que rayan en lo espiritual, en la reflexión y en la contemplación de la bello. Este experimento , con la humildad del caso, puede ayudar a aumentar la creatividad de un escritor a la hora del acto creativo. Y en ese sentido contribuir a que hayan más creaciones bellas en el mundo que generan bienestar en las personas (3 objetivo de la ODS).
</p>
<p style='text-align: center'> 👉 Don't miss this other cool space based on different styles of poetry in spanish: <a href='https://huggingface.co/spaces/hackathon-pln-es/poem-generation-es' target='_blank'>poem-generation-es</a> 👈 </p>
"""
examples = [
    ['al despertar el dinosaurio todavía estaba ahí', 20 ],
    ['La inmensidad de tu imaginación', 140],
    ['con ansias de amores inflamada', 140]
]

iface = gr.Interface(fn=make_new_sonnet,
                      title= "Generador de poesía basada en sonetos en español",
                      description = description,
                      inputs=[
                              gr.inputs.Textbox(lines=2, placeholder="Escrbe algo para comenzar", label='Escribe algo para comenzar'),
                              gr.inputs.Slider(minimum = 20, maximum = 200, default = 140, step = 5, label='Salida máxima de caracteres')],
                      outputs=[
                               gr.outputs.Textbox(label="Tu poema"),
                              ],
                      article= article,
                      examples = examples,
                      theme = 'peach'
                     )
iface.launch(enable_queue=True)