import streamlit as st
from transformers import pipeline


label2stars = {
    1: "⭐️",
    2: "⭐️⭐️",
    3: "⭐️⭐️⭐️",
    4: "⭐️⭐️⭐️⭐️",
    5: "⭐️⭐️⭐️⭐️⭐️",
    "1 star": "⭐️",
    "2 stars": "⭐️⭐️",
    "3 stars": "⭐️⭐️⭐️",
    "4 stars": "⭐️⭐️⭐️⭐️",
    "5 stars": "⭐️⭐️⭐️⭐️⭐️",
}

@st.cache_resource
def load_model(model_id):
    return pipeline("sentiment-analysis", model=model_id)


bert_model = load_model("nlptown/bert-base-multilingual-uncased-sentiment")
modernbert_model = load_model("nlptown/ModernBERT-base-sentiment-20241228")

st.title("NLP Town Sentiment Analysis")

example1 = "I have sensitive eyes and this is non-irritating. My eyes water a lot and that often leads to product getting into my eyes and causing stinging and or burning in the eye area. This cream has not bothered me at all."
example2 = "Very natural, light weight and good quality. I use on brides who are not used to lashes but still want a bit of length and definition."
example3 = "I have to admit that I'm very new to meditation and guided imagery, but I found the suggested imagery of a \"guardian\" closing the drapes in my bedroom and then sitting next to my bed while I was sleeping to be unnerving. I was wide awake at the end of the CD and was completely disappointed with this purchase. I'm glad that I am able to return it."
example4 = "Catchy promises, little delivery."

option = st.selectbox("Examples", [example1, example2, example3, example4], index=None, placeholder='Select an example or enter your text below')

query = st.text_area("Enter your text here", value=option)
click = st.button("Analyze", type="primary")

if query or click:

    bert_result = bert_model(query)[0]
    modernbert_result = modernbert_model(query)[0]

    col1, col2 = st.columns([3, 1])

    with col1:
        st.write("[nlptown/bert-base-multilingual-uncased-sentiment](https://huggingface.co/nlptown/bert-base-multilingual-uncased-sentiment):")
        st.write("[nlptown/ModernBERT-base-sentiment](https://www.nlp.town):")

    with col2:
        st.write(label2stars[bert_result['label']])
        st.write(label2stars[modernbert_result['label']])