use pipeline instead of direct model load

app.py

@@ -16,15 +16,14 @@ import os
 hf_api_key = os.getenv("API_KEY")
 login(token=hf_api_key)
 
-pipe = pipeline("text-generation", model="mistralai/Mistral-Small-24B-Base-2501")
+pipe = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
 
 # 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 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']
+        prompt = comment
+        category = pipe(prompt)[0]['label']
         results.append(category)
     df['comment_category'] = results
     return df[['customer_comment', 'comment_category']].to_html(index=False)