BioRAG / app.py
C2MV's picture
Update app.py
ce0f331 verified
raw
history blame
4.48 kB
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
from functools import wraps
# Asegúrate de que el módulo 'spaces' esté disponible y correctamente instalado.
# Si 'spaces.GPU' no está disponible, puedes omitir o implementar una alternativa.
# En este ejemplo, se asume que 'spaces.GPU' es un decorador válido en tu entorno.
try:
import spaces
except ImportError:
# Si el módulo 'spaces' no está disponible, define un decorador vacío
# para evitar errores. Esto es útil si no estás utilizando características
# específicas de Hugging Face Spaces que requieran este decorador.
def GPU(duration):
def decorator(func):
return func
return decorator
spaces = type('spaces', (), {'GPU': GPU})
# 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
# Configurar el dispositivo (GPU si 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 desde Hugging Face Hub
model_name = "pfnet/GenerRNA"
try:
print("Cargando el tokenizador desde Hugging Face Hub...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
except ValueError as e:
print(f"Error al cargar el tokenizador: {e}")
sys.exit(1)
# Cargar el modelo desde Hugging Face Hub
try:
print("Cargando el modelo GenerRNA desde Hugging Face Hub...")
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
revision="main" # Asegúrate de que la rama correcta esté especificada
).to(device)
model.eval()
print("Modelo GenerRNA cargado exitosamente.")
except Exception as e:
print(f"Error al cargar el modelo GenerRNA: {e}")
sys.exit(1)
@spaces.GPU(duration=120) # Decorador para asignar GPU durante 120 segundos
@medir_tiempo
def generar_rna_sequence(prompt, max_length=256):
"""
Función que genera una secuencia de RNA a partir de una secuencia inicial dada.
"""
try:
if not prompt.strip():
return "Por favor, ingresa una secuencia de inicio válida."
# Tokenizar la entrada
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
# Generar la secuencia
with torch.no_grad():
outputs = model.generate(
inputs,
max_length=max_length,
num_return_sequences=1,
no_repeat_ngram_size=2,
temperature=0.7,
top_k=50,
top_p=0.95,
do_sample=True
)
# Decodificar la secuencia generada
generated_sequence = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_sequence
except Exception as e:
print(f"Error durante la generación de secuencia: {e}")
return f"Error al generar la secuencia: {e}"
# Definir la interfaz de Gradio
titulo = "GenerRNA - Generador de Secuencias de RNA"
descripcion = (
"GenerRNA es un modelo generativo de RNA basado en una arquitectura Transformer. "
"Ingresa una secuencia inicial opcional y define la longitud máxima para generar nuevas secuencias de RNA."
)
iface = gr.Interface(
fn=generar_rna_sequence,
inputs=[
gr.Textbox(
lines=5,
placeholder="Ingresa una secuencia de RNA inicial (opcional)...",
label="Secuencia Inicial"
),
gr.Slider(
minimum=50,
maximum=1000,
step=50,
value=256,
label="Longitud Máxima de la Secuencia"
)
],
outputs=gr.Textbox(label="Secuencia de RNA Generada"),
title=titulo,
description=descripcion,
examples=[
[
"AUGGCUACGUAUCGACGUA"
],
[
"GCUAUGCUAGCUAGCUGAC"
]
],
cache_examples=False,
allow_flagging="never"
)
# Ejecutar la interfaz
if __name__ == "__main__":
iface.launch()