File size: 2,468 Bytes
f4af6aa 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c d92f799 09bd12c f4af6aa d92f799 f4af6aa d92f799 09bd12c d92f799 09bd12c f4af6aa 09bd12c f4af6aa 09bd12c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
from langdetect import detect
# --- 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"
# --- Cargar modelos ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
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: {msg_es}\nIA:"
# Generar respuesta en español
output = generator(
context,
max_new_tokens=150,
do_sample=True,
top_k=50,
top_p=0.9,
temperature=0.8
)[0]["generated_text"]
if "IA:" in output:
response_es = output.split("IA:")[-1].strip()
else:
response_es = output
# Traducir de vuelta al idioma original
response_final = translate_from_es(response_es, lang)
# Actualizar historial
history.append((message, response_final))
return history, ""
# --- Interfaz Gradio ---
with gr.Blocks() as demo:
gr.Markdown("# 🤖 Chatbot Multilenguaje (Traducción automática)")
chat = gr.Chatbot()
msg = gr.Textbox(placeholder="Escribe tu mensaje…")
clear_btn = gr.Button("Limpiar")
state = gr.State([])
msg.submit(answer, [state, msg], [chat, msg])
clear_btn.click(lambda: [], None, chat)
demo.launch()
|