Update app.py

app.py

@@ -1,131 +1,90 @@
 import gradio as gr
-from transformers import pipeline
 import pandas as pd
-import spaces
-
-# Load dataset
-from datasets import load_dataset
-ds = load_dataset('ZennyKenny/demo_customer_nps')
-df = pd.DataFrame(ds['train'])
-
-# Initialize model pipeline
-from huggingface_hub import login
-import os
-
-# Login using the API key stored as an environment variable
-hf_api_key = os.getenv("API_KEY")
-login(token=hf_api_key)
+from transformers import pipeline
 
-classifier = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+# Initialize model pipelines
+classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 generator = pipeline("text2text-generation", model="google/flan-t5-base")
 
+# Initialize an empty DataFrame
+df = pd.DataFrame()
+
 # Function to classify customer comments
-# https://huggingface.co/docs/hub/spaces-zerogpu
-@spaces.GPU
-def classify_comments():
+def classify_comments(categories):
+    if df.empty:
+        return "No data loaded."
+    if not categories:
+        return "No categories defined."
+
     sentiments = []
-    categories = []
-    results = []
+    comment_categories = []
     for comment in df['customer_comment']:
         sentiment = classifier(comment)[0]['label']
-        category_list = [box.value for box in category_boxes if box.value.strip() != '']
-        category_str = ', '.join([cat.strip() for cat in category_list])
+        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)
         sentiments.append(sentiment)
+        comment_categories.append(category)
     df['comment_sentiment'] = sentiments
-    df['comment_category'] = categories
+    df['comment_category'] = comment_categories
     return df[['customer_comment', 'comment_sentiment', 'comment_category']].to_html(index=False)
 
+# Function to load data from uploaded file
+def load_data(file):
+    global df
+    if file is not None:
+        try:
+            df = pd.read_csv(file.name)
+            if 'customer_comment' not in df.columns:
+                return "Error: Uploaded CSV must contain a column named 'customer_comment'"
+            return "Custom CSV loaded successfully!"
+        except Exception as e:
+            return f"Error loading CSV: {e}"
+    else:
+        return "No file uploaded."
+
+# Function to add a new category
+def add_category(categories, new_category):
+    if new_category.strip() and new_category not in categories:
+        categories.append(new_category.strip())
+    return categories, gr.update(value="")  # Clear the input box after adding
+
+# Function to remove a category
+def remove_category(categories, category_to_remove):
+    categories = [cat for cat in categories if cat != category_to_remove]
+    return categories
+
+# Function to display categories with remove buttons
+def display_categories(categories):
+    category_elements = []
+    for category in categories:
+        with gr.Row() as category_row:
+            gr.Markdown(f"- {category}")
+            remove_button = gr.Button("Remove")
+            remove_button.click(fn=remove_category, inputs=[gr.State(categories), gr.State(category)], outputs=gr.State(categories), queue=False)
+        category_elements.append(category_row)
+    return category_elements
+
 # 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
+    gr.Markdown("# NPS Comment Categorization")
 
-    category_boxes = gr.State([])  # Store category input boxes as state
-    
-    def display_categories(categories):
-        category_column.children = []  # Clear previous categories
-        for i, cat in enumerate(categories):
-            with category_column:
-                with gr.Row():
-                    gr.Markdown(f"- {cat}")
-                    remove_btn = gr.Button("X", elem_id=f"remove_{i}", interactive=True)
-                    remove_btn.click(fn=lambda x=cat: remove_category(x, categories), inputs=[], outputs=category_boxes)
-        category_column.children = []  # Reset children to clear previous categories
-        for i, cat in enumerate(categories):
-            with category_column:
-                with gr.Row():
-                    gr.Markdown(f"- {cat}")
-                    remove_btn = gr.Button("X", elem_id=f"remove_{i}", interactive=True)
-                    remove_btn.click(fn=lambda x=cat: remove_category(x, categories), inputs=[], outputs=category_boxes)
-        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=lambda x=cat: remove_category(x, categories), inputs=[], outputs=category_boxes)
-        return category_components
     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)
-    def remove_category(category, category_list):
-        category_list.remove(category)  # Remove selected category
-        return category_list
-        components = []
-        for i, cat in enumerate(categories):
-            row = gr.Row([
-                gr.Markdown(f"- {cat}"),
-                gr.Button("X", elem_id=f"remove_{i}", interactive=True).click(fn=lambda x=cat: remove_category(x, categories), inputs=[], outputs=category_boxes)
-            ])
-            components.append(row)
-        return components
-        for i, cat in enumerate(categories):
-            row = gr.Row([
-                gr.Textbox(value=cat, label=f"Category {i+1}", interactive=True),
-                gr.Button("X", elem_id=f"remove_{i}")
-            ])
-            components.append(row)
-        return components
+
     category_column = gr.Column()
-    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_boxes.change(fn=display_categories, inputs=category_boxes, outputs=category_column)
-    uploaded_file = gr.File(label="Upload CSV", type="filepath")
-    template_btn = gr.Button("Use Template")
-    gr.Markdown("# NPS Comment Categorization")
-    classify_btn = gr.Button("Classify Comments")
-    output = gr.HTML()
+    categories = gr.State([])  # Initialize an empty list for categories
 
-    def load_data(file):
-        if file is not None:
-            file.seek(0)  # Reset file pointer
-            import io
-            if file.name.endswith('.csv'):
-                file.seek(0)  # Reset file pointer
-                custom_df = pd.read_csv(file, encoding='utf-8')
-                custom_df = pd.read_csv(io.StringIO(content))
-            else:
-                return "Error: Uploaded file is not a CSV."
-            if 'customer_comment' not in custom_df.columns:
-                return "Error: Uploaded CSV must contain a column named 'customer_comment'"
-            global df
-            df = custom_df
-            return "Custom CSV loaded successfully!"
-        else:
-            return "No file uploaded."
+    add_category_btn.click(fn=add_category, inputs=[categories, category_input], outputs=[categories, category_input], queue=False)
+    categories.change(fn=display_categories, inputs=categories, outputs=category_column, queue=False)
+
+    uploaded_file = gr.File(label="Upload CSV", type="file")
+    upload_output = gr.HTML()
+    uploaded_file.change(fn=load_data, inputs=uploaded_file, outputs=upload_output, queue=False)
 
-    uploaded_file.change(fn=load_data, inputs=uploaded_file, outputs=output)
-    template_btn.click(fn=lambda: "Using Template Dataset", outputs=output)
-    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)
+    classify_btn = gr.Button("Classify Comments")
+    output = gr.HTML()
+    classify_btn.click(fn=classify_comments, inputs=categories, outputs=output, queue=False)
 
 nps.launch()