import streamlit as st
from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification
import tensorflow as tf

# Load the pre-trained model and tokenizer using the correct Hugging Face model repo ID
model_path = 'shukdevdatta123/Dreaddit_DistillBert_Stress_Model'
loaded_model = TFDistilBertForSequenceClassification.from_pretrained(model_path)
loaded_tokenizer = DistilBertTokenizer.from_pretrained(model_path)

# Define the prediction function
def predict_with_loaded_model(in_sentences):
    labels = ["non-stress", "stress"]
    inputs = loaded_tokenizer(in_sentences, return_tensors="tf", padding=True, truncation=True, max_length=512)
    predictions = loaded_model(inputs)
    predicted_labels = tf.argmax(predictions.logits, axis=-1).numpy()
    predicted_probs = tf.nn.softmax(predictions.logits, axis=-1).numpy()

    return [{"text": sentence, "confidence": probs.tolist(), "label": labels[label]} for sentence, label, probs in zip(in_sentences, predicted_labels, predicted_probs)]

# Streamlit interface
st.title("Stress Prediction with DistilBERT")

# Initialize session state variables for user input and prediction output
if 'user_input' not in st.session_state:
    st.session_state.user_input = ""
if 'prediction' not in st.session_state:
    st.session_state.prediction = None

# Add a text input box for the user to enter a sentence
user_input = st.text_area("Enter a sentence or text:", st.session_state.user_input)

# Define buttons with a sidebar layout for easy spacing
col1, col2 = st.columns([1, 1])
with col1:
    # When the user clicks "Predict", run the prediction function
    if st.button("Predict"):
        if user_input:
            # Update session state with the user input
            st.session_state.user_input = user_input
            # Make the prediction using the model
            st.session_state.prediction = predict_with_loaded_model([user_input])[0]
        else:
            st.write("Please enter a sentence to predict.")

with col2:
    # Refresh button to reset the user input and output
    if st.button("Refresh"):
        st.session_state.user_input = ""
        st.session_state.prediction = None
        st.experimental_rerun()

# Display the prediction result
if st.session_state.prediction:
    prediction = st.session_state.prediction
    st.write(f"Text: {prediction['text']}")
    st.write(f"Prediction: {prediction['label']}")
    # st.write(f"Confidence: {prediction['confidence']}")