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()