import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Title and Description
st.title("Simple DistilBERT Chatbot")
st.write("This is a basic chatbot prototype. Ask it something!")

# Load Model and Tokenizer
@st.cache_resource  # Cache for efficiency
def load_model_tokenizer():
    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
    model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
    return tokenizer, model

tokenizer, model = load_model_tokenizer()

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

# Generate Response on Button Click
if st.button("Send"):
    if not user_input:
        st.warning("Please enter some text.")
    else:
        # Preprocess and Generate Response (placeholder)
        encoded_input = preprocess_input(user_input)  
        outputs = model(**encoded_input)  

        # (TODO) Extract relevant info from outputs 

        bot_response = "I'm still under development, but I understand you said: {}".format(user_input)
        st.write("Bot: " + bot_response)