Update app.py

app.py

@@ -1,15 +1,10 @@
 import gradio as gr
 import torch
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    MarianMTModel,
-    MarianTokenizer
-)
+from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
 import time
 from functools import wraps
 import sys
-import os
+import spaces  # Asegúrate de que este módulo esté disponible y correctamente instalado
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -28,97 +23,111 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 if device == "cpu":
     print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
 
-# Obtener el token de Hugging Face desde las variables de entorno
-hf_token = os.getenv("HUGGINGFACE_TOKEN")
-if not hf_token:
-    print("Error: El token de Hugging Face no está configurado en los secretos.")
-    sys.exit(1)
-
-# Cargar el tokenizador y el modelo de generación desde HuggingFace
+# Cargar el tokenizador y el modelo desde HuggingFace
 model_name = "dmis-lab/selfbiorag_7b"
 
 try:
-    print("Cargando el tokenizador y el modelo de generación desde HuggingFace...")
-    tokenizer_gen = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
-    model_gen = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        use_auth_token=hf_token
-    ).to(device)
+    print("Cargando el tokenizador y el modelo desde HuggingFace...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 except ValueError as e:
-    print(f"Error al cargar el tokenizador de generación: {e}")
+    print(f"Error al cargar el tokenizador: {e}")
     sys.exit(1)
+
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    ).to(device)
 except Exception as e:
-    print(f"Error al cargar el modelo de generación: {e}")
+    print(f"Error al cargar el modelo: {e}")
     sys.exit(1)
 
-# Definir el modelo de traducción al Español
-translation_model_name = "Helsinki-NLP/opus-mt-en-es"
-
+# Cargar el modelo de traducción
 try:
-    print(f"Cargando el tokenizador y el modelo de traducción para Español desde HuggingFace...")
-    tokenizer_tr_es = MarianTokenizer.from_pretrained(translation_model_name)
-    model_tr_es = MarianMTModel.from_pretrained(translation_model_name).to(device)
+    print("Cargando el tokenizador y el modelo de traducción...")
+    translation_model_name = "Helsinki-NLP/opus-mt-en-es"
+    translator_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
+    translator_model = MarianMTModel.from_pretrained(translation_model_name).to(device)
 except Exception as e:
-    print(f"Error al cargar el modelo de traducción para Español: {e}")
+    print(f"Error al cargar el modelo de traducción: {e}")
     sys.exit(1)
 
+@spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
 @medir_tiempo
-def generar_y_traducir_respuesta(consulta, idioma_destino):
+def generar_respuesta(consulta):
     """
-    Función que genera una respuesta a partir de una consulta dada y la traduce al Español.
+    Función que genera una respuesta a partir de una consulta dada.
     """
     try:
         if not consulta.strip():
-            return "Por favor, ingresa una consulta válida.", ""
+            return "Por favor, ingresa una consulta válida."
         
         # Tokenizar la consulta
-        inputs = tokenizer_gen.encode(consulta, return_tensors="pt").to(device)
+        inputs = tokenizer.encode(consulta, return_tensors="pt").to(device)
         
         # Configurar los parámetros de generación
         generation_kwargs = {
             "max_new_tokens": 100,  # Ajustado a 100
-            "do_sample": False       # Generación determinista
-            # Puedes añadir otros parámetros como 'num_beams' si lo deseas
+            "do_sample": False       # Configura según la opción deseada
+            # "temperature": 0.0,    # Comenta o elimina si do_sample=False
+            # "top_p": 1.0            # Comenta o elimina si do_sample=False
         }
         
         # Generar la respuesta
         with torch.no_grad():
-            outputs = model_gen.generate(input_ids=inputs, **generation_kwargs)
-        
-        # Decodificar la respuesta en inglés con limpieza de espacios
-        respuesta_en = tokenizer_gen.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        
-        # Traducir la respuesta al Español
-        traducir_inputs = tokenizer_tr_es.encode(respuesta_en, return_tensors="pt").to(device)
+            outputs = model.generate(input_ids=inputs, **generation_kwargs)
         
-        # Realizar la traducción
-        with torch.no_grad():
-            traduccion_outputs = model_tr_es.generate(input_ids=traducir_inputs, max_length=512)
+        # Decodificar la respuesta
+        respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return respuesta
+    except Exception as e:
+        print(f"Error durante la generación de respuesta: {e}")
+        return f"Error al generar la respuesta: {e}"
+
+def traducir_texto(texto):
+    """
+    Función que traduce un texto de inglés a español.
+    """
+    try:
+        if not texto.strip():
+            return "No hay texto para traducir."
         
-        # Decodificar la traducción con limpieza de espacios
-        respuesta_traducida = tokenizer_tr_es.decode(traduccion_outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        # Tokenizar el texto a traducir
+        translated = translator_model.generate(**translator_tokenizer(texto, return_tensors="pt", padding=True).to(device))
         
-        return respuesta_en, respuesta_traducida
+        # Decodificar la traducción
+        traduccion = translator_tokenizer.decode(translated[0], skip_special_tokens=True)
+        return traduccion
     except Exception as e:
-        print(f"Error durante la generación o traducción de la respuesta: {e}")
-        return f"Error al generar la respuesta: {e}", ""
+        print(f"Error durante la traducción: {e}")
+        return f"Error al traducir el texto: {e}"
+
+def procesar_consulta(consulta, idioma):
+    """
+    Función que procesa la consulta y devuelve la respuesta original y/o traducida según el idioma seleccionado.
+    """
+    respuesta_original = generar_respuesta(consulta)
+    
+    if idioma == "Español":
+        traduccion = traducir_texto(respuesta_original)
+    else:
+        traduccion = ""
+    
+    return respuesta_original, traduccion
 
 # Definir la interfaz de Gradio
-titulo = "Generador y Traductor de Respuestas con SelfBioRAG 7B"
-descripcion = (
-    "Ingresa una consulta y el modelo generará una respuesta en inglés. "
-    "Luego, la respuesta se traducirá automáticamente al Español."
-)
+titulo = "Generador de Respuestas con SelfBioRAG 7B"
+descripcion = "Ingresa una consulta y selecciona el idioma de salida. El modelo generará una respuesta basada en el contenido proporcionado."
 
 iface = gr.Interface(
-    fn=generar_y_traducir_respuesta,
+    fn=procesar_consulta,
     inputs=[
-        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta")
+        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta"),
+        gr.Dropdown(choices=["Inglés", "Español"], value="Inglés", label="Idioma de Salida")
     ],
     outputs=[
-        gr.Textbox(label="Respuesta en Inglés"),
-        gr.Textbox(label="Respuesta Traducida al Español")
+        gr.Textbox(label="Respuesta Original (Inglés)"),
+        gr.Textbox(label="Traducción al Español")
     ],
     title=titulo,
     description=descripcion,
@@ -130,7 +139,8 @@ iface = gr.Interface(
             "Resume los puntos clave sobre el papel de las mutaciones en los genes BRCA1 y BRCA2 en el aumento del riesgo de cáncer de mama."
         ]
     ],
-    cache_examples=False
+    cache_examples=False,
+    allow_flagging="never"
 )
 
 # Ejecutar la interfaz