update unit tests

src/scgpt/embedding/config.vsh.yaml

@@ -42,27 +42,32 @@ functionality:
           example: args.json
           description: |
             Path to model config file.
-        - name: "--input_obsm_gene_tokens"
+        - name: "--obsm_gene_tokens"
+          required: true
           type: string
           default: "gene_id_tokens"
           description: |
             The key of the .obsm array containing the gene token ids
           example: values.pt
-        - name: "--input_obsm_tokenized_values"
+        - name: "--obsm_tokenized_values"
           type: string
+          required: true
           default: values_tokenized
           description: |
             The key of the .obsm array containing the count values of the tokenized genes
-        - name: "--input_obsm_padding_mask"
+        - name: "--obsm_padding_mask"
           type: string
+          required: true
           default: padding_mask
           description: |
             The key of the .obsm array containing the padding mask.
-        - name: "--input_var_gene_names"
+        - name: "--var_gene_names"
           type: string
+          required: true
           description: |
             The name of the .var column containing gene names. When no gene_name_layer is provided, the .var index will be used.
-        - name: "--input_obs_batch_label"
+        - name: "--obs_batch_label"
+          required: true
           type: string
           description: |
             The name of the adata.obs column containing the batch labels.
@@ -83,7 +88,7 @@ functionality:
           choices: ["gzip", "lzf"]
           description: |
             The compression algorithm to use for the output h5mu file.
-        - name: "--embedding_layer_key"
+        - name: "--obsm_embeddings"
           type: string
           default: "X_scGPT"
           required: false
@@ -113,7 +118,7 @@ functionality:
           type: boolean
           default: true
           description: |
-            Apply domain-specific batch normalization
+            Apply domain-specific batch normalization. When set to True, 'obs_batch_labels' must be set as well.
         - name: "--batch_size"
           type: integer
           default: 64

src/scgpt/embedding/script.py

@@ -9,29 +9,23 @@
 ## VIASH START
 par = {
     "input": "resources_test/scgpt/test_resources/Kim2020_Lung_tokenized.h5mu",
-    "input_obsm_gene_tokens": 'gene_id_tokens',
-    "input_obsm_tokenized_values": 'values_tokenized',
-    "input_obsm_padding_mask": 'padding_mask',
+    "obsm_gene_tokens": 'gene_id_tokens',
+    "obsm_tokenized_values": 'values_tokenized',
+    "obsm_padding_mask": 'padding_mask',
     "model": "resources_test/scgpt/source/best_model.pt",
     "model_config": "resources_test/scgpt/source/args.json",
     "model_vocab": "resources_test/scgpt/source/vocab.json",
     "output": "Kim2020_Lung_embedded.h5ad",
-    "input_var_gene_names": "gene_name",
-    "input_obs_batch_label": "sample",
-    "embedding_layer_key": "X_scGPT",
+    "var_gene_names": "gene_name",
+    "obs_batch_label": "sample",
+    "obsm_embeddings": "X_scGPT",
     "pad_token": "<pad>",
     "pad_value": -2,
+    "batch_size": 64,
     "modality": "rna",
     "dropout": 0.2,
-    "GEPC": True,
     "DSBN": True,
     "n_input_bins": 51,
-    "ecs_threshold": 0.8,
-    "explicit_zero_prob": True,
-    "use_fast_transformer": False,
-    "pre_norm": False,
-    "batch_size": 64,
-    "output_compression": None
 }
 ## VIASH END
 
@@ -71,16 +65,16 @@ def setup_logger():
 input_adata = mdata.mod[par["modality"]]
 adata = input_adata.copy()
 
-all_gene_ids = adata.obsm[par["input_obsm_gene_tokens"]]
-all_values = adata.obsm[par["input_obsm_tokenized_values"]]
-padding_mask = adata.obsm[par["input_obsm_padding_mask"]]
+all_gene_ids = adata.obsm[par["obsm_gene_tokens"]]
+all_values = adata.obsm[par["obsm_tokenized_values"]]
+padding_mask = adata.obsm[par["obsm_padding_mask"]]
 
 # Fetch batch ids for domain-specific batch normalization
 if par["DSBN"]:
-    if not par["input_obs_batch_label"]:
+    if not par["obs_batch_label"]:
         raise ValueError("When DSBN is set to True, you are required to provide batch labels (input_obs_batch_labels).")
     else:
-        batch_id_cats = adata.obs[par["input_obs_batch_label"]].astype("category")
+        batch_id_cats = adata.obs[par["obs_batch_label"]].astype("category")
         batch_id_labels = batch_id_cats.cat.codes.values
         batch_ids = batch_id_labels.tolist()
         batch_ids = np.array(batch_ids)
@@ -92,10 +86,10 @@ def setup_logger():
 special_tokens = [pad_token, "<cls>", "<eoc>"]
 
 # Fetching gene names
-if not par["input_var_gene_names"]:
+if not par["var_gene_names"]:
     genes = adata.var.index.astype(str).tolist()
 else:
-    genes = adata.var[par["input_var_gene_names"]].astype(str).tolist()
+    genes = adata.var[par["var_gene_names"]].astype(str).tolist()
 
 logger.info("Loading model, vocab and configs")
 # Model files
@@ -178,6 +172,6 @@ def setup_logger():
 
 logger.info("Writing output data")
 # Write output
-adata.obsm[par["embedding_layer_key"]] = cell_embeddings
+adata.obsm[par["obsm_embeddings"]] = cell_embeddings
 mdata.mod[par["modality"]] = adata
 mdata.write(par["output"], compression=par["output_compression"])

src/scgpt/embedding/test.py

@@ -123,10 +123,10 @@ def test_integration_embedding(run_component, tmp_path):
         "--model", model_file,
         "--model_vocab", vocab_file,
         "--model_config", model_config_file,
-        "--input_obs_batch_label", "sample",
-        "--input_obsm_gene_tokens", "gene_id_tokens",
-        "--input_obsm_tokenized_values", "values_tokenized",
-        "--input_obsm_padding_mask", "padding_mask",
+        "--obs_batch_label", "sample",
+        "--obsm_gene_tokens", "gene_id_tokens",
+        "--obsm_tokenized_values", "values_tokenized",
+        "--obsm_padding_mask", "padding_mask",
         "--output", output_embedding_file
     ])