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

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

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

# Load model with quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    trust_remote_code=True
)
model.config.use_cache = False

# 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")
    
    # 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}")