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)

def generar_texto(entrada):
    # Tokenizar la entrada
    input_ids = tokenizer.encode(entrada, return_tensors="pt")
    
    # Generar texto con el modelo
    output = model.generate(input_ids, max_length=100, num_return_sequences=1)
    
    # Decodificar y retornar el texto generado
    texto_generado = tokenizer.decode(output[0], skip_special_tokens=True)
    return texto_generado

# Crear la interfaz de Gradio
interfaz = gr.Interface(
    fn=generar_texto,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Escribe tu prompt aquí..."),
    outputs="text",
    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()