import gradio as gr
from transformers import pipeline
import pandas as pd

# Load the dataset
DATASET_URL = 'https://huggingface.co/datasets/ZennyKenny/demo_customer_nps/resolve/main/customer_feedback_dataset.csv'
df = pd.read_csv(DATASET_URL)

# Initialize the model pipeline
pipe = pipeline("text-generation", model="mistralai/Mistral-Small-24B-Base-2501")

# Function to classify customer comments
def classify_comments():
    results = []
    for comment in df['customer_comment']:
        prompt = f"Classify this customer feedback: '{comment}' into one of five categories."
        category = pipe(prompt, max_length=30)[0]['generated_text']
        results.append(category)
    df['comment_category'] = results
    return df[['customer_comment', '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()