Update src/app/predict.py

src/app/predict.py

@@ -33,8 +33,13 @@ def ZeroShotTextClassification(
         logging.info(f"Attempting classification with {len(labels)} labels")
 
         # Perform zero-shot classification
-        classifier = pipeline("zero-shot-classification")
-        prediction = classifier(text_input, labels, multi_label=True)
+        hypothesis_template = "This text is about {}"
+        prediction = classifier(
+            text_input,
+            labels,
+            hypothesis_template=hypothesis_template,
+            multi_label=True,
+        )
 
         # Return the classification results
         logging.info("Classification completed successfully")