Update app.py

app.py

@@ -1,90 +1,131 @@
 import gradio as gr
-import pandas as pd
 from transformers import pipeline
+import pandas as pd
+import spaces
 
-# Initialize model pipelines
-classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
-generator = pipeline("text2text-generation", model="google/flan-t5-base")
+# Load dataset
+from datasets import load_dataset
+ds = load_dataset('ZennyKenny/demo_customer_nps')
+df = pd.DataFrame(ds['train'])
 
-# Initialize an empty DataFrame
-df = pd.DataFrame()
+# Initialize model pipeline
+from huggingface_hub import login
+import os
 
-# Function to classify customer comments
-def classify_comments(categories):
-    if df.empty:
-        return "No data loaded."
-    if not categories:
-        return "No categories defined."
+# Login using the API key stored as an environment variable
+hf_api_key = os.getenv("API_KEY")
+login(token=hf_api_key)
+
+classifier = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+generator = pipeline("text2text-generation", model="google/flan-t5-base")
 
+# Function to classify customer comments
+# https://huggingface.co/docs/hub/spaces-zerogpu
+@spaces.GPU
+def classify_comments():
     sentiments = []
-    comment_categories = []
+    categories = []
+    results = []
     for comment in df['customer_comment']:
         sentiment = classifier(comment)[0]['label']
-        category_str = ', '.join(categories)
+        category_list = [box.value for box in category_boxes if box.value.strip() != '']
+        category_str = ', '.join([cat.strip() for cat in category_list])
         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'] = comment_categories
+    df['comment_category'] = 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:
-    gr.Markdown("# NPS Comment Categorization")
+    def add_category(category_list, new_category):
+        if new_category.strip() != "":
+            category_list.append(new_category.strip())  # Add new category
+        return category_list
 
+    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()
-    categories = gr.State([])  # Initialize an empty list for categories
-
-    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)
-
+    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()
-    classify_btn.click(fn=classify_comments, inputs=categories, outputs=output, queue=False)
+
+    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."
+
+    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)
 
 nps.launch()