File size: 1,622 Bytes
1133027 |
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 |
import streamlit as st
from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification
import tensorflow as tf
# Load the pre-trained model and tokenizer
model_path = 'drive-download-20241117T174204Z-001/'
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.")
|