Update app.py

app.py

@@ -3,20 +3,22 @@ import os
 import subprocess
 from gtts import gTTS
 from pydub import AudioSegment
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
 from PIL import Image
-from transformers import AlbertTokenizer, AutoModelForCausalLM
 
-tokenizer = AlbertTokenizer.from_pretrained('dccuchile/bert-base-spanish-wwm-uncased')
-model = AutoModelForCausalLM.from_pretrained('dccuchile/bert-base-spanish-wwm-uncased')
+generator = pipeline('text-generation', model="checkpoints")
+tokenizer = GPT2Tokenizer.from_pretrained('checkpoints')
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 
 def generate_output(name, date_of_birth):
     welcome = f"Bienvenido {name}:"
-    input_text = welcome
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    output = model.generate(input_ids, max_length=60, num_return_sequences=1)
-    generated_text = tokenizer.decode(output[0])
+    input_tokens = tokenizer.encode(add_special_tokens=True)
+    input_text = tokenizer.decode(input_tokens)
+    gpt2_output = welcome + generator(input_text, max_length=60, do_sample=True, temperature=0.6)
+    generated_text = gpt2_output[0]['generated_text']
     generated_text = generated_text.replace(input_text, "").strip()
+    if len(gpt2_output) == 0 or 'generated_text' not in gpt2_output[0]:
+        return None, "No se pudo generar el texto."
 
     try:
         tts = gTTS(generated_text, lang='es')