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é correctamente instalado y disponible

# 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 de generación desde HuggingFace
model_name = "dmis-lab/selfbiorag_7b"

try:
    print("Cargando el tokenizador y el modelo de generación desde HuggingFace...")
    tokenizer_gen = AutoTokenizer.from_pretrained(model_name)
    model_gen = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    ).to(device)
except ValueError as e:
    print(f"Error al cargar el tokenizador de generación: {e}")
    sys.exit(1)
except Exception as e:
    print(f"Error al cargar el modelo de generación: {e}")
    sys.exit(1)

# Definir modelos de traducción
# Diccionario de modelos de traducción según el idioma seleccionado
translation_models = {
    "Español": "Helsinki-NLP/opus-mt-en-es",
    "Portugués": "Helsinki-NLP/opus-mt-en-pt",
    "Francés": "Helsinki-NLP/opus-mt-en-fr",
    "Alemán": "Helsinki-NLP/opus-mt-en-de",
    # Añade más idiomas y sus respectivos modelos según sea necesario
}

# Cargar los tokenizadores y modelos de traducción
tokenizer_trans = {}
model_trans = {}

for lang, model_name_trans in translation_models.items():
    try:
        print(f"Cargando el tokenizador y el modelo de traducción para {lang} desde HuggingFace...")
        tokenizer = MarianTokenizer.from_pretrained(model_name_trans)
        model = MarianMTModel.from_pretrained(model_name_trans).to(device)
        tokenizer_trans[lang] = tokenizer
        model_trans[lang] = model
    except Exception as e:
        print(f"Error al cargar el modelo de traducción para {lang}: {e}")
        sys.exit(1)

@spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
@medir_tiempo
def generar_y_traducir_respuesta(consulta, idioma_destino):
    """
    Función que genera una respuesta a partir de una consulta dada y la traduce al idioma seleccionado.
    """
    try:
        if not consulta.strip():
            return "Por favor, ingresa una consulta válida.", ""
        
        # Tokenizar la consulta
        inputs = tokenizer_gen.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       # Generación determinista
            # Puedes añadir otros parámetros como 'num_beams' si lo deseas
        }
        
        # Generar la respuesta
        with torch.no_grad():
            outputs = model_gen.generate(input_ids=inputs, **generation_kwargs)
        
        # Decodificar la respuesta en inglés
        respuesta_en = tokenizer_gen.decode(outputs[0], skip_special_tokens=True)
        
        # Traducir la respuesta al idioma seleccionado
        if idioma_destino in translation_models:
            tokenizer_tr = tokenizer_trans[idioma_destino]
            model_tr = model_trans[idioma_destino]
            
            # Preparar la entrada para la traducción
            traducir_inputs = tokenizer_tr.prepare_seq2seq_batch([respuesta_en], return_tensors="pt").to(device)
            
            # Realizar la traducción
            with torch.no_grad():
                traduccion_outputs = model_tr.generate(**traducir_inputs)
            
            # Decodificar la traducción
            respuesta_traducida = tokenizer_tr.decode(traduccion_outputs[0], skip_special_tokens=True)
        else:
            respuesta_traducida = "Idioma de destino no soportado."
        
        return respuesta_en, respuesta_traducida
    except Exception as e:
        print(f"Error durante la generación o traducción de la respuesta: {e}")
        return f"Error al generar la respuesta: {e}", ""

# Definir la interfaz de Gradio
titulo = "Generador y Traductor de Respuestas con SelfBioRAG 7B"
descripcion = (
    "Ingresa una consulta y el modelo generará una respuesta en inglés. "
    "Luego, puedes seleccionar un idioma para traducir la respuesta generada."
)

idiomas_disponibles = list(translation_models.keys())

iface = gr.Interface(
    fn=generar_y_traducir_respuesta,
    inputs=[
        gr.Textbox(lines=5, placeholder="Escribe tu consulta aquí...", label="Consulta"),
        gr.Dropdown(choices=idiomas_disponibles, value="Español", label="Idioma de Traducción")
    ],
    outputs=[
        gr.Textbox(label="Respuesta en Inglés"),
        gr.Textbox(label="Respuesta Traducida")
    ],
    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.",
            "Español"
        ],
        [
            "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.",
            "Portugués"
        ]
    ],
    cache_examples=False
)

# Ejecutar la interfaz
if __name__ == "__main__":
    iface.launch()