Update src/app/predict.py

src/app/predict.py

@@ -8,16 +8,13 @@ from transformers import pipeline
 
 
 # Load the zero-shot classification model
-device = "cuda" if torch.cuda.is_available() else "cpu"
 classifier = pipeline(
-    "zero-shot-classification",
-    model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
-    device=device,
+    "zero-shot-classification", model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
 )
 
 
 def ZeroShotTextClassification(
-    text_input: str, candidate_labels: str
+    text_input: str, candidate_labels: str, multi_label: bool
 ) -> Dict[str, float]:
     """
     Performs zero-shot classification on the given text input.
@@ -25,6 +22,7 @@ def ZeroShotTextClassification(
     Args:
         - text_input: The input text to classify.
         - candidate_labels: A comma-separated string of candidate labels.
+        - multi_label: A boolean indicating whether to allow the model to choose multiple classes.
 
     Returns:
         Dictionary containing label-score pairs.
@@ -42,7 +40,7 @@ def ZeroShotTextClassification(
             text_input,
             labels,
             hypothesis_template=hypothesis_template,
-            multi_label=False,
+            multi_label=multi_label,
         )
 
         # Return the classification results