MGM2: A Foundation Model for Microbial Communities

MGM2 converts a microbial community into contextualized microbial-token representations and one fixed-size community embedding. It combines 1,536-dimensional NTv3 microbial identity vectors with relative-abundance signals in an abundance-aware Transformer.

Included Models

Each model directory contains:

• model.ckpt: an inference-only Lightning checkpoint derived from the released last.ckpt, without optimizer state or training hyperparameters;
• config.json: the matching architecture and abundance configuration.

Shared identity resources are stored once under embeddings/:

• reference_ntv3_embeddings.h5: NTv3 embeddings keyed by reference OTU ID;
• taxonomy_fallback.h5: rank-aware taxonomy prototypes for inputs without representative sequences.

Download

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="<HUGGINGFACE_REPO_ID>",
    local_dir="MGM2-release",
)

Replace <HUGGINGFACE_REPO_ID> with the ID of this model repository.

Obtain MGM2 Embeddings

The inference utilities are maintained in the MGM2 source repository.

Input with OTU or ASV sequences

The FASTA record IDs must match the feature IDs in the abundance table.

python ntv3_embedding.py \
  --fasta_path communities.fasta \
  --model_path /path/to/NTv3_650M_pre \
  --output_path community_ntv3_embeddings.h5 \
  --gpu_ids 0

python src/data/build_dataset.py \
  --otu_table community_table.csv \
  --format csv \
  --ntv3_embedding community_ntv3_embeddings.h5 \
  --top_k_otus 768 \
  --output community_input.pkl

Input with species names

Species names can be bare names such as Bacteroides fragilis, rank-prefixed names such as s__Bacteroides fragilis, or complete taxonomy paths.

python src/data/build_dataset.py \
  --otu_table community_species.csv \
  --format csv \
  --ntv3_embedding MGM2-release/embeddings/reference_ntv3_embeddings.h5 \
  --fallback_embedding MGM2-release/embeddings/taxonomy_fallback.h5 \
  --resolved_embedding community_resolved_embeddings.h5 \
  --top_k_otus 768 \
  --output community_input.pkl

Name matching is rank aware and uses only unambiguous names. The resolved HDF5 records embedding_source and matched_taxon_key for inspection.

Extract the community representation

python scripts/extract_community_embeddings.py \
  --model_dir MGM2-release/models/small \
  --data_file community_input.pkl \
  --output community_embeddings.h5 \
  --embedding_strategy cls \
  --device cuda:0

cls is the recommended general-purpose representation. mean_pool exports the mean of contextualized microbial tokens. qwen3_proj is available because the released checkpoints include the alignment head.

Output

HDF5 output contains sample_ids and an embeddings matrix with shape [number_of_communities, hidden_size]. The extractor also supports NPZ and CSV output by changing the output filename extension.

Citation

If you use MGM2 in your research, please cite the accompanying paper. Complete citation metadata will be added with the model release.