script.py

import os
import sys
import tarfile
import tempfile
import zipfile

import anndata as ad
import torch
from CellPLM.pipeline.cell_embedding import CellEmbeddingPipeline
from CellPLM.utils import set_seed

## VIASH START
# Note: this section is auto-generated by viash at runtime. To edit it, make changes
# in config.vsh.yaml and then run `viash config inject config.vsh.yaml`.

par = {
    "input": "resources_test/.../input.h5ad",
    "output": "output.h5ad",
    "model": "20231027_85M",
}
meta = {"name": "cellplm"}
## VIASH END

sys.path.append(meta["resources_dir"])
from exit_codes import exit_non_applicable
from read_anndata_partial import read_anndata

set_seed(24)
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    import warnings

    warnings.warn("Loading CellPLM models requires a GPU, this run will fail")

print("\n>>> Reading input files...", flush=True)
print(f"Input H5AD file: '{par['input']}'", flush=True)
adata = read_anndata(par["input"], X="layers/counts", obs="obs", var="var", uns="uns")

if adata.uns["dataset_organism"] != "homo_sapiens":
    exit_non_applicable(
        f"CellPLM can only be used with human data "
        f'(dataset_organism == "{adata.uns["dataset_organism"]}")'
    )

print(adata, flush=True)

print("\n>>> Getting model files...", flush=True)
# Available from https://www.dropbox.com/scl/fo/i5rmxgtqzg7iykt2e9uqm/h/ckpt?dl=0&subfolder_nav_tracking=1
if os.path.isdir(par["model"]):
    model_temp = None
    model_dir = par["model"]
else:
    model_temp = tempfile.TemporaryDirectory()
    model_dir = model_temp.name

    if zipfile.is_zipfile(par["model"]):
        print("Extracting CellPLM models from .zip...", flush=True)
        with zipfile.ZipFile(par["model"], "r") as zip_file:
            zip_file.extractall(model_dir)
    elif tarfile.is_tarfile(par["model"]) and par["model"].endswith(".tar.gz"):
        print("Extracting CellPLM models from .tar.gz...", 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(
            "The 'model' argument should be a directory a .zip file or a .tar.gz file"
        )

print(f"Model directory: '{model_dir}'", flush=True)

print("\n>>> Creating embedding model pipeline...", flush=True)
pipeline = CellEmbeddingPipeline(
    pretrain_prefix=par["model_name"], pretrain_directory=model_dir
)

print("\n>>> Embedding data...", flush=True)
embedding = pipeline.predict(adata, device=device)
embedding = embedding.cpu().numpy()

print("\n>>> Storing output...", flush=True)
output = ad.AnnData(
    obs=adata.obs[[]],
    var=adata.var[[]],
    obsm={
        "X_emb": embedding,
    },
    uns={
        "dataset_id": adata.uns["dataset_id"],
        "normalization_id": adata.uns["normalization_id"],
        "method_id": meta["name"],
    },
)
print(output)

print("\n>>> Writing output to file...", flush=True)
print(f"Output H5AD file: '{par['output']}'", flush=True)
output.write_h5ad(par["output"], compression="gzip")

if model_temp is not None:
    print("\n>>> Cleaning up temporary directories...", flush=True)
    model_temp.cleanup()

print("\n>>> Done!", flush=True)