Merge remote-tracking branch 'origin/main'

geneformer/gene_name_id_dict_gc95M.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8b0fd0521406ed18b2e341ef0acb5f53aa1a62457a07ca5840e1c142f46dd326
-size 2038812
+oid sha256:fabfa0c2f49c598c59ae432a32c3499a5908c033756c663b5e0cddf58deea8e1
+size 1660882

geneformer/mtl/data.py

@@ -112,7 +112,7 @@ def preload_and_process_data(config):
     # Validate that the mappings match
     validate_label_mappings(config)
 
-    return (*train_data, *val_data[:2])  # Return train and val data along with mappings
+    return (*train_data[:2], *val_data)  # Return train and val data along with mappings
 
 
 def validate_label_mappings(config):

geneformer/tokenizer.py

@@ -88,6 +88,7 @@ def sum_ensembl_ids(
     collapse_gene_ids,
     gene_mapping_dict,
     gene_token_dict,
+    custom_attr_name_dict,
     file_format="loom",
     chunk_size=512,
 ):
@@ -104,6 +105,13 @@ def sum_ensembl_ids(
                 "ensembl_id_collapsed" not in data.ra.keys()
             ), "'ensembl_id_collapsed' column already exists in data.ra.keys()"
             
+            assert (
+                "n_counts" in data.ca.keys()
+            ), "'n_counts' column missing from data.ca.keys()"
+
+            if custom_attr_name_dict is not None:
+                for label in custom_attr_name_dict:
+                    assert label in data.ca.keys(), f"Attribute `{label}` not present in dataset features"
 
             # Get the ensembl ids that exist in data
             ensembl_ids = data.ra.ensembl_id
@@ -208,6 +216,13 @@ def sum_ensembl_ids(
         assert (
             "ensembl_id_collapsed" not in data.var.columns
         ), "'ensembl_id_collapsed' column already exists in data.var"
+        assert (
+            "n_counts" in data.obs.columns
+        ), "'n_counts' column missing from data.obs"
+
+        if custom_attr_name_dict is not None:
+            for label in custom_attr_name_dict:
+                assert label in data.obs.columns, f"Attribute `{label}` not present in data.obs"
 
 
         # Get the ensembl ids that exist in data
@@ -461,6 +476,7 @@ class TranscriptomeTokenizer:
             self.collapse_gene_ids,
             self.gene_mapping_dict,
             self.gene_token_dict,
+            self.custom_attr_name_dict, 
             file_format="h5ad",
             chunk_size=self.chunk_size,
         )
@@ -537,6 +553,7 @@ class TranscriptomeTokenizer:
             self.collapse_gene_ids,
             self.gene_mapping_dict,
             self.gene_token_dict,
+            self.custom_attr_name_dict,
             file_format="loom",
             chunk_size=self.chunk_size,
         )