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
# https://huggingface.co/docs/hub/spaces-zerogpu
@spaces.GPU
def classify_comments():
    sentiments = []
    categories = []
    results = []
    for comment in df['customer_comment']:
        # Classify the sentiment first
        sentiment = classifier(comment)[0]['label']
        prompt = f"What category best describes this comment? '{comment}' Please answer using only the name of the category: Product Experience, Customer Support, Price of Service, Other."
        category = generator(prompt, max_length=30)[0]['generated_text']
        categories.append(category)
        sentiments.append(sentiment)
    df['comment_sentiment'] = sentiments
    df['comment_category'] = categories
    return df[['customer_comment', 'comment_sentiment', 'comment_category']].to_html(index=False)

# Gradio Interface
with gr.Blocks() as nps:
    gr.Markdown("# NPS Comment Categorization")
    classify_btn = gr.Button("Classify Comments")
    output = gr.HTML()

    classify_btn.click(fn=classify_comments, outputs=output)

nps.launch()