|
import streamlit as st |
|
from transformers import pipeline |
|
|
|
|
|
model_name = "ale-dp/distilbert-base-uncased-finetuned-emotion" |
|
text_classifier = pipeline('text-classification', model=model_name) |
|
|
|
|
|
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() |
|
|