Remade the test script to add easily new models

models/ZeroShotClassifier/HuggingFace/zero_shot_classification.py

@@ -36,12 +36,19 @@ def create_classifier(model = MODELS[0]):
     return pipeline("zero-shot-classification", model=model, device=device)
 
 
-def classify(classifier, sequence, lables = LABELS, debug = False):
-    results = classifier(sequence, lables)
+def classify(classifier, sequence, lables, treshold, debug = False):
+    predictions = {}
 
     if debug:
         print(f"Sequence: {sequence}")
+    
+    for label in lables:
+            results = classifier(sequence, [label])
+
+            if debug:
                 print(f"Labels: {results['labels']}")
                 print(f"Scores: {results['scores']}")
 
-    return results
+            predictions[label] = results["scores"][0] > treshold
+
+    return predictions

testModel/metrics.py

@@ -22,10 +22,19 @@ def add_confusion_matrices(confusion_matrix, tmp_confusion_matrix):
     return confusion_matrix
 
 def get_tpr(confusion_matrix):
-    return confusion_matrix[0][0] / (confusion_matrix[0][0] + confusion_matrix[1][0])
+    denominator = (confusion_matrix[0][0] + confusion_matrix[1][0])
+    if denominator == 0:
+        return 0
+    return confusion_matrix[0][0] / denominator
 
 def get_tnr(confusion_matrix):
-    return confusion_matrix[1][1] / (confusion_matrix[1][1] + confusion_matrix[0][1])
+    denominator = (confusion_matrix[1][1] + confusion_matrix[0][1])
+    if denominator == 0:
+        return 0
+    return confusion_matrix[1][1] / denominator
 
 def get_precision(confusion_matrix):
-    return confusion_matrix[0][0] / (confusion_matrix[0][0] + confusion_matrix[0][1])
\ No newline at end of file
+    denominator = (confusion_matrix[0][0] + confusion_matrix[0][1])
+    if denominator == 0:
+        return 0
+    return confusion_matrix[0][0] / denominator
\ No newline at end of file