import streamlit as st
import tensorflow as tf
from transformers import BertTokenizer, TFBertForSequenceClassification
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Load the BERT tokenizer and model
tokenizer = BertTokenizer.from_pretrained('config.json')  # Path to folder containing config.json
model = TFBertForSequenceClassification.from_pretrained('tf_model.h5', from_pt=True)  # Path to folder containing tf_model.h5

def predict(text):
    # Encode the text using the BERT tokenizer
    input_ids = tokenizer.encode(text, add_special_tokens=True, max_length=128, truncation=True)
    input_ids = pad_sequences([input_ids], maxlen=128, truncating='post', padding='post')
    
    # Convert to tensors
    input_ids = tf.convert_to_tensor(input_ids)
    
    # Get predictions
    logits = model(input_ids)[0]
    
    # Apply softmax to calculate probabilities
    probabilities = tf.nn.softmax(logits, axis=1).numpy()[0]
    
    return probabilities

# Streamlit UI
st.title("Stress Categorization with BERT")
st.write("Enter the text to analyze for stress levels:")

# Text input
user_input = st.text_area("Text", height=150)

if st.button("Predict"):
    # Make prediction
    probabilities = predict(user_input)
    
    # Display probabilities
    st.write("Probabilities:")
    st.write(f"Stressed: {probabilities[1]:.4f}")
    st.write(f"Not Stressed: {probabilities[0]:.4f}")
    
    # Display the most likely class
    if probabilities[0] > probabilities[1]:
        st.success("Prediction: Not Stressed")
    else:
        st.error("Prediction: Stressed")

    # Assuming you have an accuracy metric available (replace with actual accuracy if available)
    accuracy = 0.95  # Example accuracy
    st.write(f"Model Accuracy: {accuracy * 100:.2f}%")