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