feat: Update app

app.py

@@ -11,6 +11,15 @@ classifier = pipeline(
 )
 
 
+# Set the description for the interface
+DESCRIPTION = """Classify text in Danish, Swedish or Norwegian into categories, without
+any training data!
+
+Note that the models will most likely *not* work as well as a finetuned model on your
+specific data, but they can be used as a starting point for your own classification
+task ✨"""
+
+
 def classification(task: str, doc: str) -> str:
     """Classify text into categories.
 
@@ -93,40 +102,55 @@ def classification(task: str, doc: str) -> str:
                 "sport",
             ]
 
-    # Else if the task is offensive text detection, classify the text into offensive
-    # or not offensive
-    elif task == "Offensive text detection":
+    # Else if the task is spam detection, classify the text into spam or not spam
+    elif task == "Spam detection":
         if language == "sv":
-            hypothesis_template = "Detta exempel er {}."
-            candidate_labels = ["stötande", "inte stötande"]
+            hypothesis_template = "Det här e-postmeddelandet ser {}"
+            candidate_labels = {
+                "ut som ett skräppostmeddelande": "Spam",
+                "inte ut som ett skräppostmeddelande": "Inte spam",
+            }
         elif language == "no":
-            hypothesis_template = "Dette eksemplet er {}."
-            candidate_labels = ["støtende", "ikke støtende"]
+            hypothesis_template = "Denne e-posten ser {}"
+            candidate_labels = {
+                "ut som en spam-e-post": "Spam",
+                "ikke ut som en spam-e-post": "Ikke spam",
+            }
         else:
-            hypothesis_template = "Dette eksempel er {}."
-            candidate_labels = ["anstødig tale", "ikke anstødig tale"]
+            hypothesis_template = "Denne e-mail ligner {}"
+            candidate_labels = {
+                "en spam e-mail": "Spam",
+                "ikke en spam e-mail": "Ikke spam",
+            }
 
     # Else the task is not supported, so raise an error
     else:
         raise ValueError(f"Task {task} not supported.")
 
+    # If `candidate_labels` is a list then convert it to a dictionary, where the keys
+    # are the entries in the list and the values are the keys capitalized
+    if isinstance(candidate_labels, list):
+        candidate_labels = {label: label.capitalize() for label in candidate_labels}
+
     # Run the classifier on the text
     result = classifier(
-        doc, candidate_labels=candidate_labels, hypothesis_template=hypothesis_template
+        doc,
+        candidate_labels=list(candidate_labels.keys()),
+        hypothesis_template=hypothesis_template,
     )
 
     print(result)
 
     # Return the predicted label
     return (
-        f"{result['labels'][0].capitalize()}\n"
+        f"{candidate_labels[result['labels'][0]]}\n"
         f"({confidence_str}: {result['scores'][0]:.0%})"
     )
 
 # Create a dropdown menu for the task
 dropdown = gr.inputs.Dropdown(
     label="Task",
-    choices=["Sentiment classification", "News topic classification", "Offensive text detection"],
+    choices=["Sentiment classification", "News topic classification", "Spam detection"],
     default="Sentiment classification",
 )
 
@@ -136,7 +160,7 @@ interface = gr.Interface(
     inputs=[dropdown, gr.inputs.Textbox(label="Text")],
     outputs=gr.outputs.Label(type="text"),
     title="Scandinavian zero-shot text classification",
-    description="Classify text in Danish, Swedish or Norwegian into categories, without any training data!",
+    description=DESCRIPTION,
 )
 
 # Run the app