import torch
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer

# Model name
model_name = "ybelkada/falcon-7b-sharded-bf16"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

# Load model in CPU mode
device = "cpu"  # Hugging Face Spaces does not provide free GPUs
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # Use float16 for lower memory usage
    device_map=device
)

# Streamlit UI
st.title("🦜 Falcon-7B Chatbot")
st.write("Ask me anything!")

# Store chat history
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# User input
user_input = st.text_input("You:", "")

if user_input:
    # Tokenize input
    inputs = tokenizer(user_input, return_tensors="pt")
    inputs.pop("token_type_ids", None)  # Remove token_type_ids to avoid errors
    inputs = {key: value.to(device) for key, value in inputs.items()}  # Move inputs to device
    
    # Generate response
    with torch.no_grad():
        output = model.generate(**inputs, max_length=200, do_sample=True, top_k=50, top_p=0.95)
    
    # Decode response
    response = tokenizer.decode(output[:, inputs["input_ids"].shape[-1]:][0], skip_special_tokens=True)
    
    # Store and display chat history
    st.session_state.chat_history.append(("You", user_input))
    st.session_state.chat_history.append(("Bot", response))

# Display chat history
for sender, message in st.session_state.chat_history:
    st.write(f"**{sender}:** {message}")