Add function to get number of model embeddings

geneformer/perturber_utils.py

@@ -156,8 +156,12 @@ def quant_layers(model):
     return int(max(layer_nums)) + 1
 
 
+def get_model_emb_dims(model):
+    return model.config.hidden_size
+
+
 def get_model_input_size(model):
-    return int(re.split("\(|,", str(model.bert.embeddings.position_embeddings))[1])
+    return model.config.max_position_embeddings
 
 
 def flatten_list(megalist):