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> 
   El 28 de Junio de 2022 la Comisión de la verdad, una entidad del estado Colombiano que busca el esclarecimiento de los patrones y causas explicativas del conflicto armado interno, liberó el <a href="https://www.comisiondelaverdad.co/hay-futuro-si-hay-verdad">volumen testimonial</a> titulado <i>Cuando los pájaros No Cantaban. Historias del conflicto armado colombianos</i>. Una serie de historias del conflicto armado donde el lenguaje en español, con su diversidad y riqueza lingüística de las víctimas es protagonista. En este experimento fue entrenado un modelo gpt-2 con las más de 400 historias de víctimas del conflicto armado colombiano que conforma el volúmen testimonial del informe.
</p>
<p>
 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>
"""

article = """
<p>
<h3>Sobre el límite de la ficción en las narrativas históricas</h3>
Partiendo de que la empatía con las historias es algo natural en el contexto latinoamericano y de Colombia, este experimento parte de preguntarse ¿cuál es el límite de la ficción en las narrativas históricas?, ¿hasta qué punto se valen las licencias de la ficción si logra empatizar con el ciudadano del común sobre temas de importancia nacional?, ¿Un modelo de inteligencia artificial entrenado con el lenguaje de las víctimas, puede ayudar a aumentar el alcance de divulgación y empatía ciudadana con la memoria histórica, o por el contrario contribuye a distorsionarla? 
</p>
<p>
Sin pretensiones de arrojar afirmaciones concluyentes, la formulación de esas preguntas tienen como motivación un objetivo de interés nacional: la búsqueda de nuevas maneras de apoyar la difusión y  conocimiento de la memoria histórica Colombiana. Algo fundamental para construir una nación en torno a objetivos comunes.
</p>
<p>
Este espacio hace parte del proyecto open source Ask2Democracy (<a href= "https://github.com/jorge-henao/ask_to_democracy">repo en github con FastAPI</a>) que busca ayudar con el entendimiento de temas relevantes para el país, como las pasadas elecciones presidenciales 2022 y la memoria histórica Colombiana.  Para este experimento fue construido el <a href="https://huggingface.co/datasets/jorge-henao/historias_conflicto_colombia" target='_blank'>dataset con las historias portado en formato Hugging Face. </a>
</p> 
 
"""

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

iface = gr.Interface(fn=make_new_story,
                      title= "Ask2Democracy - Generador de historias basado en los testimonios del 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 = 1000, default = 280, step = 5, label='Tamaño de la historia')],
                      outputs=[
                               gr.outputs.Textbox(label="Tu historia"),
                              ],
                      examples = examples,
                      article = article,
                      theme = 'peach'
                     )
iface.launch(enable_queue=True)