From 06233ffff38fdcfa8d94b748c5a3c6a5c3f101fb Mon Sep 17 00:00:00 2001 From: Luke Zappia Date: Mon, 9 Dec 2024 09:43:05 +0100 Subject: [PATCH 1/4] Add model path argument to scGPT --- src/methods/scgpt/config.vsh.yaml | 6 +++ src/methods/scgpt/script.py | 62 ++++++++++++++++++++++++------- 2 files changed, 55 insertions(+), 13 deletions(-) diff --git a/src/methods/scgpt/config.vsh.yaml b/src/methods/scgpt/config.vsh.yaml index a6a72835..b22bf79f 100644 --- a/src/methods/scgpt/config.vsh.yaml +++ b/src/methods/scgpt/config.vsh.yaml @@ -29,6 +29,12 @@ arguments: description: String giving the scGPT model to use choices: ["scGPT_human", "scGPT_CP"] default: "scGPT_human" + - name: --model_path + type: file + description: | + Path to the directory containing the specified scGPT model or a + .zip/.tar.gz archive to extract. If not given the model will be downloaded. + required: false - name: --n_hvg type: integer default: 3000 diff --git a/src/methods/scgpt/script.py b/src/methods/scgpt/script.py index 37297958..356a0e88 100644 --- a/src/methods/scgpt/script.py +++ b/src/methods/scgpt/script.py @@ -1,5 +1,8 @@ +import os import sys +import tarfile import tempfile +import zipfile import anndata as ad import gdown @@ -43,23 +46,55 @@ print(adata, flush=True) -print(f"\n>>> Downloading '{par['model']}' model...", flush=True) -model_drive_ids = { - "scGPT_human": "1oWh_-ZRdhtoGQ2Fw24HP41FgLoomVo-y", - "scGPT_CP": "1_GROJTzXiAV8HB4imruOTk6PEGuNOcgB", -} -drive_path = f"https://drive.google.com/drive/folders/{model_drive_ids[par['model']]}" -model_dir = tempfile.TemporaryDirectory() -print(f"Downloading from '{drive_path}'", flush=True) -gdown.download_folder(drive_path, output=model_dir.name, quiet=True) -print(f"Model directory: '{model_dir.name}'", flush=True) +if par["model_path"] is None: + print(f"\n>>> Downloading '{par['model']}' model...", flush=True) + model_drive_ids = { + "scGPT_human": "1oWh_-ZRdhtoGQ2Fw24HP41FgLoomVo-y", + "scGPT_CP": "1_GROJTzXiAV8HB4imruOTk6PEGuNOcgB", + } + drive_path = ( + f"https://drive.google.com/drive/folders/{model_drive_ids[par['model']]}" + ) + model_temp = tempfile.TemporaryDirectory() + model_dir = model_temp.name + print(f"Downloading from '{drive_path}'", flush=True) + gdown.download_folder(drive_path, output=model_dir, quiet=True) +else: + if os.path.isdir(par["model_path"]): + print(f"\n>>> Using model directory...", flush=True) + model_temp = None + model_dir = par["model_path"] + else: + model_temp = tempfile.TemporaryDirectory() + model_dir = model_temp.name + + if zipfile.is_zipfile(par["model_path"]): + print(f"\n>>> Extracting model from .zip...", flush=True) + print(f".zip path: '{par['model_path']}'", flush=True) + with zipfile.ZipFile(par["model_path"], "r") as zip_file: + zip_file.extractall(model_dir) + elif tarfile.is_tarfile(par["model_path"]) and par["model_path"].endswith( + ".tar.gz" + ): + print(f"\n>>> Extracting model from .tar.gz...", flush=True) + print(f".tar.gz path: '{par['model_path']}'", flush=True) + with tarfile.open(par["model_path"], "r:gz") as tar_file: + tar_file.extractall(model_dir) + model_dir = os.path.join(model_dir, os.listdir(model_dir)[0]) + else: + raise ValueError( + f"The 'model_path' argument should be a directory a .zip file or a .tar.gz file" + ) + +print(f"Model directory: '{model_dir}'", flush=True) +print(os.listdir(model_dir), flush=True) print("\n>>> Embedding data...", flush=True) device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Device: '{device}'", flush=True) embedded = scgpt.tasks.embed_data( adata, - model_dir.name, + model_dir, gene_col="feature_name", batch_size=64, use_fast_transformer=False, # Disable fast-attn as not installed @@ -86,7 +121,8 @@ print(f"Output H5AD file: '{par['output']}'", flush=True) output.write_h5ad(par["output"], compression="gzip") -print("\n>>> Cleaning up temporary directories...", flush=True) -model_dir.cleanup() +if model_temp is not None: + print("\n>>> Cleaning up temporary directories...", flush=True) + model_temp.cleanup() print("\n>>> Done!", flush=True) From 106b7ae75682b4d6af7d418c62379a06795d0791 Mon Sep 17 00:00:00 2001 From: Luke Zappia Date: Mon, 9 Dec 2024 09:47:41 +0100 Subject: [PATCH 2/4] Use cached scGPT model in benchmark workflow --- src/workflows/run_benchmark/main.nf | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/workflows/run_benchmark/main.nf b/src/workflows/run_benchmark/main.nf index afcb968c..52522162 100644 --- a/src/workflows/run_benchmark/main.nf +++ b/src/workflows/run_benchmark/main.nf @@ -29,7 +29,9 @@ methods = [ scalex, scanorama, scanvi, - scgpt, + scgpt.run( + args: [model_path: file("s3://openproblems-work/cache/scGPT_human.zip")] + ), scimilarity.run( args: [model: file("s3://openproblems-work/cache/scimilarity-model_v1.1.tar.gz")] ), From 0ba1508c969bad324b67ee05f308f37ea090d865 Mon Sep 17 00:00:00 2001 From: Luke Zappia Date: Mon, 9 Dec 2024 09:59:51 +0100 Subject: [PATCH 3/4] Make scGPT inherit from base method --- src/methods/scgpt/config.vsh.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/methods/scgpt/config.vsh.yaml b/src/methods/scgpt/config.vsh.yaml index b22bf79f..83f226ad 100644 --- a/src/methods/scgpt/config.vsh.yaml +++ b/src/methods/scgpt/config.vsh.yaml @@ -1,4 +1,4 @@ -__merge__: ../../api/comp_method.yaml +__merge__: ../../api/base_method.yaml name: scgpt label: scGPT From 869d09d5e45fccc03b33b2b61a17936b867003e2 Mon Sep 17 00:00:00 2001 From: Luke Zappia Date: Tue, 10 Dec 2024 13:31:20 +0100 Subject: [PATCH 4/4] Swap scGPT model argument names --- src/methods/scgpt/config.vsh.yaml | 11 ++++++----- src/methods/scgpt/script.py | 26 +++++++++++++------------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/src/methods/scgpt/config.vsh.yaml b/src/methods/scgpt/config.vsh.yaml index 83f226ad..263a490b 100644 --- a/src/methods/scgpt/config.vsh.yaml +++ b/src/methods/scgpt/config.vsh.yaml @@ -24,16 +24,17 @@ info: model: "scGPT_CP" arguments: - - name: --model + - name: --model_name type: string - description: String giving the scGPT model to use + description: String giving the name of the scGPT model to use choices: ["scGPT_human", "scGPT_CP"] default: "scGPT_human" - - name: --model_path + - name: --model type: file description: | - Path to the directory containing the specified scGPT model or a - .zip/.tar.gz archive to extract. If not given the model will be downloaded. + Path to the directory containing the scGPT model specified by model_name + or a .zip/.tar.gz archive to extract. If not given the model will be + downloaded. required: false - name: --n_hvg type: integer diff --git a/src/methods/scgpt/script.py b/src/methods/scgpt/script.py index 356a0e88..dda48777 100644 --- a/src/methods/scgpt/script.py +++ b/src/methods/scgpt/script.py @@ -15,6 +15,7 @@ par = { "input": "resources_test/.../input.h5ad", "output": "output.h5ad", + "model_name": "scGPT_human", "model": "scGPT_human", "n_hvg": 3000, } @@ -46,48 +47,47 @@ print(adata, flush=True) -if par["model_path"] is None: - print(f"\n>>> Downloading '{par['model']}' model...", flush=True) +if par["model"] is None: + print(f"\n>>> Downloading '{par['model_name']}' model...", flush=True) model_drive_ids = { "scGPT_human": "1oWh_-ZRdhtoGQ2Fw24HP41FgLoomVo-y", "scGPT_CP": "1_GROJTzXiAV8HB4imruOTk6PEGuNOcgB", } drive_path = ( - f"https://drive.google.com/drive/folders/{model_drive_ids[par['model']]}" + f"https://drive.google.com/drive/folders/{model_drive_ids[par['model_name']]}" ) model_temp = tempfile.TemporaryDirectory() model_dir = model_temp.name print(f"Downloading from '{drive_path}'", flush=True) gdown.download_folder(drive_path, output=model_dir, quiet=True) else: - if os.path.isdir(par["model_path"]): + if os.path.isdir(par["model"]): print(f"\n>>> Using model directory...", flush=True) model_temp = None - model_dir = par["model_path"] + model_dir = par["model"] else: model_temp = tempfile.TemporaryDirectory() model_dir = model_temp.name - if zipfile.is_zipfile(par["model_path"]): + if zipfile.is_zipfile(par["model"]): print(f"\n>>> Extracting model from .zip...", flush=True) - print(f".zip path: '{par['model_path']}'", flush=True) - with zipfile.ZipFile(par["model_path"], "r") as zip_file: + print(f".zip path: '{par['model']}'", flush=True) + with zipfile.ZipFile(par["model"], "r") as zip_file: zip_file.extractall(model_dir) - elif tarfile.is_tarfile(par["model_path"]) and par["model_path"].endswith( + elif tarfile.is_tarfile(par["model"]) and par["model"].endswith( ".tar.gz" ): print(f"\n>>> Extracting model from .tar.gz...", flush=True) - print(f".tar.gz path: '{par['model_path']}'", flush=True) - with tarfile.open(par["model_path"], "r:gz") as tar_file: + print(f".tar.gz path: '{par['model']}'", flush=True) + with tarfile.open(par["model"], "r:gz") as tar_file: tar_file.extractall(model_dir) model_dir = os.path.join(model_dir, os.listdir(model_dir)[0]) else: raise ValueError( - f"The 'model_path' argument should be a directory a .zip file or a .tar.gz file" + f"The 'model' argument should be a directory a .zip file or a .tar.gz file" ) print(f"Model directory: '{model_dir}'", flush=True) -print(os.listdir(model_dir), flush=True) print("\n>>> Embedding data...", flush=True) device = "cuda" if torch.cuda.is_available() else "cpu"