import streamlit as st
from transformers import pipeline

# Set up chatbot pipeline (you can change the model)
chatbot = pipeline("text-generation", model="microsoft/DialoGPT-medium", cache_dir="/app/cache")

st.title("Simple Chatbot")
st.write("Chat with an AI below:")

# Session state to store the chat history
if "history" not in st.session_state:
    st.session_state.history = []

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

if st.button("Send"):
    if user_input:
        st.session_state.history.append(f"You: {user_input}")

        # Get response from model
        response = chatbot(user_input, max_length=100, pad_token_id=50256)[0]["generated_text"]
        reply = response.split(user_input)[-1].strip()
        st.session_state.history.append(f"Bot: {reply}")
    else:
        st.warning("Please enter a message.")

# Display conversation
for line in st.session_state.history:
    st.write(line)