Update app.py

app.py

@@ -1,54 +1,75 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import random
 
 # Cargar el modelo y el tokenizador
-model_name = "bigscience/bloom-560m"
+model_name = "microsoft/DialoGPT-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def chatbot(input, history):
-    history = history or []
-    history.append({"role": "user", "content": input})
+# Caché simple para respuestas frecuentes
+response_cache = {}
+
+# Lista de respuestas predefinidas para variar la conversación
+fallback_responses = [
+    "Interesante. ¿Puedes decirme más sobre eso?",
+    "Entiendo. ¿Cómo te hace sentir eso?",
+    "¿Qué te llevó a pensar en eso?",
+    "Es una perspectiva interesante. ¿Has considerado otras alternativas?",
+    "Me gustaría saber más. ¿Puedes elaborar un poco?",
+]
+
+def get_response(input_text, conversation_history):
+    # Verificar si la respuesta está en caché
+    if input_text in response_cache:
+        return response_cache[input_text]
     
-    chat_history = ""
-    for message in history:
-        if message["role"] == "user":
-            chat_history += f"Human: {message['content']}\n"
-        else:
-            chat_history += f"AI: {message['content']}\n"
+    # Limitar la longitud de la conversación
+    if len(conversation_history) > 5:
+        conversation_history = conversation_history[-5:]
     
-    chat_history += "AI:"
+    # Preparar el input para el modelo
+    bot_input_ids = tokenizer.encode(conversation_history + input_text + tokenizer.eos_token, return_tensors='pt')
     
-    input_ids = tokenizer.encode(chat_history, return_tensors="pt")
-    attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
+    # Generar respuesta
+    chat_response_ids = model.generate(
+        bot_input_ids, 
+        max_length=1000, 
+        pad_token_id=tokenizer.eos_token_id,
+        no_repeat_ngram_size=3,
+        do_sample=True,
+        top_k=100,
+        top_p=0.7,
+        temperature=0.8
+    )
     
-    max_length = input_ids.shape[1] + 50
+    chat_response = tokenizer.decode(chat_response_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
     
-    output = model.generate(
-        input_ids, 
-        attention_mask=attention_mask,
-        max_length=max_length,
-        num_return_sequences=1,
-        no_repeat_ngram_size=2,
-        temperature=0.7
-    )
+    # Si la respuesta es vacía o muy corta, usar una respuesta predefinida
+    if not chat_response or len(chat_response.split()) < 3:
+        chat_response = random.choice(fallback_responses)
     
-    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
+    # Guardar en caché
+    response_cache[input_text] = chat_response
     
-    history.append({"role": "assistant", "content": response.strip()})
+    return chat_response
+
+def chatbot(input_text, history):
+    history = history or []
+    conversation_history = " ".join([f"{h[0]} {h[1]}" for h in history])
     
-    # Convertir el historial al formato que Gradio espera
-    gradio_history = [[m["content"], h["content"]] for m, h in zip(history[::2], history[1::2])]
+    response = get_response(input_text, conversation_history)
     
-    return gradio_history, history
+    history.append((input_text, response))
+    return history, history
 
 iface = gr.Interface(
     fn=chatbot,
     inputs=["text", "state"],
     outputs=["chatbot", "state"],
-    title="Tu Compañero AI con BLOOM",
-    description="Un chatbot de IA diseñado para simular conversaciones personales, utilizando el modelo BLOOM.",
+    title="Tu Compañero AI Mejorado",
+    description="Un chatbot de IA diseñado para simular conversaciones personales de manera rápida y coherente.",
 )
 
 iface.launch()