Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer, MarianMTModel, MarianTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
 import time
 from functools import wraps
 import sys
@@ -24,7 +24,7 @@ if device == "cpu":
     print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")
 
 # Cargar el tokenizador y el modelo principal desde HuggingFace
-model_name = "yangheng/OmniGenome-52M"
+model_name = "dmis-lab/selfbiorag_7b"
 
 try:
     print("Cargando el tokenizador y el modelo desde HuggingFace...")
@@ -34,7 +34,7 @@ except ValueError as e:
     sys.exit(1)
 
 try:
-    model = AutoModel.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         model_name, 
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     ).to(device)
@@ -65,15 +65,20 @@ def generar_respuesta(consulta):
         # Tokenizar la consulta
         inputs = tokenizer.encode(consulta, return_tensors="pt").to(device)
         
-        # Realizar una pasada hacia adelante (forward pass)
+        # Configurar los parámetros de generación
+        generation_kwargs = {
+            "max_new_tokens": 100,  # Ajustado a 100
+            "do_sample": False       # No usar sampling
+            # "temperature": 0.6,    # Eliminado para evitar advertencias
+            # "top_p": 0.9            # Eliminado para evitar advertencias
+        }
+        
+        # Generar la respuesta
         with torch.no_grad():
-            outputs = model(inputs)
+            outputs = model.generate(input_ids=inputs, **generation_kwargs)
         
-        # Procesar la salida según las capacidades del modelo OmniGenome
-        # Aquí deberás implementar la lógica específica basada en lo que OmniGenome retorna
-        # Por ejemplo, si retorna embeddings, podrías convertirlos a texto de alguna manera
-        # Dado que OmniGenome no está diseñado para generación de texto, este es un placeholder
-        respuesta = "Salida del modelo OmniGenome no está diseñada para generación de texto."
+        # 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}")
@@ -111,7 +116,7 @@ def procesar_consulta(consulta, idioma):
     return respuesta_original, traduccion
 
 # Definir la interfaz de Gradio
-titulo = "Generador de Respuestas con OmniGenome"
+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(
@@ -121,7 +126,7 @@ iface = gr.Interface(
         gr.Dropdown(choices=["Inglés", "Español"], value="Español", label="Idioma de Salida")
     ],
     outputs=[
-        gr.Textbox(label="Respuesta Original"),
+        gr.Textbox(label="Respuesta Original (Inglés)"),
         gr.Textbox(label="Traducción al Español")
     ],
     title=titulo,