Khalaji, Ali
Add class probs
bd2f6d2
import streamlit as st
from transformers import pipeline
# Load classification model from Hugging Face
model_name = "ale-dp/distilbert-base-uncased-finetuned-emotion"
text_classifier = pipeline('text-classification', model=model_name)
# Define class labels
class_labels = ["Sadness", "Joy", "Love", "Anger", "Fear", "Surprise"]
def main():
st.title("Ordinal Emotion Classifier")
user_input = st.text_area("Enter text:")
if st.button("Classify"):
if user_input:
results = classify_text(user_input)
display_results(results)
else:
st.warning("Please enter some text to classify.")
def classify_text(text):
results = text_classifier(text, return_all_scores=True)
scores_list = results[0]
total_score = sum(score['score'] for score in scores_list)
labeled_probabilities = {}
for score in scores_list:
label = score['label']
probability = (score['score'] / total_score) * 100
labeled_probabilities[label] = probability
return labeled_probabilities
def display_results(results):
st.subheader("Prediction:")
for label, probability in results.items():
st.write(f"{label.lower()}: {probability:.2f}%")
if __name__ == "__main__":
main()