BioRAG / app.py
C2MV's picture
Update app.py
50560b1 verified
raw
history blame
5.47 kB
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
import time
from functools import wraps
import sys
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):
@wraps(func)
def wrapper(*args, **kwargs):
inicio = time.time()
resultado = func(*args, **kwargs)
fin = time.time()
tiempo_transcurrido = fin - inicio
print(f"Tiempo de ejecución de '{func.__name__}': {tiempo_transcurrido:.2f} segundos")
return resultado
return wrapper
# Verificar si CUDA está disponible
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.")
# Cargar el tokenizador y el modelo desde HuggingFace
model_name = "dmis-lab/selfbiorag_7b"
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)
# Cargar el modelo de traducción
try:
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: {e}")
sys.exit(1)
@spaces.GPU(duration=120) # Decorador para asignar GPU durante 120 segundos
@medir_tiempo
def generar_respuesta(consulta):
"""
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."
# 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
"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.generate(input_ids=inputs, **generation_kwargs)
# 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."
# Tokenizar el texto a traducir
translated = translator_model.generate(**translator_tokenizer(texto, return_tensors="pt", padding=True).to(device))
# 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 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 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=procesar_consulta,
inputs=[
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 Original (Inglés)"),
gr.Textbox(label="Traducción al Español")
],
title=titulo,
description=descripcion,
examples=[
[
"Clasifica el siguiente informe de radiología según la parte del cuerpo a la que se refiere (por ejemplo, pecho, abdomen, cerebro, etc.): Los discos intervertebrales en L4-L5 y L5-S1 muestran signos de degeneración con leve abultamiento que comprime la raíz nerviosa adyacente."
],
[
"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,
allow_flagging="never"
)
# Ejecutar la interfaz
if __name__ == "__main__":
iface.launch()