app.py

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)

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):
    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[['customer_comment', 'comment_sentiment', 'comment_category']].to_html(index=False)

# Gradio Interface
with gr.Blocks() as nps:
    # State to store categories
    categories = gr.State([])

    # 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):
        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")
        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
            if file.name.endswith('.csv'):
                custom_df = pd.read_csv(file, encoding='utf-8')
            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."

    # Function to use template categories
    def use_template():
        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)