import streamlit as st
from transformers import pipeline

# Load models

# Distilled Sentiment Classifier
# Link: https://huggingface.co/lxyuan/distilbert-base-multilingual-cased-sentiments-student
distilled_sentiment_classifier = pipeline(
    model="lxyuan/distilbert-base-multilingual-cased-sentiments-student", 
    return_all_scores=True
)

# Emotion Classifier
# Link: https://huggingface.co/SamLowe/roberta-base-go_emotions
emotion_text_classifier = pipeline("text-classification", model="SamLowe/roberta-base-go_emotions")

# Named Entity Recognition
# Link: https://huggingface.co/mdarhri00/named-entity-recognition
named_entity_classifier = pipeline("token-classification", model="mdarhri00/named-entity-recognition")

# Toxicity Classifier
# Link: https://huggingface.co/s-nlp/roberta_toxicity_classifier
toxicity_classifier = pipeline("text-classification", model="s-nlp/roberta_toxicity_classifier")

# Streamlit app
def main():
    st.title("HuggingFace Model Demo App")

    # User input for text
    user_text = st.text_area("Enter some text:")

    if user_text:
        # Available Models
        analysis_type = st.selectbox("Select Analysis Type", ["", "Sentiment Analysis", "Emotion Analysis", "Named Entity Recognition", "Toxicity Analysis"])
        
        # Run custom and display outputs
        st.header("Function Output:")

        if analysis_type == "Sentiment Analysis":
            st.subheader("Sentiment Analysis:")
            
            for labels_and_scores in data:
                for entry in labels_and_scores:
                    label = entry["label"]
                    score = entry["score"]
                    st.write(f"Label: {label}, Score: {score}")
            
            # Find the label with the highest score
            max_score_entry = max(data, key=lambda x: x['score'])
            max_score_label = max_score_entry['label']
            max_score = max_score_entry['score']
            
            # Display the label with the highest score
            st.write(f"Label with the highest score: {max_score_label}, Highest Score: {max_score}")

    
        if analysis_type == "Emotion Analysis":    
            st.subheader("Emotion Analysis:")
            
            # Parse JSON data
            data = emotion_text_classifier(user_text)
            
            # Extract and display label and score for each entry
            for entry in data:
                label = entry["label"]
                score = entry["score"]
                st.write(f"Label: {label}, Score: {score}")

        
        if analysis_type == "Named Entity Recognition":
            st.subheader("Named Entity Recognition:")
            
            # Parse JSON data
            data = named_entity_classifier(user_text)
            #st.write(data)
            # Extract and display data
            for labels in data:
                for entry in labels:
                    entity_group = entry["entity_group"]
                    score = entry["score"]
                    word = entry["word"]
                    start = entry["start"]
                    end = entry["end"]
                    st.write(f"Word: {word}, Entity Group: {entity_group}, Score: {score}, Start: {start}, End: {end}")

        
        if analysis_type == "Toxicity Analysis":
            st.subheader("Toxicity Analysis:")
            
            # Parse JSON data
            data = toxicity_classifier(user_text)
            #st.write(data)
            # Extract and display label and score values
            for entry in data:
                label = entry["label"]
                score = entry["score"]
                st.write(f"Label: {label}, Score: {score}")

if __name__ == "__main__":
    main()