stress_model / app.py
The-Adnan-Syed's picture
Create app.py
3097f15 verified
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}%")