Update geneformer/emb_extractor.py (#453)

- Update geneformer/emb_extractor.py (2c8d3f5d8ebb362ad102cb7e924d84a39b7349c8)
- Update geneformer/emb_extractor.py (62074538699215f6f0f8aca01d2e8f974386d800)


Co-authored-by: Han Chen <hchen725@users.noreply.huggingface.co>

geneformer/emb_extractor.py

@@ -596,6 +596,12 @@ 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 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(
                 filtered_input_data, cell_state, self.nproc