Update geneformer/emb_extractor.py

geneformer/emb_extractor.py

@@ -411,7 +411,7 @@ class EmbExtractor:
         self,
         model_type="Pretrained",
         num_classes=0,
-        emb_mode="cell",
+        emb_mode="cls",
         cell_emb_style="mean_pool",
         gene_emb_style="mean_pool",
         filter_data=None,
@@ -596,11 +596,11 @@ class EmbExtractor:
         filtered_input_data = pu.load_and_filter(
             self.filter_data, self.nproc, input_data_file
         )
-        
+
         # Check to make sure that all the labels exist in the tokenized data:
         if self.emb_label is not None:
             for label in self.emb_label:
-                assert label in list(filtered_input_data.features), f"Attribute `{label}` not present in dataset features" 
+                assert label in filtered_input_data.features.keys(), f"Attribute `{label}` not present in dataset features"
 
         if cell_state is not None:
             filtered_input_data = pu.filter_by_dict(
@@ -725,6 +725,12 @@ class EmbExtractor:
             )
             raise
 
+        if self.emb_label is not None:
+            logger.error(
+                "For extracting state embs, emb_label should be None since labels are based on state embs dict keys."
+            )
+            raise
+        
         state_embs_dict = dict()
         state_key = cell_states_to_model["state_key"]
         for k, v in cell_states_to_model.items():