Update app.py

app.py

@@ -21,73 +21,59 @@ generator = pipeline("text2text-generation", model="google/flan-t5-base")
 
 # Function to classify customer comments
 @spaces.GPU
-def classify_comments(category_boxes):
+def classify_comments(categories):
     sentiments = []
-    categories = []
+    assigned_categories = []
     for comment in df['customer_comment']:
+        # Classify sentiment
         sentiment = classifier(comment)[0]['label']
-        category_list = [box for box in category_boxes if box.strip() != '']
-        category_str = ', '.join([cat.strip() for cat in category_list])
+        # Generate category
+        category_str = ', '.join(categories)
         prompt = f"What category best describes this comment? '{comment}' Please answer using only the name of the category: {category_str}."
         category = generator(prompt, max_length=30)[0]['generated_text']
-        categories.append(category)
+        assigned_categories.append(category)
         sentiments.append(sentiment)
     df['comment_sentiment'] = sentiments
-    df['comment_category'] = categories
+    df['comment_category'] = assigned_categories
     return df[['customer_comment', 'comment_sentiment', 'comment_category']].to_html(index=False)
 
 # Gradio Interface
 with gr.Blocks() as nps:
-    def add_category(category_list, new_category):
-        if new_category.strip() != "":
-            category_list.append(new_category.strip())  # Add new category
-        return category_list
+    # State to store categories
+    categories = gr.State([])
 
-    def remove_category(category, category_list):
-        category_list.remove(category)  # Remove selected category
-        return category_list
+    # Function to add a category
+    def add_category(categories, new_category):
+        if new_category.strip() != "" and len(categories) < 5:  # Limit to 5 categories
+            categories.append(new_category.strip())
+        return categories, f"Categories: {', '.join(categories)}"
 
+    # Function to display categories
     def display_categories(categories):
-        category_components = []
-        for i, cat in enumerate(categories):
-            with gr.Row():
-                gr.Markdown(f"- {cat}")
-                remove_btn = gr.Button("X", elem_id=f"remove_{i}", interactive=True)
-                remove_btn.click(
-                    fn=remove_category,
-                    inputs=[gr.State(cat), category_boxes],
-                    outputs=category_boxes
-                )
-            category_components.append(gr.Row())
-        return category_components
-
-    category_boxes = gr.State([])  # Store category input boxes as state
-    category_column = gr.Column()
+        return gr.Column.update(visible=True, value=[gr.Markdown(f"- {cat}") for cat in categories])
 
+    # UI for adding categories
     with gr.Row():
         category_input = gr.Textbox(label="New Category", placeholder="Enter category name")
         add_category_btn = gr.Button("Add Category")
-        add_category_btn.click(
-            fn=add_category,
-            inputs=[category_boxes, category_input],
-            outputs=category_boxes
-        )
-        category_boxes.change(
-            fn=display_categories,
-            inputs=category_boxes,
-            outputs=category_column
-        )
+        category_status = gr.Markdown("Categories: None")
 
+    # Display added categories
+    category_display = gr.Column(visible=False)
+
+    # File upload and template buttons
     uploaded_file = gr.File(label="Upload CSV", type="filepath")
     template_btn = gr.Button("Use Template")
     gr.Markdown("# NPS Comment Categorization")
+
+    # Classify button
     classify_btn = gr.Button("Classify Comments")
     output = gr.HTML()
 
+    # Function to load data from uploaded CSV
     def load_data(file):
         if file is not None:
             file.seek(0)  # Reset file pointer
-            import io
             if file.name.endswith('.csv'):
                 custom_df = pd.read_csv(file, encoding='utf-8')
             else:
@@ -100,14 +86,35 @@ with gr.Blocks() as nps:
         else:
             return "No file uploaded."
 
-    uploaded_file.change(fn=load_data, inputs=uploaded_file, outputs=output)
-    template_btn.click(fn=lambda: "Using Template Dataset", outputs=output)
-
+    # Function to use template categories
     def use_template():
-        return ["Product Experience", "Customer Support", "Price of Service", "Other"]
-
-    template_btn.click(fn=use_template, outputs=category_boxes)
-    category_boxes.change(fn=display_categories, inputs=category_boxes, outputs=category_column)
-    classify_btn.click(fn=classify_comments, inputs=category_boxes, outputs=output)
+        template_categories = ["Product Experience", "Customer Support", "Price of Service", "Other"]
+        return template_categories, f"Categories: {', '.join(template_categories)}"
+
+    # Event handlers
+    add_category_btn.click(
+        fn=add_category,
+        inputs=[categories, category_input],
+        outputs=[categories, category_status]
+    )
+    categories.change(
+        fn=display_categories,
+        inputs=categories,
+        outputs=category_display
+    )
+    uploaded_file.change(
+        fn=load_data,
+        inputs=uploaded_file,
+        outputs=output
+    )
+    template_btn.click(
+        fn=use_template,
+        outputs=[categories, category_status]
+    )
+    classify_btn.click(
+        fn=classify_comments,
+        inputs=categories,
+        outputs=output
+    )
 
 nps.launch(share=True)