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 <a href="https://arxiv.org/abs/1706.03762" target='blank'>patrones observados</a> en la <i>hechicería</i> 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>
<p>
Contacto: Jorge Henao 🇨🇴 <a href="https://twitter.com/jhenaotw" target='_blank'>Twitter</a> <a href="https://www.linkedin.com/in/henaojorge" target='_blank'/>LinkedIn</a>
</p>
"""
article = """
<p>
<h3>¿Por qué sonetos?</h3>
Este espacio 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 Hugging 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 en el texto generado.
</p>
<p>
<h3>Sobre el impacto de la inteligencia artificial en el arte</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 puede significar, que la inteligencia artificial se involucre en una expresión artística que raya en lo espiritual, la reflexión y la contemplación de lo bello, sobrepasa el alcance de este experimento. ¿Cuál es el impacto del arte generado con inteligencia artificial?, ¿cómo se puede medir ese impacto?, éstas son algunas de las preguntas a lugar.
<br/>Por otro lado, está la simbiosis humano-máquina a la hora de co-crear en conjunto. Las posibilidades mostradas en este experimento, también  ilustran un escenario donde la inteligencia artificial aumenta y/o complementa la creatividad de un escritor a la hora del acto creativo. Contribuyendo así, a que hayan más creaciones bellas en el mundo que generan bienestar en las personas (3 objetivo ODS).
La discusión está abierta y este espacio pretende contribuir con esa reflexión.
</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.Textbox(lines=2, placeholder="Escrbe algo para comenzar", label='Escribe algo para comenzar'),
                              gr.Slider(minimum = 20, maximum = 200, value = 140, step = 5, label='Salida máxima de caracteres')],
                      outputs=[
                               gr.Textbox(label="Tu poema"),
                              ],
                      article= article,
                      examples = examples,
                      theme = 'peach'
                     )
iface.queue(max_size=10)
iface.launch()