Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+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,8 +46,43 @@ 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'}
 
-def make_poem(author, sentiment, words, text):
-    return [author, sentiment, words, text]
+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, text, sentiment, words):
+  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]
+
+    pre_output = pre_output.replace('.', ' ')
+    poem += '\n' + pre_output
+  return poem
 
 iface = gr.Interface(
     fn=make_poem,