eval_downstream.py

Author: natuan

src/deepsparse/transformers/eval_downstream.py

@@ -122,7 +122,10 @@ def mnli_eval(args):
     )
     print(f"Engine info: {text_classify.engine}")
 
-    label_map = _get_label2id(text_classify.config_path)
+    try:
+        label_map = _get_label2id(text_classify.config_path)
+    except KeyError:
+        label_map = {"entailment": 0, "neutral": 1, "contradiction": 2}
 
     for idx, sample in _enumerate_progress(mnli_matched, args.max_samples):
         pred = text_classify([[sample["premise"], sample["hypothesis"]]])
@@ -162,7 +165,10 @@ def qqp_eval(args):
     )
     print(f"Engine info: {text_classify.engine}")
 
-    label_map = _get_label2id(text_classify.config_path)
+    try:
+        label_map = _get_label2id(text_classify.config_path)
+    except KeyError:
+        label_map = {"not_duplicate": 0, "duplicate": 1, "LABEL_0": 0, "LABEL_1": 1}
 
     for idx, sample in _enumerate_progress(qqp, args.max_samples):
         pred = text_classify([[sample["question1"], sample["question2"]]])
@@ -193,7 +199,10 @@ def sst2_eval(args):
     )
     print(f"Engine info: {text_classify.engine}")
 
-    label_map = _get_label2id(text_classify.config_path)
+    try:
+        label_map = _get_label2id(text_classify.config_path)
+    except KeyError:
+        label_map = {"negative": 0, "positive": 1, "LABEL_0": 0, "LABEL_1": 1}
 
     for idx, sample in _enumerate_progress(sst2, args.max_samples):
         pred = text_classify(