import streamlit as st
from transformers import pipeline

def main():
    sentiment_pipeline = pipeline(model="frankai98/Yelp_Bert")

    st.title("Sentiment Analysis with HuggingFace Spaces")
    st.write("Enter a sentence to analyze its sentiment:")

    user_input = st.text_input("")
    if user_input:
        result = sentiment_pipeline(user_input)
        sentiment_label = result[0]["label"]
        confidence = result[0]["score"]
        
        # Extract the number from the label (assuming format is "LABEL_X")
        if sentiment_label.startswith("LABEL_"):
            try:
                # Extract the number after "LABEL_" and convert to int
                label_num = int(sentiment_label.split("_")[1])
                # Add 1 to display as stars
                stars = label_num + 1
                st.write(f"Sentiment: {stars} stars")
            except (ValueError, IndexError):
                # Fallback if the format is unexpected
                st.write(f"Sentiment: {sentiment_label}")
        else:
            # Direct display if label doesn't match expected format
            st.write(f"Sentiment: {sentiment_label}")
            
        st.write(f"Confidence: {confidence:.2f}")

if __name__ == "__main__":
    main()