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)
tuned_model = 'jorge-henao/gpt2-small-spanish-historias-conflicto-col'
sonnets_pipe = pipeline('text2text-generation', model=tuned_model, tokenizer=tokenizer)

def make_new_story(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> 
    <br></br>
    <br>Este espacio hace parte de un proyecto open source que busca ayudar con el entendimiento de temas relevantes para el país, como estas elecciones y la memoria histórica Colombiana. <a href= "https://github.com/jorge-henao/ask_to_democracy">repo en github con FastAPI</a></br>
    Por: 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> 
"""

examples = [
    ['cuando salí no había nadie', 20 ],
    ['La última vez que la vi', 140],
    ['LLegaron y mi vida se fue', 140]
]

iface = gr.Interface(fn=make_new_story,
                      title= "Generador de historias basadas en el conflicto Colombiano",
                      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"),
                              ],
                      examples = examples,
                      theme = 'default'
                     )
iface.launch(enable_queue=True)