import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
from huggingface_hub import login

# Iniciar sesión en Hugging Face utilizando un token guardado en una variable de entorno
def login_to_huggingface():
    read_key = os.getenv("Read_Token")  # Asegúrate de tener la variable de entorno configurada
    if read_key:
        login(token=read_key)
    else:
        st.error("No se encontró el token de Hugging Face. Asegúrate de que la variable de entorno 'HUGGINGFACE_TOKEN' esté configurada.")

# Inicializar el modelo generador de texto
@st.cache_resource  # Cachear el modelo para no recargarlo repetidamente
def load_model():
    model_id = "meta-llama/Llama-3.2-1B"  # Cambia al modelo correcto disponible
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    return tokenizer, model

# Iniciar sesión
login_to_huggingface()

# Cargar el modelo y el tokenizador
tokenizer, model = load_model()

# Crear la interfaz de usuario
st.title("Generador de Texto Llama-2")
st.write("Introduce un texto para que el modelo lo continúe:")

# Campo de entrada para el prompt
user_input = st.text_input("Escribe tu prompt aquí:", "")

# Botón para generar texto
if st.button("Generar texto"):
    if user_input:
        # Tokenizar el input y generar el texto
        input_ids = tokenizer.encode(user_input, return_tensors='pt').to(model.device)
        with st.spinner("Generando..."):
            output = model.generate(
                input_ids,
                do_sample=True,
                max_length=100,
                num_return_sequences=1,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decodificar el texto generado
        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        
        # Mostrar el texto generado
        st.subheader("Texto generado:")
        st.write(generated_text)
    else:
        st.error("Por favor, introduce un prompt antes de generar texto.")