Update app.py

app.py

@@ -1,22 +1,25 @@
 from transformers import *
 import gradio as gr
 import torch
-from googletrans import Translator
 model_category = AutoModelForCausalLM.from_pretrained("dquisi/story_spanish_gpt2_by_category")
 tokenizer_category = AutoTokenizer.from_pretrained("dquisi/story_spanish_gpt2_by_category")
 story_gen = pipeline("text-generation", "pranavpsv/gpt2-genre-story-generator")
-translator = Translator()
 tokenizer = AutoTokenizer.from_pretrained("dquisi/story_spanish_gpt2_v2")
 model = AutoModelForCausalLM.from_pretrained("dquisi/story_spanish_gpt2_v2")
+task_name_en_es = f"translation_en_to_es"
+model_name_en_es = f"Helsinki-NLP/opus-mt-en-es"
+task_name_es_en = f"translation_es_to_en"
+model_name_es_en = f"Helsinki-NLP/opus-mt-es-en"
+translator_es_en  = pipeline(task_name_en_es, model=model_name_en_es, tokenizer=model_name_en_es)
+translator_en_es = pipeline(task_name_es_en, model=model_name_es_en, tokenizer=model_name_es_en)
 def generate_story_translate(texto,longitud=250,categoria='superheroe'):
-  input_translate = translator.translate(texto, src='es', dest='en')
-  input_translate_category = translator.translate(categoria, src='es', dest='en')
-  query = "<BOS> <{0}> {1}".format(input_translate_category.text, input_translate.text)
+  translate_en = translator_es_en(texto)[0]["translation_text"]
+  translate_cat_en = translator_es_en(texto)[0]["translation_text"]
+  query = "<BOS> <{0}> {1}".format(translate_cat_en, translate_en)
   generated_text = story_gen(query, max_length=longitud,do_sample=True,repetition_penalty=1.1, temperature=1.2, top_p=0.95, top_k=50)
   generated_text = generated_text[0]['generated_text']
   generated_text = generated_text.split('> ')[2]
-  generate_translate = translator.translate(generated_text, src='en', dest='es')
-  return generate_translate.text
+  return translator_en_es(generated_text)[0]["translation_text"]
 def generate_story_custom(texto,longitud=250):
   query = "<BOS> <{0}>".format(texto)
   input_ids = tokenizer(query, return_tensors="pt")["input_ids"]
@@ -28,10 +31,10 @@ def generate_story_custom_category(texto,longitud=250,categoria='superheroe'):
   output = model_category.generate(input_ids, max_length=longitud,do_sample=True,repetition_penalty=1.1, temperature=1.2, top_p=0.95, top_k=50)
   return tokenizer_category.decode(output[0])
 
-contexto = gr.inputs.Textbox(lines=10, placeholder="Ingresar palabras claves para generar un cuento")
+contexto = gr.inputs.Textbox(lines=5, placeholder="Ingresar palabras claves para generar un cuento")
 categoria = gr.inputs.Textbox(lines=1, placeholder="Ingresar Categoria")
 longitud = gr.inputs.Slider(50, 500)
-opciones = gr.inputs.CheckboxGroup(["Generar", "Generar por Categoria", "Generar 2"])
+opciones = gr.inputs.CheckboxGroup(["Generar", "Generar por Categoria", "Generar V 2.0"])
 resultado = gr.outputs.HTML(label="Resultado")
 def generate_storie(contexto,opciones,categoria,longitud):
   resultado=""
@@ -40,7 +43,7 @@ def generate_storie(contexto,opciones,categoria,longitud):
     cuentos.append(generate_story_custom(contexto,longitud))
   if "Generar por Categoria" in opciones:
     cuentos.append(generate_story_custom_category(contexto,longitud,categoria))
-  if "Generar 2" in opciones:
+  if "Generar V 2.0" in opciones:
     cuentos.append(generate_story_translate(contexto,longitud,categoria))  
   resultado += "<p><b>Generados:</b> "+'<b>Cuento: <b/>'.join(cuentos)+"</p>" 
   return resultado