classifier: iterates through mailing list

1 files changed, 3 insertions, 4 deletions

diff --git a/classification/test.py b/classification/test.py
index a8bc7e1e..bcd6b439 100644..100755
--- a/classification/test.py
+++ b/classification/test.py
@@ -2,16 +2,15 @@ from os import listdir, path
 
 directory : str = "./test_input"
 
-def test(classifier):
+def test(classifier, categories):
     for name in listdir(directory):
         if name == "README.md":
             continue
 
         with open(path.join(directory, name), "r") as file:
-            sequence_to_classify = file.read()
+            text = file.read()
 
-        candidate_labels = ['semantic', 'other', 'mistranslation', 'instruction']
-        result = classifier(sequence_to_classify, candidate_labels, multi_label=True)
+        result = classifier(text, categories, multi_label=True)
 
         print(name)
         for label, score in zip(result["labels"], result["scores"]):