Add the logic to generate the poem

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import random
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 author_set = {'Leopoldo Lugones', 'Nacho Buzón', 'Octavio Paz', 'Luis Cañizal de la Fuente', 'Juan de Salinas', 'Vicente Huidobro', 'Hilario Barrero',
               'Ramón de Campoamor', 'Anna Ajmátova', 'Víctor Hugo López Cancino', 'Ramón María del Valle-Inclán', 'Infantiles', 'Jorge Luis Borges',
@@ -45,12 +47,52 @@ author_set = {'Leopoldo Lugones', 'Nacho Buzón', 'Octavio Paz', 'Luis Cañizal
               'Jorge Teillier', 'Félix María de Samaniego', 'Nicolás Fernández de Moratín', 'Juan Boscán', 'Manuel María Flores', 'Gutierre de Cetina',
               'Alfonsina Storni', 'José Luis Rey Cano', 'Jorge Manrique', 'Nicanor Parra'}
 
+model_name = 'hackathon-pln-es/poem-gen-spanish-t5-small'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+
+def make_poem(author, sentiment, words, text):
+  num_lines=5
+  poem = text
+  prev_output = ''
+  l_words = [x.strip() for x in words.split(',')]
+
+  # Add empty strings to words
+  if num_lines > len(l_words):
+    diff = num_lines - len(l_words)
+    l_words += [''] * diff
+
+  random.shuffle(l_words)
+
+  for i in range(num_lines):
+    word = l_words[i]
+    if word == '':
+      input_text = f"""poema: estilo: {author} && sentimiento: {sentiment} && texto: {poem} """
+    else:
+      input_text = f"""poema: estilo: {author} && sentimiento: {sentiment} && palabras: {word} && texto: {poem} """
+    inputs = tokenizer(input_text, return_tensors="pt")
+
+    outputs = model.generate(inputs["input_ids"],
+                            do_sample = True,
+                            max_length = 30,
+                            repetition_penalty = 20.0,
+                            top_k = 50,
+                            top_p = 0.92)
+    detok_outputs = [tokenizer.decode(x, skip_special_tokens=True) for x in outputs]
+    pre_output = detok_outputs[0]
+
+    poem += '\n' + pre_output
+  # audio = tts_es(poem)
+  # return poem, audio
+  return poem
+
 
 def poem_generate(author, sentiment, words, text):
-    debug_poem = 'Y cuando se despertó, el dinosaurio seguía ahí'
+    poem_txt = make_poem(author, sentiment, words, text)
     # Pending: Translate poem to English, so that text can be the input of the latentdiffussion
-    gen_image = poem_to_image(debug_poem)
-    return debug_poem, gen_image
+    poem_img = poem_to_image(poem_txt)
+    return poem_txt, poem_img
 
 
 def poem_to_image(poem):
@@ -91,7 +133,6 @@ with gr.Blocks() as demo:
             b1 = gr.Button("Generate Poem & Illustration")
             
         output_poem_txt = gr.Textbox(lines=7, label='Poema generado')
-        output_poem_txt2 = gr.Textbox(lines=7, label='Poema generado 2')
         output_image = gr.Image(type="filepath", shape=(256, 256))
 
     b1.click(poem_generate, inputs=[input_author, input_sentiment, input_include_words, input_initial_text], outputs=[output_poem_txt, output_image])