Update app.py

app.py

@@ -1,75 +1,48 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import random
 
 # Cargar el modelo y el tokenizador
-model_name = "microsoft/DialoGPT-small"
+model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 
-# Caché simple para respuestas frecuentes
-response_cache = {}
+def generate_response(prompt, max_length=200):
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs, 
+            max_length=max_length, 
+            num_return_sequences=1,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.strip()
 
-# 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]
-    
-    # Limitar la longitud de la conversación
-    if len(conversation_history) > 5:
-        conversation_history = conversation_history[-5:]
-    
-    # Preparar el input para el modelo
-    bot_input_ids = tokenizer.encode(conversation_history + input_text + tokenizer.eos_token, return_tensors='pt')
-    
-    # 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
-    )
-    
-    chat_response = tokenizer.decode(chat_response_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    
-    # 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)
-    
-    # Guardar en caché
-    response_cache[input_text] = chat_response
-    
-    return chat_response
-
-def chatbot(input_text, history):
+def chatbot(message, history):
     history = history or []
-    conversation_history = " ".join([f"{h[0]} {h[1]}" for h in history])
     
-    response = get_response(input_text, conversation_history)
+    # Construir el prompt en el formato que Mixtral espera
+    prompt = "Eres un asistente AI amigable y útil. Responde de manera concisa y coherente.\n\n"
+    for human, ai in history:
+        prompt += f"Human: {human}\nAssistant: {ai}\n"
+    prompt += f"Human: {message}\nAssistant:"
+    
+    response = generate_response(prompt)
     
-    history.append((input_text, response))
+    history.append((message, response))
     return history, history
 
 iface = gr.Interface(
     fn=chatbot,
     inputs=["text", "state"],
     outputs=["chatbot", "state"],
-    title="Tu Compañero AI Mejorado",
-    description="Un chatbot de IA diseñado para simular conversaciones personales de manera rápida y coherente.",
+    title="Tu Compañero AI con Mixtral",
+    description="Un chatbot de IA avanzado utilizando el modelo Mixtral-8x7B-Instruct-v0.1 para conversaciones coherentes y naturales.",
 )
 
 iface.launch()