embs_df with all model embeddings

geneformer/emb_extractor.py

@@ -49,10 +49,8 @@ def get_embs(
     if summary_stat is None:
         embs_list = []
     elif summary_stat is not None:
-        # test embedding extraction for example cell and extract # emb dims
-        example = filtered_input_data.select([i for i in range(1)])
-        example.set_format(type="torch")
-        emb_dims = test_emb(model, example["input_ids"], layer_to_quant)
+        # get # of emb dims
+        emb_dims = pu.get_model_emb_dims(model)
         if emb_mode == "cell":
             # initiate tdigests for # of emb dims
             embs_tdigests = [TDigest() for _ in range(emb_dims)]
@@ -239,14 +237,6 @@ def tdigest_median(embs_tdigests, emb_dims):
     return [embs_tdigests[i].percentile(50) for i in range(emb_dims)]
 
 
-def test_emb(model, example, layer_to_quant):
-    with torch.no_grad():
-        outputs = model(input_ids=example.to("cuda"))
-
-    embs_test = outputs.hidden_states[layer_to_quant]
-    return embs_test.size()[2]
-
-
 def label_cell_embs(embs, downsampled_data, emb_labels):
     embs_df = pd.DataFrame(embs.cpu().numpy())
     if emb_labels is not None:
@@ -632,13 +622,15 @@ class EmbExtractor:
 
         if self.exact_summary_stat == "exact_mean":
             embs = embs.mean(dim=0)
+            emb_dims = pu.get_model_embedding_dimensions(model)
             embs_df = pd.DataFrame(
-                embs_df[0:255].mean(axis="rows"), columns=[self.exact_summary_stat]
+                embs_df[0:emb_dims-1].mean(axis="rows"), columns=[self.exact_summary_stat]
             ).T
         elif self.exact_summary_stat == "exact_median":
             embs = torch.median(embs, dim=0)[0]
+            emb_dims = pu.get_model_embedding_dimensions(model)
             embs_df = pd.DataFrame(
-                embs_df[0:255].median(axis="rows"), columns=[self.exact_summary_stat]
+                embs_df[0:emb_dims-1].median(axis="rows"), columns=[self.exact_summary_stat]
             ).T
 
         if cell_state is not None: