add category output

app.py

@@ -16,18 +16,24 @@ import os
 hf_api_key = os.getenv("API_KEY")
 login(token=hf_api_key)
 
-pipe = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+classifier = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+generator = pipeline("text-generation", model="mrm8488/t5-base-finetuned-question-generation-ap")
 
 # Function to classify customer comments
-@spaces.GPU
 def classify_comments():
+    sentiments = []
+    categories = []
     results = []
     for comment in df['customer_comment']:
-        prompt = comment
-        category = pipe(prompt)[0]['label']
-        results.append(category)
-    df['comment_category'] = results
-    return df[['customer_comment', 'comment_category']].to_html(index=False)
+        # Classify the sentiment first
+        sentiment = pipe(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: