Spaces:
Sleeping
Sleeping
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}%") | |