import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")

# Set the pad_token_id if not already set
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Streamlit app setup
st.title("Friendly Bot")
st.write("Type a message and press Enter to chat with the bot.")

# Initialize conversation history in session state
if "history" not in st.session_state:
    st.session_state["history"] = []

# Function to get the bot response
def get_bot_response(user_input):
    # Tokenize input and create an attention mask
    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
    attention_mask = inputs['attention_mask']
    
    # Generate response using the attention mask and pad_token_id
    outputs = model.generate(
        inputs['input_ids'], 
        attention_mask=attention_mask, 
        max_new_tokens=150, 
        pad_token_id=tokenizer.pad_token_id
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    return response

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

# Chat interface
if st.button("Send"):
    if user_input:
        # Add user message to history
        st.session_state.history.append(("User", user_input))
        
        # Get bot response and add it to history
        bot_response = get_bot_response(user_input)
        st.session_state.history.append(("Bot", bot_response))

# Display conversation history
for speaker, message in st.session_state.history:
    if speaker == "User":
        st.write(f"You: {message}")
    else:
        st.write(f"Bot: {message}")