File size: 2,528 Bytes
ed2bd59 59d81fb ed2bd59 982521b ed2bd59 982521b c158890 982521b c158890 982521b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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"):
# Clear both user input and prediction output
st.session_state.user_input = ""
st.session_state.prediction = None
st.rerun() # Refresh the app
# 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']}")
|