ensure gene class dict only inverted when using gene classifier

geneformer/classifier_utils.py

@@ -115,13 +115,12 @@ def label_classes(classifier, data, gene_class_dict, nproc):
 
     class_id_dict = dict(zip(label_set, [i for i in range(len(label_set))]))
     id_class_dict = {v: k for k, v in class_id_dict.items()}
-    inverse_gene_class_dict = {}
-    # Iterate over each key and list of values in the original dictionary
-    for key, value_list in gene_class_dict.items():
-        # Iterate over each value in the list
-        for value in value_list:
-            # Assign the value as a key and the original key as its value in the new dictionary
-            inverse_gene_class_dict[value] = key
+
+    if classifier == "gene":
+        inverse_gene_class_dict = {}
+        for key, value_list in gene_class_dict.items():
+            for value in value_list:
+                inverse_gene_class_dict[value] = key
 
     def classes_to_ids(example):
         if classifier == "cell":