Update app.py

app.py

@@ -4,14 +4,22 @@ from transformers import pipeline
 # Load the sentiment analysis model pipeline
 sentiment_classifier = pipeline("text-classification",model='Ryleeeee/CustomSentimentModel', return_all_scores=True)
 
-# Load the text summarization model pipeline
-product_categorizer = pipeline("zero-shot-classification", model="MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33")
+# 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("Setiment classification: positive, netural, negative")
+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"}
+predictions3 = [output_dic[int(o['label'].split("_")[1])] for o in output3]
+
+
+
 
 # User can enter the customer review
 review = st.text_area("Enter the customer review", "")
@@ -44,5 +52,5 @@ if st.button("Classify Sentiment"):
 
         # Perform text summarization when the review sentiment is classified as negative           
         if sentiment_result[1] == 'negative':
-            category = product_category(review)
-            st.write("Product category: ", category)
+            category = output_dic[int(product_category(review)["label"].split("_")[1])]
+            st.write(f"Category of the faulty product is {category}")