Update app.py

app.py

@@ -15,7 +15,8 @@ AVAILABLE_MODELS = [
 def classify_text(
         model_name: str, 
         text: str, 
-        labels: str
+        labels: str,
+        multi_label: bool = False,
     ) -> Dict[str, float]:
     """
     Classifies the input text into one of the provided labels using a zero-shot classification model.
@@ -39,7 +40,7 @@ def classify_text(
 
         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)
+        result = classifier(text, candidate_labels=labels_list, multi_label=multi_label)
         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."
@@ -102,8 +103,11 @@ with gr.Blocks(css=css) as iface:
     gr.Markdown("# Zero-Shot Text Classifier")
     gr.Markdown("Select a model, enter text, and a set of labels to classify the text using a zero-shot classification model.")
 
-    # Dropdown to select a model
-    model_dropdown = gr.Dropdown(AVAILABLE_MODELS, label="Choose Model")
+    with gr.Row():
+        # Dropdown to select a model
+        model_dropdown = gr.Dropdown(AVAILABLE_MODELS, label="Choose Model")
+        # Checkbox for multi-label classification
+        multi_label = gr.Checkbox(label="True", value=False, info="Check for multi-label classification, uncheck for single-label (multi-class).")
 
     # Input fields for text and labels
     with gr.Row():
@@ -115,7 +119,7 @@ with gr.Blocks(css=css) as iface:
 
     # Classification button
     submit_button = gr.Button("Classify", elem_classes=["classify-button"])
-    submit_button.click(fn=classify_text, inputs=[model_dropdown, text_input, label_input], outputs=output_label)
+    submit_button.click(fn=classify_text, inputs=[model_dropdown, text_input, label_input, multi_label], outputs=output_label)
 
     # Example input/output pairs
     gr.Examples(examples, inputs=[text_input, label_input])