Update app.py

app.py

@@ -1,31 +1,53 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+from langdetect import detect
 
-MODEL_NAME = "distilgpt2"  # Modelo ligero para CPU Basic
+# --- CONFIG ---
+MODEL_NAME = "distilgpt2"  # Modelo ligero para CPU
+TRANSLATE_TO_ES_MODEL = "Helsinki-NLP/opus-mt-mul-es"
+TRANSLATE_FROM_ES_MODEL = "Helsinki-NLP/opus-mt-es-mul"
 
-def load_model(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    gen = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=-1  # CPU
-    )
-    return gen
+# --- Cargar modelos ---
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-generator = load_model(MODEL_NAME)
+translator_to_es = pipeline("translation", model=TRANSLATE_TO_ES_MODEL, device=-1)
+translator_from_es = pipeline("translation", model=TRANSLATE_FROM_ES_MODEL, device=-1)
 
+# --- Funciones de traducción ---
+def translate_to_es(text):
+    try:
+        lang = detect(text)
+    except:
+        lang = "es"
+    if lang != "es":
+        translated = translator_to_es(text)[0]["translation_text"]
+        return translated, lang
+    return text, lang
+
+def translate_from_es(text, lang):
+    if lang != "es":
+        translated = translator_from_es(text)[0]["translation_text"]
+        return translated
+    return text
+
+# --- Función principal del chatbot ---
 def answer(history, message):
     if not message.strip():
         return history, ""
 
+    # Detectar idioma y traducir a español si es necesario
+    msg_es, lang = translate_to_es(message)
+
+    # Construir contexto
     context = ""
     for user, bot in history[-6:]:
         context += f"Usuario: {user}\nIA: {bot}\n"
-    context += f"Usuario: {message}\nIA:"
+    context += f"Usuario: {msg_es}\nIA:"
 
+    # Generar respuesta en español
     output = generator(
         context,
         max_new_tokens=150,
@@ -36,17 +58,22 @@ def answer(history, message):
     )[0]["generated_text"]
 
     if "IA:" in output:
-        response = output.split("IA:")[-1].strip()
+        response_es = output.split("IA:")[-1].strip()
     else:
-        response = output
+        response_es = output
+
+    # Traducir de vuelta al idioma original
+    response_final = translate_from_es(response_es, lang)
 
-    history.append((message, response))
+    # Actualizar historial
+    history.append((message, response_final))
     return history, ""
 
+# --- Interfaz Gradio ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Chatbot gratuito hecho por ti")
+    gr.Markdown("# 🤖 Chatbot Multilenguaje (Traducción automática)")
     chat = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Escribe un mensaje…")
+    msg = gr.Textbox(placeholder="Escribe tu mensaje…")
     clear_btn = gr.Button("Limpiar")
     state = gr.State([])