import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig
from huggingface_hub import login

# Set your HuggingFace token
hf_token = st.secrets["HF_TOKEN25"]  # Using Streamlit secrets
try:
    login(token=hf_token)
    st.success("Successfully logged in to Hugging Face!")
except Exception as e:
    st.error(f"Error logging in to Hugging Face: {str(e)}")

st.title("LLaMA Chatbot")

@st.cache_resource
def load_model():
    try:
        model_path = "Alaaeldin/llama2-app"
        st.info("Loading model... This might take a minute.")
        
        tokenizer = AutoTokenizer.from_pretrained(model_path, token=hf_token)
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16,
            device_map="auto",
            load_in_8bit=True,
            token=hf_token
        )
        st.success("✅ Model loaded successfully!")
        return model, tokenizer
    except Exception as e:
        st.error(f"❌ Error loading model: {str(e)}")
        return None, None

model, tokenizer = load_model()

# Add a text input
if model and tokenizer:
    user_input = st.text_input("Your message:", "")
    if st.button("Send"):
        st.write("User:", user_input)