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 para el modelo principal
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 principal 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 en CPU
try:
    print("Cargando el tokenizador y el modelo de traducción en CPU...")
    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("cpu")  # Forzar a CPU
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       # 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.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))
        
        # 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()