File size: 3,656 Bytes
5a798cc 820a0dd c3a5bd0 b956a25 820a0dd 5a798cc b956a25 820a0dd 5a798cc c3a5bd0 3d55aee c3a5bd0 407e335 c3a5bd0 b956a25 c3a5bd0 3d55aee c3a5bd0 5a798cc c3a5bd0 5a798cc |
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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
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)
@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}"
# Definir la interfaz de Gradio
titulo = "Generador de Respuestas con SelfBioRAG 7B"
descripcion = "Ingresa una consulta y el modelo generará una respuesta basada en el contenido proporcionado."
iface = gr.Interface(
fn=generar_respuesta,
inputs=gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí..."),
outputs=gr.Textbox(),
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
)
# Ejecutar la interfaz
if __name__ == "__main__":
iface.launch()
|