import streamlit as st
from transformers import pipeline

# Load the sentiment analysis model pipeline
sentiment_classifier = pipeline("text-classification",model='Ryleeeee/CustomSentimentModel', return_all_scores=True)

# Load the product category classification model pipeline
product_categorizer = pipeline("text-classification", model="Ryleeeee/CustomProductCategoryModel")

# Streamlit application title and background image
st.image("./header.png", use_column_width=True)
st.markdown("<h1 style='text-align: center;'>Customer Review Analysis</h1>", unsafe_allow_html=True)

st.write("Sentiment classification: positive, netural, negative")
st.write("Product category classification: books, mobile, mobile accessories, refrigerator, smartTv")

product_dic = {0: "books", 1: "mobile", 2: "mobile accessories", 3: "refrigerator", 4: "smartTv"}

# User can enter the customer review
review = st.text_area("Enter the customer review", "")

def sentiment_class(text):
    results = sentiment_classifier(text)[0]
    max_score = float('-inf')
    max_label = ''
    for result in results:
        if result['score'] > max_score:
            max_score = result['score']
            max_label = result['label']
    return max_score, max_label

def product_category(text):
    results = product_categorizer(text)[0]
    return results


# Perform sentiment analysis when the user clicks the "Classify Sentiment" button
if st.button("Classify Sentiment"):
    # Check if the user has entered review
    if review is None or review.strip() == '':
        st.warning("Please enter a customer review first.")
    else:
        # Perform sentiment analysis on the input text
        sentiment_result = sentiment_class(review)
        st.write("Review sentiment: ", sentiment_result[1])
        st.write("Prediction score: ", sentiment_result[0])

        # Perform text summarization when the review sentiment is classified as negative           
        if sentiment_result[1] == 'negative':
            category = product_dic[int(product_category(review)["label"].split("_")[1])]
            predic_score = product_category(review)["score"]
            st.write("Category of the faulty product: ", category)
            st.write("Prediction score: ", predic_score)