import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Especifica el nombre del modelo
model_name = "BSC-LT/ALIA-40b"

# Cargar el tokenizador y el modelo
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16  # Reduce uso de memoria a la mitad
)

# Verificar en qué dispositivo está el modelo
print(f"Modelo cargado en: {model.device}")



def generar_texto(entrada):
    # Tokenizar la entrada
    input_ids = tokenizer.encode(entrada, return_tensors="pt").input_ids.to("cuda")
    
    # Generar texto con el modelo
    output = model.generate(
        input_ids,
        max_length=100,         # Reduce la longitud de salida
        temperature=0.7,        # Controla la aleatoriedad
        top_p=0.9,              # Reduce la búsqueda del modelo
        num_return_sequences=1, # Evita múltiples respuestas simultáneas
        do_sample=True          # Sampling para rapidez
    )

    
    # Decodificar y retornar el texto generado
    texto_generado = tokenizer.decode(output[0], skip_special_tokens=True)
    return texto_generado

# Crear la interfaz de Gradio usando la sintaxis actualizada
interfaz = gr.Interface(
    fn=generar_texto,
    inputs=gr.Textbox(lines=2, placeholder="Escribe tu prompt aquí..."),
    outputs=gr.Textbox(),
    title="Generador de Texto con ALIA-40b",
    description="Este modelo genera texto utilizando ALIA-40b, un modelo LLM entrenado por BSC-LT."
)



if __name__ == "__main__":
    interfaz.launch()