Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
 from functools import wraps
+import sys
 
 # Decorador para medir el tiempo de ejecución
 def medir_tiempo(func):
@@ -23,37 +24,50 @@ if device == "cpu":
 
 # Cargar el tokenizador y el modelo desde HuggingFace
 model_name = "dmis-lab/selfbiorag_7b"
-print("Cargando el tokenizador y el modelo desde HuggingFace...")
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-).to(device)
+
+try:
+    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: {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: {e}")
+    sys.exit(1)
 
 @medir_tiempo
 def generar_respuesta(consulta):
     """
     Función que genera una respuesta a partir de una consulta dada.
     """
-    # Tokenizar la consulta
-    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
-        "temperature": 0.0,
-        "top_p": 1.0,
-        "do_sample": False,
-        "skip_special_tokens": True
-    }
-    
-    # Generar la respuesta
-    with torch.no_grad():
-        outputs = model.generate(**inputs, **generation_kwargs)
-    
-    # Decodificar la respuesta
-    respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return respuesta
+    try:
+        # Tokenizar la consulta
+        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
+            "temperature": 0.0,
+            "top_p": 1.0,
+            "do_sample": False,
+            "skip_special_tokens": True
+        }
+        
+        # Generar la respuesta
+        with torch.no_grad():
+            outputs = model.generate(**inputs, **generation_kwargs)
+        
+        # Decodificar la respuesta
+        respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return respuesta
+    except Exception as e:
+        return f"Error al generar la respuesta: {e}"
 
 # Definir la interfaz de Gradio
 titulo = "Generador de Respuestas con SelfBioRAG 7B"
@@ -61,8 +75,8 @@ descripcion = "Ingresa una consulta y el modelo generará una respuesta basada e
 
 iface = gr.Interface(
     fn=generar_respuesta,
-    inputs=gr.inputs.Textbox(lines=5, placeholder="Escribe tu consulta aquí..."),
-    outputs=gr.outputs.Textbox(),
+    inputs=gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí..."),
+    outputs=gr.Textbox(),
     title=titulo,
     description=descripcion,
     examples=[