Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import torch
 from transformers import pipeline
 from typing import Dict
 
@@ -27,10 +28,21 @@ def classify_text(
     Returns:
         Dict[str, float]: A dictionary mapping each label to its classification score.
     """
-    classifier = pipeline("zero-shot-classification", model=model_name)
-    labels_list = [label.strip() for label in labels.split(",")]
-    result = classifier(text, candidate_labels=labels_list)
-    return {label: score for label, score in zip(result["labels"], result["scores"])}
+    if not text.strip():
+        return "Error: Please enter some text to classify."
+    if not labels.strip():
+        return "Error: Please enter some labels to classify the text."
+    
+    try:
+        # Set device: 0 if GPU available, else -1 for CPU
+        device = 0 if torch.cuda.is_available() else -1
+
+        classifier = pipeline("zero-shot-classification", model=model_name, device=device)
+        labels_list = [label.strip() for label in labels.split(",")]
+        result = classifier(text, candidate_labels=labels_list)
+        return {label: score for label, score in zip(result["labels"], result["scores"])}
+    except Exception as _:
+        return "Error: An unexpected error occurred. Please try again later."
 
 # Example Input with Mutually Exclusive Labels from News Articles
 examples = [