add auth

app.py

@@ -1,20 +1,29 @@
 import gradio as gr
 from transformers import pipeline
-from datasets import load_dataset
 import pandas as pd
 
-# Load the dataset from Hugging Face
+# Load the dataset
+DATASET_URL = 'https://huggingface.co/datasets/ZennyKenny/demo_customer_nps/resolve/main/customer_feedback_dataset.csv'
+from datasets import load_dataset
 ds = load_dataset('ZennyKenny/demo_customer_nps')
 df = pd.DataFrame(ds['train'])
 
 # Initialize the 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)
+
 pipe = pipeline("text-generation", model="mistralai/Mistral-Small-24B-Base-2501")
 
 # Function to classify customer comments
+@spaces.GPU
 def classify_comments():
     results = []
     for comment in df['customer_comment']:
-        prompt = f"Classify this customer feedback: '{comment}' into one of five categories."
+        prompt = f"Classify this customer feedback: '{comment}' into one of the following categories: Price of Service, Quality of Customer Support, Product Experience. Please only respond with the category name and nothing else."
         category = pipe(prompt, max_length=30)[0]['generated_text']
         results.append(category)
     df['comment_category'] = results