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

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

# When the user clicks "Predict", run the prediction function
if st.button("Predict"):
    if user_input:
        # Make the prediction using the model
        prediction = predict_with_loaded_model([user_input])[0]
        st.write(f"Text: {prediction['text']}")
        st.write(f"Prediction: {prediction['label']}")
        # st.write(f"Confidence: {prediction['confidence']}")
    else:
        st.write("Please enter a sentence to predict.")