diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f1c3368..638197c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## NEW FUNCTIONALITY +* Add new method: MMContext, which is a pre-trained multi-modal embedding model, that maps free text annotations and count data into a shared latent space and is trained with contrastive learning. For now included only as a zeroshot method. + +## NEW FUNCTIONALITY + * Add new method: CellMapper, which is a k-NN based approach to map cells across representations and can be used for label projection. Two versions are included here, one based on PCA or CCA embeddings (`linear`) and one based on an scvi embedding (`scvi`) (PR #22) ## MAJOR CHANGES diff --git a/src/methods/mmcontext_zeroshot/config.vsh.yaml b/src/methods/mmcontext_zeroshot/config.vsh.yaml new file mode 100644 index 00000000..ae2444c9 --- /dev/null +++ b/src/methods/mmcontext_zeroshot/config.vsh.yaml @@ -0,0 +1,76 @@ +__merge__: ../../api/base_method.yaml + +name: mmcontext_zeroshot +label: mmcontext (zero shot) +summary: "Zero-shot reference mapping using a pretrained mmcontext model." +description: | + mmcontext is a pretrained model for single-cell context modeling. In this zero-shot + variant, a pretrained mmcontext SentenceTransformer model is used to embed cells and + assign cell-type labels in a zero-shot fashion, given a set of label names from the + test data. The training data is not used in this method. + +references: + doi: 10.64898/2025.12.08.692934 + +links: + documentation: https://github.com/mengerj/mmcontext/blob/main/README.md + repository: https://github.com/mengerj/mmcontext + +info: + preferred_normalization: counts + +arguments: + - name: --model_id + type: string + description: > + Hugging Face model ID for the pretrained mmcontext SentenceTransformer model. + If not provided, a default model ID will be used. + required: false + - name: --batch_size + type: integer + default: 64 + description: Batch size used for mmcontext inference. + - name: --cache_dir + type: string + description: > + Directory used by mmcontext and adata_hf_datasets to cache intermediate + files and model artifacts. + required: true + - name: --initial_embedding + type: string + default: gs10k + description: > + Name of the initial embedding method used by InitialEmbedder (e.g. gs10k). + - name: --n_top + type: integer + default: 5 + description: > + Number of top candidate labels for mmcontext to consider when annotating + cells; the best scoring label will be used as the prediction. + - name: --force_model + type: boolean + default: false + description: > + Skip the sanity-check that `--initial_embedding` is contained within `--model_id`. + Only use this if you are sure the initial embedder matches the one used to train + the mmcontext model checkpoint. + +resources: + - type: python_script + path: script.py + - path: /src/utils/exit_codes.py + +engines: + - type: docker + image: openproblems/base_python:1 + setup: + - type: docker + run: | + pip install "mmcontext @ git+https://github.com/mengerj/mmcontext.git@main" + +runners: + - type: executable + - type: nextflow + directives: + label: [midtime, midmem, midcpu] + diff --git a/src/methods/mmcontext_zeroshot/script.py b/src/methods/mmcontext_zeroshot/script.py new file mode 100644 index 00000000..6d36b91c --- /dev/null +++ b/src/methods/mmcontext_zeroshot/script.py @@ -0,0 +1,175 @@ +import sys +import pathlib +import time + +import anndata as ad +import numpy as np +import scanpy as sc +import torch +from adata_hf_datasets import InitialEmbedder +from adata_hf_datasets.dataset import AnnDataSetConstructor +from mmcontext.embed.model_utils import prepare_model_and_embed +from mmcontext.eval import OmicsQueryAnnotator +from sentence_transformers import SentenceTransformer + +## VIASH START +par = { + "input_train": "resources_test/task_label_projection/cxg_immune_cell_atlas/train.h5ad", + "input_test": "resources_test/task_label_projection/cxg_immune_cell_atlas/test.h5ad", + "output": "output.h5ad", + # mmcontext-specific parameters + "model_id": None, + "batch_size": 64, + "cache_dir": "mmcontext_cache", + "initial_embedding": "gs10k", + "n_top": 5, + "force_model": False, +} +meta = { + "name": "mmcontext_zeroshot", + # "resources_dir" will be injected by Viash when built; default here for local testing + "resources_dir": "target/executable/methods/mmcontext_zeroshot", +} +## VIASH END + +sys.path.append(meta["resources_dir"]) + + + +def compute_mmcontext_embeddings( + adata_path: str, + adata: ad.AnnData, + cache_dir: str, + model_id: str, + initial_embedding: str, +) -> np.ndarray: + """ + Compute mmcontext embeddings for all cells in `adata` using the same + pipeline as your external script, adapted for normalized-then-log data. + """ + cache_dir_path = pathlib.Path(cache_dir) + cache_dir_path.mkdir(parents=True, exist_ok=True) + + # Initial embeddings (e.g. gs10k) + emb_key = f"X_{initial_embedding}" + print("=== Computing initial embeddings ===", flush=True) + embedder = InitialEmbedder(method=initial_embedding) + embedder.prepare(adata) + embeddings = embedder.embed(adata) + adata.obsm[emb_key] = embeddings + + # Persist processed AnnData so AnnDataSetConstructor can link to it + processed_path = adata_path.replace(".h5ad", "_pp.h5ad") + print(f"=== Writing processed AnnData to {processed_path} ===", flush=True) + adata.write_h5ad(processed_path) + + print("=== Building AnnData dataset for mmcontext ===", flush=True) + adata.obs["sample_idx"] = adata.obs.index + constructor = AnnDataSetConstructor(dataset_format="single") + constructor.add_anndata( + adata, + sentence_keys=["sample_idx"], + adata_link=processed_path, + ) + ds = constructor.get_dataset() + + print("=== Loading MMContext model ===", flush=True) + model = SentenceTransformer(model_id, trust_remote_code=True) + + print("=== Computing MMContext embeddings ===", flush=True) + emb_df, _ = prepare_model_and_embed( + model, + ds, + index_col="sample_idx", + adata_download_dir=str(cache_dir_path), + main_col="cell_sentence_1", + layer_key=emb_key, + ) + + # Ensure embeddings align with adata.obs index order + emb_df_sorted = emb_df.set_index("sample_idx").reindex( + list(adata.obs.index) + ).reset_index() + embedding_matrix = np.vstack(emb_df_sorted["embedding"].to_numpy()) + adata.obsm["mmcontext_emb"] = embedding_matrix + return embedding_matrix + + +print("Reading input files", flush=True) +input_train = ad.read_h5ad(par["input_train"]) +input_test = ad.read_h5ad(par["input_test"]) + +# Example organism check; adjust or remove depending on mmcontext’s capabilities. +#if input_train.uns.get("dataset_organism") not in (None, "homo_sapiens"): +# exit_non_applicable( +# "mmcontext_zeroshot currently only supports human data " +# f'(dataset_organism == "{input_train.uns["dataset_organism"]}")' +# ) + +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f"Device: '{device}'", flush=True) + +print("Preparing test data for mmcontext", flush=True) +if "counts" in input_test.layers: + input_test.X = input_test.layers["counts"] + +# Updated preprocessing: normalize first, then log-transform +print("=== Preprocessing test data (normalize -> log1p) ===", flush=True) +sc.pp.normalize_total(input_test) +sc.pp.log1p(input_test) + +# Determine model ID +default_model_id = "jo-mengr/mmcontext-pubmedbert-gs10k" +model_id = par.get("model_id") or default_model_id +print(f"Using mmcontext model ID: {model_id}", flush=True) + +# Quick safety check: initial embedding should match the checkpoint naming. +if (not par.get("force_model", False)) and (par["initial_embedding"] not in model_id): + raise ValueError( + "The configured initial embedding method does not appear to match the mmcontext model.\n" + f"- initial_embedding: '{par['initial_embedding']}'\n" + f"- model_id: '{model_id}'\n\n" + "The initial embedding method needs to match the one that was used to train the model.\n" + "If you are absolutely sure they match, re-run with `--force_model true` to skip this check." + ) + +start_time = time.time() + +embeddings = compute_mmcontext_embeddings( + adata_path=par["input_test"], + adata=input_test, + cache_dir=par["cache_dir"], + model_id=model_id, + initial_embedding=par["initial_embedding"], +) + +print("=== Creating annotator and predicting labels ===", flush=True) +annotator_model = SentenceTransformer(model_id, trust_remote_code=True) +annotator = OmicsQueryAnnotator(annotator_model) + +# Use unique labels from the test data as the zero-shot label set +label_space = input_test.obs["label"].unique() +annotator.annotate_omics_data( + input_test, + labels=label_space, + emb_key="mmcontext_emb", + n_top=par["n_top"], +) + +# OmicsQueryAnnotator writes the best label to obs["best_label"] +input_test.obs["label_pred"] = input_test.obs["best_label"] +input_test.obs["label_pred"] = input_test.obs["label_pred"].astype("category") + +elapsed_time = time.time() - start_time +print(f"mmcontext zero-shot prediction completed in {elapsed_time:.2f} seconds", flush=True) + +print("Write output AnnData to file", flush=True) +output = ad.AnnData( + obs=input_test.obs[["label_pred"]], + uns={ + "method_id": meta["name"], + "dataset_id": input_test.uns["dataset_id"], + "normalization_id": input_test.uns["normalization_id"], + }, +) +output.write_h5ad(par["output"], compression="gzip")