Update app.py

app.py

@@ -1,105 +1,178 @@
-import os
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import pipeline
+import pandas as pd
 import spaces
 
-# Load API token from environment variable
-API_TOKEN = os.getenv("API_TOKEN")
-
-if not API_TOKEN:
-    raise ValueError("API_TOKEN environment variable is not set.")
-
-client = InferenceClient(model="mistralai/Mistral-Nemo-Instruct-2407", token=API_TOKEN)
-
-SYSTEM_MESSAGE = ("""
-Your name is Monty, and you represent montebello.ai. Your goal is to engage customers about montebello.ai’s AI solutions and encourage them to either schedule a call by clicking the "Call Me" button or by calling (415)-862-1695.
-
-**About montebello.ai:**
-We provide affordable, tailored AI solutions for small businesses, helping them save time, reduce costs, and improve efficiency.
-
-**Our Services:**
-- **AI Phone Assistant:** Set up in days for as little as $200/month to handle calls, customer inquiries, and appointment scheduling.
-- **Data Insights:** Analyze sales, inventory, and customer behavior for smarter business decisions.
-- **Productivity Boost:** Streamline workflows to free up your team’s time.
-- **Customer Experience:** Enhance engagement with AI-powered chatbots and personalized marketing.
-
-**Your Approach:**
-1. **Start by Asking:** 
-   - 'What would you like to know about montebello.ai?'
-
-2. **Provide Specific Examples:** 
-   - 'An AI Phone Agent can save you hours each week by handling common inquiries.’
-   - 'An AI Email Agent can reduce the amount of incoming emails that your team needs to handle by a significant amount.'
-   - 'An AI Appointment Agent can help you manage your time in the most effective way possible, freeing you up to work with your customers.'
-   
-
-3. **Highlight Affordability and Speed:** 
-   - 'We can set up you AI Agent in just a few days for as little as $200/month.’
-
-4. **Encourage Action:** 
-   - 'To experience our AI Phone Agent firsthand, call (415)-862-1695.’
-   - 'To experience our AI Email Agent firsthand, write us at [email protected].’
-
-**Tone and Style:**
-- Friendly, energetic, and professional.
-- Keep responses very concise and focused on solving their problems.
-- Avoid technical jargon unless the customer is familiar with AI.
-
-**Closing:** 
-End by reiterating how montebello.ai can help and encourage them to take the next step by calling us, emailing us, or setting up an appointment.
-""")
-
+# Load dataset
+from datasets import load_dataset
+ds = load_dataset('ZennyKenny/demo_customer_nps')
+df = pd.DataFrame(ds['train'])
 
-MAX_TOKENS = 256
-TEMPERATURE = 0.7
-TOP_P = 0.95
-
-# @spaces.GPU
-def respond(message, history: list[tuple[str, str]]):
-    messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=MAX_TOKENS,
-        stream=True,
-        temperature=TEMPERATURE,
-        top_p=TOP_P,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
+# Initialize model pipeline
+from huggingface_hub import login
+import os
 
-# Custom CSS
+# 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
+@spaces.GPU
+def classify_comments(categories):
+    global df  # Ensure we're modifying the global DataFrame
+    sentiments = []
+    assigned_categories = []
+    for comment in df['customer_comment']:
+        # Classify sentiment
+        sentiment = classifier(comment)[0]['label']
+        # 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']
+        assigned_categories.append(category)
+        sentiments.append(sentiment)
+    df['comment_sentiment'] = sentiments
+    df['comment_category'] = assigned_categories
+    return df.to_html(index=False)  # Return all fields with appended sentiment and category
+
+# 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:**\n" + "\n".join([f"- {cat}" for cat in categories])
+
+# Function to reset categories
+def reset_categories():
+    return [], "**Categories:**\n- None"
+
+# Function to load data from uploaded CSV
+def load_data(file):
+    global df  # Ensure we're modifying the global DataFrame
+    if file is not None:
+        file.seek(0)  # Reset file pointer
+        if file.name.endswith('.csv'):
+            custom_df = pd.read_csv(file, encoding='utf-8')
+        else:
+            return "Error: Uploaded file is not a CSV."
+        # Check for required columns
+        required_columns = ['customer_comment']
+        if not all(col in custom_df.columns for col in required_columns):
+            return f"Error: Uploaded CSV must contain the following column: {', '.join(required_columns)}"
+        df = custom_df
+        return "Custom CSV loaded successfully!"
+    else:
+        return "No file uploaded."
+
+# Function to use template categories
+def use_template():
+    template_categories = ["Product Experience", "Customer Support", "Price of Service", "Other"]
+    return template_categories, f"**Categories:**\n" + "\n".join([f"- {cat}" for cat in template_categories])
+
+# Custom CSS for button styling
 custom_css = """
-body {
-    background-color: #f9d029 !important;
+button {
+    background-color: #4c6bfd !important;
+    border: none !important;
+    color: white !important;
+    padding: 10px 20px !important;
+    text-align: center !important;
+    text-decoration: none !important;
+    display: inline-block !important;
+    font-size: 16px !important;
+    border-radius: 5px !important;
+    transition: background-color 0.3s ease !important;
+}
+
+button:hover {
+    background-color: #3a52cc !important;
 }
 footer {
     visibility: hidden;
 }
-.submit-button {
-    background-color: #f9d029 !important;
-    color: #000000 !important;
-    border: none !important;
-    border-radius: 5px;
-    padding: 10px 20px;
-    font-weight: bold;
-}
-.submit-button:hover {
-    background-color: #e6b800 !important;
-}
 """
 
-montebello = gr.ChatInterface(respond, css=custom_css)
-
-if __name__ == "__main__":
-    montebello.launch()
+# Gradio Interface
+with gr.Blocks(css=custom_css) as nps:
+    # State to store categories
+    categories = gr.State([])
+
+    # App title
+    gr.Markdown("# 🎉 Customer Comment Classifier 🎉")
+
+    # Short explanation
+    gr.Markdown("""
+    This app classifies customer comments into categories and assigns sentiment labels (Positive/Negative). 
+    You can upload your own dataset or use the provided template. The app will append the generated 
+    `comment_sentiment` and `comment_category` fields to your dataset.
+    """)
+
+    # File upload and instructions
+    with gr.Row():
+        with gr.Column(scale=1):
+            uploaded_file = gr.File(label="📂 Upload CSV", type="filepath", scale=1)
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            **📝 Instructions:**
+            - Upload a CSV file with at least one column: `customer_comment`.
+            - If you don't have your own data, click **Use Template** to load a sample dataset.
+            """)
+            template_btn = gr.Button("✨ Use Template", size="sm")
+
+    gr.Markdown("---")
+
+    # Category section
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Category input and buttons
+            category_input = gr.Textbox(label="📝 New Category", placeholder="Enter category name", scale=1)
+            with gr.Row():
+                add_category_btn = gr.Button("➕ Add Category", size="sm")
+                reset_btn = gr.Button("🔄 Reset Categories", size="sm")
+            # Category display
+            category_status = gr.Markdown("**📂 Categories:**\n- None")
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            **📝 Instructions:**
+            - Enter a category name and click **Add Category** to add it to the list.
+            - Click **Reset Categories** to clear the list.
+            - The `customer_comment` field will be categorized based on the categories you provide.
+            """)
+
+    gr.Markdown("---")
+
+    # Classify button and output
+    with gr.Row():
+        with gr.Column(scale=1):
+            classify_btn = gr.Button("🔍 Classify", size="sm")
+        with gr.Column(scale=3):  # Center the container and make it 75% of the window width
+            output = gr.HTML()
+
+    # Event handlers
+    add_category_btn.click(
+        fn=add_category,
+        inputs=[categories, category_input],
+        outputs=[categories, category_input, category_status]
+    )
+    reset_btn.click(
+        fn=reset_categories,
+        outputs=[categories, category_status]
+    )
+    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)