Fix mapping preview

text_classification.py

@@ -91,22 +91,27 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
             return column_mapping, prediction_result, None
 
         if isinstance(column_mapping["label"], dict):
-            for model_label in id2label_mapping.keys():
-                id2label_mapping[model_label] = column_mapping["label"][str(label2id[model_label])]
+            # Use the column mapping passed by user
+            for i, model_label in column_mapping["label"].items():
+                id2label_mapping[model_label] = dataset_labels[int(i)]
     elif None in id2label_mapping.values():
         column_mapping["label"] = {
             i: None for i in id2label.keys()
         }
         return column_mapping, prediction_result, None
 
+    id2label_mapping = {
+        v: k for k, v in id2label_mapping.items()
+    }
     id2label_df = pd.DataFrame({
-        "ID": [i for i in id2label.keys()],
-        "Model labels": [id2label[label] for label in id2label.keys()],
-        "Dataset labels": [id2label_mapping[id2label[label]] for label in id2label.keys()],
+        "ID": list(range(len(dataset_labels))),
+        "Labels": dataset_labels,
+        "Labels in original model": [f"{id2label_mapping[label]}({label2id[id2label_mapping[label]]})" for label in dataset_labels],
     })
     if "label" not in column_mapping.keys():
+        # Column mapping should contain original model labels
         column_mapping["label"] = {
-            i: id2label_mapping[id2label[i]] for i in id2label.keys()
+            str(i): id2label_mapping[label] for i, label in zip(id2label.keys(), dataset_labels)
         }
 
     return column_mapping, prediction_result, id2label_df