import streamlit as st
import os
import logging
from transformers import AutoModelForCausalLM, AutoTokenizer

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

st.title("Meta LLaMA Text Generation")

@st.cache_resource
def load_model():
    model_name = "meta-llama/Meta-Llama-3-8B"
    access_token = os.getenv('hf')
    
    if not access_token:
        st.error("Hugging Face access token is not set. Please set the environment variable 'hf'.")
        return None, None
    
    logger.info("Loading tokenizer and model...")
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token)
    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=access_token)
    return tokenizer, model

tokenizer, model = load_model()

if tokenizer is not None and model is not None:
    prompt = st.text_input("Enter a prompt:", "Once upon a time")

    if st.button("Generate Text"):
        try:
            inputs = tokenizer(prompt, return_tensors="pt")
            outputs = model.generate(**inputs, max_length=50)
            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            st.write(generated_text)
        except Exception as e:
            st.error(f"An error occurred: {e}")
            logger.error(f"An error occurred during text generation: {e}")
else:
    st.error("Failed to load the model. Check the logs for more details.")