Add waypoint skill
This commit is contained in:
Timothy Kassis
@@ -55,7 +55,7 @@ These skills enable your AI agent to seamlessly work with specialized scientific
|
||||
|
||||
## 📦 What's Included
|
||||
|
||||
This repository provides **162 scientific and research skills** organized into the following categories:
|
||||
This repository provides **163 scientific and research skills** organized into the following categories:
|
||||
|
||||
- **100+ Scientific & Financial Databases** - A unified database-lookup skill provides deterministic, provenance-rich access to 78 public databases (PubChem, ChEMBL, UniProt, COSMIC, ClinicalTrials.gov, FRED, USPTO, and more), plus dedicated skills for DepMap, Imaging Data Commons, PrimeKG, NCATS ARAX, U.S. Treasury Fiscal Data, Hugging Science, OneKGPd, and Genomic Intelligence. Multi-database packages like BioServices (~40 bioinformatics services), BioPython (39 NCBI sub-databases via Entrez), and gget (20+ genomics databases) add further coverage
|
||||
- **70+ Optimized Python Package Skills** - Explicitly defined, version-aware workflows for RDKit, Scanpy, PyTorch Lightning, scikit-learn, PyTDC, PathML, pydicom, NeuroKit2, PufferLib, QuTiP, GeoPandas, pymatgen, BioPython, Qiskit, Molecular Dynamics (OpenMM/MDAnalysis), and others. The agent can still use *any* Python package; these skills provide stronger, safer guidance for the packages listed
|
||||
@@ -448,13 +448,13 @@ networks, and search GEO for similar patterns.
|
||||
|
||||
## 📚 Available Skills
|
||||
|
||||
This repository contains **162 scientific and research skills** organized across multiple domains. Each skill provides comprehensive documentation, code examples, and best practices for working with scientific libraries, databases, and tools.
|
||||
This repository contains **163 scientific and research skills** organized across multiple domains. Each skill provides comprehensive documentation, code examples, and best practices for working with scientific libraries, databases, and tools.
|
||||
|
||||
### Skill Categories
|
||||
|
||||
> **Note:** The Python package and integration skills listed below are *explicitly defined* skills — curated with documentation, examples, and best practices for stronger, more reliable performance. They are not a ceiling: the agent can install and use *any* Python package or call *any* API, even without a dedicated skill. The skills listed simply make common workflows faster and more dependable.
|
||||
|
||||
#### 🧬 **Bioinformatics & Genomics** (26 skills)
|
||||
#### 🧬 **Bioinformatics & Genomics** (27 skills)
|
||||
- RNA-seq pipelines: Bulk RNA-seq (end-to-end FASTQ -> counts -> DE -> enrichment orchestrator)
|
||||
- Sequence analysis: BioPython, pysam, scikit-bio, BioServices
|
||||
- Single-cell analysis: Scanpy, AnnData, scvi-tools, scVelo (RNA velocity), Arboreto, Cellxgene Census
|
||||
@@ -464,6 +464,7 @@ This repository contains **162 scientific and research skills** organized across
|
||||
- Differential expression: PyDESeq2
|
||||
- Functional enrichment: Pathway Enrichment (ORA, GSEA/preranked, ssGSEA via gseapy + g:Profiler; GO, KEGG, Reactome, WikiPathways, MSigDB)
|
||||
- Phylogenetics: ETE Toolkit, Phylogenetics (MAFFT, IQ-TREE 2, FastTree)
|
||||
- Microbiome foundation models: Waypoint (Outpost Bio's open Waypoint-6m/45m/170m checkpoints, the Atlas 539k-sample MGnify pretraining corpus, and the eight-task Compass benchmark — embedding, fine-tuning, benchmarking, and pretraining on taxonomic abundance profiles, with MetaPhlAn/Kraken2/QIIME 2 conversion)
|
||||
|
||||
#### 🧪 **Cheminformatics & Drug Discovery** (10 skills)
|
||||
- Molecular manipulation: RDKit, Datamol, Molfeat
|
||||
|
||||
BIN
docs/images/waypoint-bio.png
Normal file
BIN
docs/images/waypoint-bio.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
273
skills/waypoint-bio/SKILL.md
Normal file
273
skills/waypoint-bio/SKILL.md
Normal file
@@ -0,0 +1,273 @@
|
||||
---
|
||||
name: waypoint-bio
|
||||
description: Use when working with Outpost Bio's open microbiome foundation models - the Waypoint checkpoints (Waypoint-6m, Waypoint-45m, Waypoint-170m), the Atlas pretraining corpus, the Compass eight-task benchmark, or the `waypoint` CLI from the `waypoint-bio` package. Covers embedding microbiome samples, fine-tuning on taxonomic abundance data, benchmarking a checkpoint on Compass, pretraining a GPT-2 model on taxonomic abundance profiles, and converting MetaPhlAn, Kraken2, QIIME 2, or MGnify abundance tables into waypoint format.
|
||||
license: MIT
|
||||
compatibility: Requires Python 3.10+ with `waypoint-bio` (pulls torch, transformers, datasets, peft, scikit-learn). Needs network access and a Hugging Face token with access granted to the gated outpost-bio repos. A GPU is strongly recommended for pretraining and benchmarking.
|
||||
metadata:
|
||||
version: "1.0"
|
||||
skill-author: K-Dense Inc.
|
||||
upstream-version: "waypoint-bio 1.0.2 (PyPI); GitHub main 1.0.4"
|
||||
last-reviewed: "2026-08-17"
|
||||
openclaw:
|
||||
primaryEnv: HF_TOKEN
|
||||
envVars:
|
||||
- name: HF_TOKEN
|
||||
required: true
|
||||
description: Hugging Face read token with access to the gated outpost-bio/Waypoint-*, outpost-bio/Atlas, and outpost-bio/Compass repos.
|
||||
---
|
||||
|
||||
# Waypoint: Outpost Bio's Open Microbiome Foundation Models
|
||||
|
||||
## Overview
|
||||
|
||||
Outpost Bio open-sourced three artefacts under Apache 2.0, described in
|
||||
[Treloar et al., bioRxiv 2026.05.02.722381](https://www.biorxiv.org/content/10.64898/2026.05.02.722381v2):
|
||||
|
||||
| Artefact | What it is | Hugging Face |
|
||||
| --- | --- | --- |
|
||||
| **Waypoint** | GPT-2-style causal LMs over taxonomic tokens, 6M–170M params | `outpost-bio/Waypoint-6m`, `-45m`, `-170m` |
|
||||
| **Atlas** | 539,308 microbiome samples scraped from MGnify (485,377 pretrain / 53,931 benchmark) | `outpost-bio/Atlas` |
|
||||
| **Compass** | Eight downstream tasks over four studies | `outpost-bio/Compass` |
|
||||
|
||||
The unifying idea: a microbiome sample is a *sentence*. Each taxon is one token, tokens are ordered
|
||||
by descending abundance z-score, and the model is trained with next-token prediction. A pretrained
|
||||
checkpoint then supplies sample-level embeddings or a fine-tuning backbone for prediction tasks.
|
||||
|
||||
All of it is driven by one CLI, `waypoint`, with five subcommands: `prepare-dataset`, `embed`,
|
||||
`finetune`, `benchmark`, `pretrain`.
|
||||
|
||||
## When to use
|
||||
|
||||
- Embedding 16S/shotgun taxonomic profiles into fixed-size vectors for clustering, visualisation, or
|
||||
a downstream classifier.
|
||||
- Fine-tuning a Waypoint checkpoint to predict a phenotype, treatment, or continuous readout from
|
||||
community composition.
|
||||
- Scoring your own microbiome model against Compass so the number is comparable to the paper.
|
||||
- Pretraining a taxonomic language model on Atlas or on your own corpus.
|
||||
- Converting profiler output (MetaPhlAn, Kraken2/Bracken, QIIME 2, MGnify TSVs) into the input format
|
||||
these tools expect.
|
||||
|
||||
**Do not reach for this** when you have fewer than ~1,000 labelled samples — see
|
||||
[Scientific caveats](#scientific-caveats). A random forest on relative abundances is the better tool
|
||||
there, and the paper says so.
|
||||
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
pip install waypoint-bio # installs the `waypoint` command
|
||||
```
|
||||
|
||||
Atlas, Compass, and every Waypoint checkpoint are **gated**. Access is auto-approved, but you must
|
||||
click through once per repo and then authenticate:
|
||||
|
||||
1. Request access on each repo page you need: [Waypoint-6m](https://huggingface.co/outpost-bio/Waypoint-6m),
|
||||
[Waypoint-45m](https://huggingface.co/outpost-bio/Waypoint-45m),
|
||||
[Waypoint-170m](https://huggingface.co/outpost-bio/Waypoint-170m),
|
||||
[Atlas](https://huggingface.co/datasets/outpost-bio/Atlas),
|
||||
[Compass](https://huggingface.co/datasets/outpost-bio/Compass).
|
||||
2. Authenticate locally:
|
||||
|
||||
```bash
|
||||
hf auth login # or: export HF_TOKEN=hf_...
|
||||
```
|
||||
|
||||
A 401/403 from any subcommand almost always means access was never requested on that specific repo —
|
||||
a token alone is not enough. Use a read-scoped token. The tokenizer loads via
|
||||
`trust_remote_code=True`, so pin a `revision` if you need the remote code fixed across runs.
|
||||
|
||||
## The waypoint data format
|
||||
|
||||
Everything except `prepare-dataset` consumes **waypoint format**: a `.parquet` / `.csv` / `.tsv`
|
||||
whose rows are samples, with two aligned list-columns plus any label columns you need.
|
||||
|
||||
| Column | Type | Notes |
|
||||
| --- | --- | --- |
|
||||
| `Taxa` | `list[str]` | Full lineage strings, `;`-separated: `k__Bacteria; p__Firmicutes; ...; g__Lactobacillus` |
|
||||
| `Relative Abundances` | `list[float]` | Same length as `Taxa`, same order |
|
||||
| *(any)* | scalar | Targets, covariates, or a `Split` column |
|
||||
|
||||
Prefer parquet. CSV/TSV stores the lists as `repr` strings and round-trips through `ast.literal_eval`.
|
||||
|
||||
**Give full lineages, not bare names.** The tokenizer extracts the genus segment (`g__`) from each
|
||||
lineage and falls back to the most specific higher rank when genus is missing. Bare names disable
|
||||
that fallback entirely.
|
||||
|
||||
## Workflow
|
||||
|
||||
### 1. Get your data into waypoint format
|
||||
|
||||
If you already have a sample × taxa (or taxa × sample) abundance matrix with lineage labels:
|
||||
|
||||
```bash
|
||||
waypoint prepare-dataset \
|
||||
--input abundance_matrix.tsv \
|
||||
--metadata sample_labels.csv \
|
||||
--output dataset.parquet
|
||||
```
|
||||
|
||||
Orientation is auto-detected from the first column header (`taxonomy`, `lineage`, `taxon`, `otu`,
|
||||
`#otu id` ⇒ taxa-as-rows); override with `--orientation`. Rows are normalised to sum to 1 unless you
|
||||
pass `--no_normalize`, and zeros are dropped unless you pass `--keep_zeros`.
|
||||
|
||||
`prepare-dataset` cannot read profiler output directly — MetaPhlAn uses `|` separators, Kraken2
|
||||
reports encode the hierarchy as indentation, and QIIME 2/SILVA prefixes the domain `d__` instead of
|
||||
`k__` (which the tokenizer silently ignores). Use the bundled converter for those:
|
||||
|
||||
```bash
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input merged_metaphlan.tsv --format metaphlan \
|
||||
--output dataset.parquet
|
||||
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input reports/*.kreport --format kraken \
|
||||
--output dataset.parquet
|
||||
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input feature-table.tsv --format qiime2 \
|
||||
--output dataset.parquet
|
||||
```
|
||||
|
||||
See `references/data-preparation.md` for every input layout, rank handling, and the `d__`/`|` gotchas.
|
||||
|
||||
### 2. Check vocabulary coverage before anything else
|
||||
|
||||
Waypoint's vocabulary is fixed at pretraining time from Atlas. Taxa absent from it become `<unk>` and
|
||||
are **silently dropped** by `waypoint embed`; the paper names this as the models' main limitation. A
|
||||
sample whose taxa are all out-of-vocabulary yields a degenerate `[BOS][EOS]` embedding.
|
||||
|
||||
```bash
|
||||
python scripts/vocab_coverage.py --model outpost-bio/Waypoint-6m --data dataset.parquet
|
||||
```
|
||||
|
||||
It reports per-sample and abundance-weighted coverage and flags samples below a threshold. Treat
|
||||
median abundance-weighted coverage under ~0.8 as a reason to re-examine your taxonomy labels before
|
||||
trusting any downstream number.
|
||||
|
||||
### 3. Embed samples
|
||||
|
||||
```bash
|
||||
waypoint embed \
|
||||
--model outpost-bio/Waypoint-6m \
|
||||
--data dataset.parquet \
|
||||
--output embeddings.parquet
|
||||
```
|
||||
|
||||
Output is indexed by sample ID with columns `dim_0 … dim_{H-1}` (`H` = 256 for 6m, 512 for 45m,
|
||||
768 for 170m). Defaults: `--pooling last_token`, `--batch_size 32`, `--max_length 512`, device
|
||||
auto-detected (`cuda` → `mps` → `cpu`).
|
||||
|
||||
Keep `--pooling last_token` unless you have a reason to change it: it matches how the checkpoints
|
||||
were pretrained and how `benchmark` and `finetune` pool. `mean` is a reasonable alternative for
|
||||
unsupervised use; `first_token`/`cls_token` return the BOS position and carry little signal in a
|
||||
causal LM.
|
||||
|
||||
### 4. Fine-tune on your labels
|
||||
|
||||
```bash
|
||||
# classification
|
||||
waypoint finetune \
|
||||
--model outpost-bio/Waypoint-45m \
|
||||
--data dataset.parquet \
|
||||
--output_dir outputs/ft_disease \
|
||||
--task_type classification \
|
||||
--target "Disease Status" \
|
||||
--config configs/finetune_classification.yaml
|
||||
|
||||
# regression, with a categorical covariate one-hot appended to the pooled embedding
|
||||
waypoint finetune \
|
||||
--model outpost-bio/Waypoint-45m \
|
||||
--data dataset.parquet \
|
||||
--output_dir outputs/ft_degradation \
|
||||
--task_type regression \
|
||||
--target "Degradation Rate" \
|
||||
--covariate_column Drug \
|
||||
--config configs/finetune_regression.yaml
|
||||
```
|
||||
|
||||
Config paths resolve against the bundled `waypoint_bio/configs/` tree, so `configs/...` works from
|
||||
any directory without cloning.
|
||||
|
||||
Defaults worth overriding for small datasets: `warmup_steps: 1000` (drop to ~50 so warmup finishes
|
||||
before early stopping), `num_epochs: 1` in the shipped configs (raise it — early stopping on
|
||||
validation loss is what actually terminates training), and `use_lora: true` when VRAM is tight
|
||||
(~1% of parameters trained; adapters are merged back before saving, so the checkpoint stays a plain
|
||||
`AutoModel`).
|
||||
|
||||
Splits default to a random 80/10/10. **Set `split_column` to a `Split` column whenever samples are
|
||||
correlated** — repeated measures, one donor sampled over time, technical replicates — or a random
|
||||
split leaks and the test score is meaningless.
|
||||
|
||||
Outputs land in `--output_dir`: `best_model/` (loadable by `embed`/`benchmark`),
|
||||
`test_metrics.json`, `training_log.csv` + `.html`, and `finetune_results.json`.
|
||||
|
||||
### 5. Benchmark on Compass
|
||||
|
||||
```bash
|
||||
waypoint benchmark --model outpost-bio/Waypoint-6m --output_dir outputs/benchmark
|
||||
waypoint benchmark --model outputs/pretrain/best_model --tasks 1 6 --output_dir outputs/smoke
|
||||
```
|
||||
|
||||
Fine-tunes a fresh head per task and writes `benchmark_results.json`. Classification tasks score
|
||||
macro-F1; the one regression task scores R² clamped to [0, 1]; `final_score` is the unweighted mean
|
||||
across tasks. Full task table, metric keys, and result-file schema: `references/compass-benchmark.md`.
|
||||
|
||||
### 6. Pretrain
|
||||
|
||||
```bash
|
||||
waypoint pretrain \
|
||||
--model_config configs/models/gpt2-45m.yaml \
|
||||
--pretrain_config configs/pretraining.yaml \
|
||||
--output_dir outputs/pretrain_45m
|
||||
```
|
||||
|
||||
Downloads Atlas, builds a taxonomic tokenizer from the corpus, computes per-token abundance
|
||||
mean/std for z-score ordering, then trains with next-token prediction and early stopping. Add
|
||||
`--data my_corpus.parquet` to pretrain on your own waypoint-format corpus instead, and
|
||||
`--max_samples N` for a smoke test.
|
||||
|
||||
Nine architectures ship, from `gpt2-6m.yaml` (8 layers, 256 hidden) to `gpt2-170m.yaml` (24 layers,
|
||||
768 hidden); per-head dimension is fixed at 64 throughout. `references/cli-reference.md` has the
|
||||
full table and every config key.
|
||||
|
||||
## Scientific caveats
|
||||
|
||||
These are load-bearing. Ignoring them produces numbers that look fine and mean nothing.
|
||||
|
||||
- **Below ~1,000 labelled examples, Waypoint underperforms a random forest on raw abundances.** The
|
||||
paper's crossover against the RF baseline sits near **10,000** training examples. Fit the baseline
|
||||
first; only adopt the transformer if it wins on your data.
|
||||
- **Out-of-vocabulary taxa are dropped, not flagged.** Every Compass dataset carries some. Run
|
||||
`scripts/vocab_coverage.py` and report the coverage alongside your results.
|
||||
- **45M, not 170M, was the best benchmark model.** Pretraining loss keeps falling with scale, but
|
||||
downstream Compass score does not — start at 6m or 45m and only scale up if it demonstrably helps.
|
||||
- **Genus-level tokenisation is the default**, so species-level distinctions are collapsed. Changing
|
||||
`taxon_rank` requires re-pretraining, not just re-tokenising.
|
||||
- **Compositional data.** Relative abundances are constrained to sum to 1; differences in one taxon
|
||||
induce apparent changes in others. This affects interpretation of any per-taxon attribution.
|
||||
- **Batch and study effects dominate microbiome data.** Atlas spans MGnify pipelines v1.0–v5.0 and
|
||||
four sequencing modalities. Never let a study or run boundary coincide with your label boundary.
|
||||
- **Not a clinical or diagnostic tool.** The model cards state this explicitly.
|
||||
|
||||
## References
|
||||
|
||||
- `references/cli-reference.md` — every subcommand flag, every config key, the model-size table.
|
||||
- `references/compass-benchmark.md` — the eight tasks, filters, metrics, `benchmark_results.json` schema.
|
||||
- `references/data-preparation.md` — waypoint format, profiler conversions, taxonomy string rules.
|
||||
- `references/python-api.md` — using the tokenizer, datasets, heads, and checkpoints from Python.
|
||||
|
||||
## Scripts
|
||||
|
||||
- `scripts/profiler_to_waypoint.py` — MetaPhlAn / Kraken2 / QIIME 2 / generic lineage tables → waypoint format.
|
||||
- `scripts/vocab_coverage.py` — tokenizer coverage report for a waypoint-format file.
|
||||
|
||||
## Upstream
|
||||
|
||||
Code [github.com/Outpost-Bio/waypoint](https://github.com/Outpost-Bio/waypoint) ·
|
||||
package `waypoint-bio` ·
|
||||
paper [bioRxiv 2026.05.02.722381](https://www.biorxiv.org/content/10.64898/2026.05.02.722381v2) ·
|
||||
community [Waypoint Slack](https://join.slack.com/t/outpostbio-waypoint/shared_invite/zt-3w6ivgtba-WJOCkdxiISxQpwVq9ZZxTA) ·
|
||||
contact `waypoint@outpost.bio`.
|
||||
|
||||
Cite Treloar, N. J., Ur-Rehman, S., Yang, J., & Outpost Bio (2026). *Learning the Language of the
|
||||
Microbiome with Transformers.* bioRxiv. Per-artefact DOIs are listed at
|
||||
[outpost.bio/citations](https://www.outpost.bio/citations).
|
||||
210
skills/waypoint-bio/references/cli-reference.md
Normal file
210
skills/waypoint-bio/references/cli-reference.md
Normal file
@@ -0,0 +1,210 @@
|
||||
# `waypoint` CLI reference
|
||||
|
||||
Targets `waypoint-bio` 1.0.2 (PyPI) / 1.0.4 (GitHub main, commit `f45eee6`, 2026-07-16).
|
||||
|
||||
```
|
||||
waypoint {pretrain,benchmark,finetune,embed,prepare-dataset} ...
|
||||
```
|
||||
|
||||
Config paths are resolved first against the working directory, then against the bundled
|
||||
`waypoint_bio/configs/` tree inside the installed wheel. So `--config configs/benchmark.yaml`
|
||||
works from anywhere without cloning the repo. The same fallback applies to the bundled example
|
||||
data (`examples/abundance_matrix.tsv`, `examples/finetune_classification.parquet`, …).
|
||||
|
||||
---
|
||||
|
||||
## `waypoint prepare-dataset`
|
||||
|
||||
Converts a sample × taxa abundance matrix into waypoint format.
|
||||
|
||||
| Flag | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `--input` | *required* | `.csv` / `.tsv` abundance matrix. |
|
||||
| `--output` | *required* | `.parquet` recommended; `.csv` supported. |
|
||||
| `--orientation` | `auto` | `auto`, `samples_as_rows`, `taxa_as_rows`. |
|
||||
| `--taxonomy_format` | `full` | `full` for lineage strings; a rank name (`genus`, `species`, …) to prefix bare names. |
|
||||
| `--no_normalize` | off | Skip row-normalisation to relative abundances. |
|
||||
| `--keep_zeros` | off | Keep zero-abundance entries in each sample's lists. |
|
||||
| `--metadata` | none | CSV/TSV/parquet of per-sample metadata, indexed by sample ID, merged in as extra columns. |
|
||||
|
||||
`auto` treats the file as taxa-as-rows when the first column header is `taxonomy`, `lineage`,
|
||||
`taxon`, `otu`, or `#otu id` (case-insensitive); otherwise samples-as-rows with the first column
|
||||
as the sample ID.
|
||||
|
||||
`--taxonomy_format genus` prefixes bare column names with `g__`. It disables higher-rank fallback,
|
||||
because a bare name carries no lineage to fall back to — prefer real lineage strings.
|
||||
|
||||
---
|
||||
|
||||
## `waypoint embed`
|
||||
|
||||
One fixed-size vector per sample from a pretrained checkpoint. No fine-tuning, no labels needed.
|
||||
|
||||
| Flag | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `--model` | `outpost-bio/Waypoint-6m` | Hub id or local checkpoint directory. |
|
||||
| `--data` | *required* | Waypoint-format `.parquet` / `.csv` / `.tsv`. |
|
||||
| `--output` | *required* | `.parquet`, or `.csv` if the path ends in `.csv`. |
|
||||
| `--pooling` | `last_token` | `last_token`, `mean`, `first_token`, `cls_token`. |
|
||||
| `--batch_size` | `32` | |
|
||||
| `--max_length` | `512` | Truncates after ordering, so the least informative taxa are lost first. |
|
||||
| `--device` | auto | `cuda`, `mps`, or `cpu`; auto-detects in that order. |
|
||||
|
||||
Output columns are `dim_0 … dim_{H-1}`, indexed by sample ID. Hidden size `H` is 256 (6m),
|
||||
512 (45m), 768 (170m).
|
||||
|
||||
**Behaviour worth knowing:** tokens that map to `<unk>` are *dropped* before ordering, not encoded.
|
||||
A row with no in-vocabulary taxa still produces an output row, but its sequence is `[BOS][EOS]` and
|
||||
the embedding is meaningless. Run `scripts/vocab_coverage.py` first.
|
||||
|
||||
Ordering: by descending abundance z-score when `token_std_means.parquet` is present (it ships with
|
||||
every published checkpoint and with `waypoint pretrain` output), otherwise by descending raw
|
||||
relative abundance.
|
||||
|
||||
---
|
||||
|
||||
## `waypoint finetune`
|
||||
|
||||
Fine-tunes a checkpoint on your own labelled waypoint-format data.
|
||||
|
||||
| Flag | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `--model` | *required* | Hub id or local checkpoint. |
|
||||
| `--data` | *required* | Waypoint-format file containing `--target`. |
|
||||
| `--output_dir` | *required* | |
|
||||
| `--task_type` | *required* | `classification` or `regression`. |
|
||||
| `--target` | *required* | Target column name. |
|
||||
| `--covariate_column` | none | Categorical column, one-hot encoded and concatenated to the pooled embedding before the head. |
|
||||
| `--config` | task default | Flat YAML; defaults to the bundled classification/regression config. |
|
||||
|
||||
### Fine-tuning config keys
|
||||
|
||||
```yaml
|
||||
split_column: null # column holding train/validation/test; null = random split
|
||||
val_fraction: 0.1
|
||||
test_fraction: 0.1
|
||||
|
||||
max_length: 512 # must match the checkpoint's pretraining context
|
||||
pooling_strategy: last_token
|
||||
filter_unk_taxa: true # drop out-of-vocabulary taxa rather than feed <unk>
|
||||
|
||||
seed: 42
|
||||
learning_rate: 0.00003
|
||||
num_epochs: 1 # raise this; early stopping is what should terminate training
|
||||
batch_size: 64
|
||||
warmup_steps: 1000 # lower to ~50 for small datasets
|
||||
weight_decay: 0.001
|
||||
eval_strategy: steps
|
||||
eval_steps: 400
|
||||
logging_steps: 5
|
||||
patience: 5 # eval steps without improvement before early stopping
|
||||
save_total_limit: 1
|
||||
|
||||
use_lora: false
|
||||
lora_r: 8
|
||||
lora_alpha: 16 # convention: 2 * r
|
||||
lora_dropout: 0.05
|
||||
lora_target_modules: [c_attn, c_proj] # GPT-2 fused QKV and output projection
|
||||
lora_bias: none
|
||||
lora_fan_in_fan_out: true # required for GPT-2 Conv1D layouts
|
||||
```
|
||||
|
||||
`num_epochs: 1` in the shipped configs is tuned for the large Compass tasks. On a few-thousand-row
|
||||
dataset one epoch is a handful of optimizer steps and the model barely moves — raise `num_epochs`
|
||||
and let `patience` stop it. Likewise `eval_steps: 400` may never fire; lower it so early stopping
|
||||
and best-checkpoint selection can actually work.
|
||||
|
||||
LoRA adapters are merged back into the base transformer before saving, so `best_model/` loads with
|
||||
a plain `AutoModel.from_pretrained` and works with `waypoint embed` and `waypoint benchmark`.
|
||||
|
||||
### Outputs
|
||||
|
||||
| Path | Contents |
|
||||
| --- | --- |
|
||||
| `best_model/` | Fine-tuned base transformer in standard HF format, plus tokenizer and `token_std_means.parquet`. |
|
||||
| `best_model/finetuned_model_state.pt` | Full torch state dict: transformer + head + covariate embedding. |
|
||||
| `validation_metrics.json`, `test_metrics.json` | Per-split scores, benchmark-equivalent. |
|
||||
| `training_log.csv`, `training_log.html` | Every row of `trainer.state.log_history`; the HTML is an interactive plotly line plot. |
|
||||
| `finetune_results.json` | Run config, label maps, covariate map, val/test scores. |
|
||||
|
||||
---
|
||||
|
||||
## `waypoint benchmark`
|
||||
|
||||
| Flag | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `--model` | `outpost-bio/Waypoint-6m` | Hub id or local checkpoint. |
|
||||
| `--config` | bundled `configs/benchmark.yaml` | Shared by all eight tasks. |
|
||||
| `--output_dir` | `outputs/benchmark` | |
|
||||
| `--tasks` | all 8 | Space-separated task numbers, e.g. `--tasks 1 6`. |
|
||||
| `--seed` | `42` | |
|
||||
| `--max_samples` | none | Caps each split; use for smoke tests only, never for a reported score. |
|
||||
|
||||
`configs/benchmark.yaml` is the fine-tuning config applied identically to every task:
|
||||
`learning_rate: 3e-5`, `num_epochs: 1`, `batch_size: 64`, `warmup_steps: 1000`,
|
||||
`weight_decay: 0.001`, `patience: 5`, `pooling_strategy: last_token`, `eval_steps: 400`,
|
||||
`filter_unk_taxa: true`, `seed: 42`. Change it and your score is no longer comparable to the paper.
|
||||
|
||||
The paper reports means over three independent runs. A single run is noisy; vary `--seed` and
|
||||
report the spread.
|
||||
|
||||
---
|
||||
|
||||
## `waypoint pretrain`
|
||||
|
||||
| Flag | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `--model_config` | `configs/models/gpt2-6m.yaml` | Architecture YAML. |
|
||||
| `--pretrain_config` | `configs/pretraining.yaml` | Hyperparameter YAML. |
|
||||
| `--output_dir` | `outputs/pretrain` | Best checkpoint written to `<output_dir>/best_model/`. |
|
||||
| `--max_samples` | none | Limit training samples for a quick test. |
|
||||
| `--data` | none | Local waypoint-format corpus instead of downloading Atlas. |
|
||||
|
||||
Steps: download the Atlas `pretrain` split → build a taxonomic tokenizer from the corpus →
|
||||
compute per-token abundance mean/std for z-score ordering → train GPT-2 with next-token prediction
|
||||
and early stopping → save `best_model/`.
|
||||
|
||||
### `configs/pretraining.yaml`
|
||||
|
||||
```yaml
|
||||
training_type: next_token_prediction
|
||||
taxon_rank: genus # tokenization rank; changing it means re-pretraining
|
||||
fallback_to_higher_rank: true # use the most specific higher rank when genus is absent
|
||||
max_length: 512
|
||||
learning_rate: 0.001
|
||||
warmup_steps: 1000
|
||||
weight_decay: 0.001
|
||||
batch_size: 32
|
||||
num_epochs: 100
|
||||
patience: 10
|
||||
eval_steps: 3261
|
||||
save_steps: 3261
|
||||
logging_steps: 100
|
||||
val_split: 0.1
|
||||
seed: 42
|
||||
```
|
||||
|
||||
### Architectures
|
||||
|
||||
All share `model_type: gpt2`, `n_positions: 512`, and a fixed per-head dimension of 64.
|
||||
|
||||
| Config | Layers | Hidden | Heads | ~Params |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `gpt2-6m.yaml` | 8 | 256 | 4 | 6M |
|
||||
| `gpt2-6m-mgm.yaml` | 8 | 256 | 8 | 6M — matches the MGM baseline architecture |
|
||||
| `gpt2-10m.yaml` | 8 | 320 | 5 | 10M |
|
||||
| `gpt2-18m.yaml` | 10 | 384 | 6 | 18M |
|
||||
| `gpt2-29m.yaml` | 12 | 448 | 7 | 29M |
|
||||
| `gpt2-45m.yaml` | 14 | 512 | 8 | 45M |
|
||||
| `gpt2-79m.yaml` | 16 | 640 | 10 | 79M |
|
||||
| `gpt2-85m-gpt-small.yaml` | 12 | 768 | 12 | 85M — GPT-2 small geometry |
|
||||
| `gpt2-170m.yaml` | 24 | 768 | 12 | 170M |
|
||||
|
||||
Only 6m, 45m, and 170m are published as checkpoints. The rest exist so the paper's scaling study is
|
||||
reproducible; `gpt2-6m-mgm` isolates the effect of head count against the MGM baseline.
|
||||
|
||||
Parameter counts exclude token and positional embeddings, so the Hub's reported sizes are larger
|
||||
(the 6m checkpoint reports ~10.1M, the 45m ~51.8M).
|
||||
|
||||
Pretraining Atlas end to end is a multi-GPU-day job. Validate the pipeline with
|
||||
`--max_samples 5000` before committing to a full run.
|
||||
124
skills/waypoint-bio/references/compass-benchmark.md
Normal file
124
skills/waypoint-bio/references/compass-benchmark.md
Normal file
@@ -0,0 +1,124 @@
|
||||
# Compass: the eight-task microbiome benchmark
|
||||
|
||||
`outpost-bio/Compass` on the Hugging Face Hub — gated, Apache 2.0, ~605 MB, ~62.8k rows across four
|
||||
Hub configurations. Eight tasks are derived from those four configurations by filtering and by
|
||||
choosing different target columns.
|
||||
|
||||
Every configuration exposes `train` / `validation` / `test` splits and carries a `Split` column
|
||||
recording the same assignment.
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
ds = load_dataset("outpost-bio/Compass", "mgnify-biomes") # requires access + HF_TOKEN
|
||||
```
|
||||
|
||||
## The four source datasets
|
||||
|
||||
| Config | Source | Rows (train/val/test) | Extra columns |
|
||||
| --- | --- | --- | --- |
|
||||
| `mgnify-biomes` | MGnify metagenomic profiles across gut, skin, oral, marine, freshwater, soil, engineered systems | 33,121 / 4,139 / 4,139 | `Biome 1`–`Biome 5`, `Run Accession`, `Data Type`, `Sequencing Method`, `Pipeline Version`, `Study Accession` |
|
||||
| `handuo` | Han, Duo et al. — 16S amplicon study of drug–microbiome interactions in stool-derived communities | 3,168 / 396 / 396 | `SIC Name`, `Control`, `ATC Class`, `Sample ID` |
|
||||
| `mastrorilli` | Mastrorilli et al. — drug degradation by gut communities | 9,282 / 3,084 / 3,053 | `Degradation Rate`, `Drug`, `Sample ID` |
|
||||
| `roswall` | Roswall et al. — longitudinal infant gut cohort | 2,031 total | `Timepoint`, `Delivery Mode`, `Sample ID` |
|
||||
|
||||
All configs carry `Taxa` and `Relative Abundances` as aligned list columns.
|
||||
|
||||
## The eight tasks
|
||||
|
||||
As defined in `waypoint_bio/benchmark.py`:
|
||||
|
||||
| # | Internal id | Config | Targets | Type | Pre-filter |
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| 1 | `1_biome` | `mgnify-biomes` | `Biome 1`–`Biome 5` | classification (5 outputs) | none |
|
||||
| 2 | `2_biome_gut` | `mgnify-biomes` | `Biome 4`, `Biome 5` | classification (2 outputs) | `Biome 3 == "Digestive system"` |
|
||||
| 3 | `3_sic` | `handuo` | `SIC Name` | classification | `SIC Name` starts with `SIC`, excludes `control` and `seed` |
|
||||
| 4 | `4_drug_non_drug` | `handuo` | `Control` | binary classification | none |
|
||||
| 5 | `5_drug_class` | `handuo` | `ATC Class` | classification | `ATC Class` not null |
|
||||
| 6 | `6_drug_degradation` | `mastrorilli` | `Degradation Rate` | regression | none; `Drug` used as covariate |
|
||||
| 7 | `7_infant_age` | `roswall` | `Timepoint` | classification | none |
|
||||
| 8 | `8_birth_mode` | `roswall` | `Delivery Mode` | binary classification | none |
|
||||
|
||||
What each asks, in plain terms:
|
||||
|
||||
1. **Biome classification** — predict all five levels of the MGnify biome ontology at once
|
||||
(e.g. `root → Host-associated → Human → Digestive system → Large intestine`).
|
||||
2. **Gut biome classification** — same, restricted to digestive-system samples, predicting only the
|
||||
two finest levels. Harder: the easy environmental separations are gone.
|
||||
3. **SIC classification** — identify which stool-derived in-vitro community a drug-perturbed sample
|
||||
came from.
|
||||
4. **Drug vs. control** — did this community receive a drug?
|
||||
5. **Drug class** — recover the ATC class of the applied drug from the resulting composition.
|
||||
6. **Drug degradation** — regress the degradation rate from composition plus drug identity. The
|
||||
`Drug` covariate is one-hot encoded and concatenated to the pooled embedding.
|
||||
7. **Infant age** — predict the sampling timepoint from an infant gut sample.
|
||||
8. **Birth mode** — vaginal vs. caesarean delivery.
|
||||
|
||||
## Scoring
|
||||
|
||||
- **Classification:** macro-averaged F1 — F1 per class, averaged with equal weight. Chosen so the
|
||||
metric is not dominated by majority classes. Where a task has several target columns (1 and 2),
|
||||
the per-target macro-F1s are averaged.
|
||||
- **Regression (task 6):** R², clamped to `[0, 1]` so it shares a scale with the F1 scores. A
|
||||
negative R² therefore reads as `0.0`, not as "worse than the mean".
|
||||
- **Final score:** unweighted arithmetic mean of the eight task scores.
|
||||
|
||||
Supplementary metrics are computed and stored but do not enter the score: one-vs-one macro ROC-AUC,
|
||||
macro PR-AUC (pairwise average precision over the same OVO pairs), balanced accuracy, plain
|
||||
accuracy; and MSE, Pearson, Spearman for regression.
|
||||
|
||||
## `benchmark_results.json`
|
||||
|
||||
```
|
||||
benchmark_results.json
|
||||
├── model string — the value passed to --model
|
||||
├── final_score number — mean of every results[].score
|
||||
└── results array, one object per task
|
||||
├── task string — "1_biome", "6_drug_degradation", ...
|
||||
├── task_type "classification" | "regression"
|
||||
├── score number — macro F1, or R² clamped to [0,1]
|
||||
└── metrics object — keys depend on task_type
|
||||
```
|
||||
|
||||
`metrics` keys are suffixed with the target column name:
|
||||
|
||||
| Task type | Keys |
|
||||
| --- | --- |
|
||||
| `classification` | `accuracy_<target>`, `balanced_accuracy_<target>`, `f1_macro_<target>`; with probabilities, binary `roc_auc_<target>` / `pr_auc_<target>` or multiclass `roc_auc_macro_ovo_<target>` / `pr_auc_macro_ovo_<target>`. Means: `f1_macro_mean`, optionally `roc_auc_mean`, `pr_auc_mean`. |
|
||||
| `regression` | `mse_<target>`, `r2_<target>`, usually `pearson_<target>` and `spearman_<target>`. Mean: `r2_mean`. |
|
||||
|
||||
Example:
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "outpost-bio/Waypoint-6m",
|
||||
"final_score": 0.71,
|
||||
"results": [
|
||||
{"task": "1_biome", "task_type": "classification", "score": 0.65,
|
||||
"metrics": {"f1_macro_mean": 0.65, "roc_auc_mean": 0.81, "pr_auc_mean": 0.74}},
|
||||
{"task": "6_drug_degradation", "task_type": "regression", "score": 0.42,
|
||||
"metrics": {"mse_Degradation Rate": 0.019, "r2_Degradation Rate": 0.44, "r2_mean": 0.44}}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
The numbers above are the illustrative values from the upstream README, not measured results.
|
||||
|
||||
## Interpreting a benchmark run
|
||||
|
||||
**Baselines matter more than the absolute score.** The paper compares Waypoint against classical
|
||||
baselines (random forest and logistic regression on relative abundances) and against MGM, the prior
|
||||
microbiome foundation model. Two findings shape how a Compass number should be read:
|
||||
|
||||
- Waypoint beats the random-forest baseline from roughly **10,000 training examples upward**, and
|
||||
*loses* to it below about 1,000. Report the training-set size next to any score.
|
||||
- Baselines can use every taxon; the transformer sees only its fixed vocabulary. The paper's fair
|
||||
comparison is the `(no unk)` baseline, with out-of-vocabulary taxa stripped from the baseline's
|
||||
input too. Compare against that, not against a baseline given the full table.
|
||||
|
||||
**Scale does not monotonically help.** Pretraining loss falls all the way to 170M, but the best
|
||||
Compass score in the paper came from the **45M** model. Non-pretrained transformers get *worse* as
|
||||
they grow — the gain from scale is a property of pretraining, not of capacity.
|
||||
|
||||
**Reproducibility.** Use the bundled `configs/benchmark.yaml` unchanged, do not pass `--max_samples`,
|
||||
and run at least three seeds. Comparing a run that changed the learning rate or capped splits against
|
||||
published numbers is not a comparison.
|
||||
200
skills/waypoint-bio/references/data-preparation.md
Normal file
200
skills/waypoint-bio/references/data-preparation.md
Normal file
@@ -0,0 +1,200 @@
|
||||
# Preparing data for Waypoint
|
||||
|
||||
## Waypoint format
|
||||
|
||||
Rows are samples. Two aligned list-columns, plus whatever labels you need.
|
||||
|
||||
| Column | Type | Required |
|
||||
| --- | --- | --- |
|
||||
| `Taxa` | `list[str]` — full lineage strings | yes |
|
||||
| `Relative Abundances` | `list[float]` — same length and order as `Taxa` | yes |
|
||||
| `Split` | `str` — `train` / `validation` / `test` | only when using `split_column` |
|
||||
| *(any)* | scalar targets and covariates | as needed |
|
||||
|
||||
The DataFrame index holds the sample ID and is preserved through `embed`.
|
||||
|
||||
Use `.parquet`. CSV/TSV works but stores each list as its Python `repr`, parsed back with
|
||||
`ast.literal_eval` — brittle and large.
|
||||
|
||||
```python
|
||||
import pandas as pd
|
||||
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"Taxa": [["k__Bacteria; p__Firmicutes; c__Bacilli; o__Lactobacillales; f__Lactobacillaceae; g__Lactobacillus",
|
||||
"k__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Bacteroidales; f__Bacteroidaceae; g__Bacteroides"]],
|
||||
"Relative Abundances": [[0.41, 0.59]],
|
||||
"Group": ["Case"],
|
||||
},
|
||||
index=pd.Index(["sample_001"], name="sample_id"),
|
||||
)
|
||||
df.to_parquet("dataset.parquet")
|
||||
```
|
||||
|
||||
## How taxonomy strings are read
|
||||
|
||||
`TaxonomicTokenizer` splits each lineage on `;`, strips whitespace, and inspects each segment's
|
||||
three-character prefix:
|
||||
|
||||
| Prefix | Rank |
|
||||
| --- | --- |
|
||||
| `s__` | species |
|
||||
| `g__` | genus |
|
||||
| `f__` | family |
|
||||
| `o__` | order |
|
||||
| `c__` | class |
|
||||
| `p__` | phylum |
|
||||
| `k__` | kingdom |
|
||||
|
||||
With `taxon_rank: genus` and `fallback_to_higher_rank: true` (the published defaults), each lineage
|
||||
becomes one token:
|
||||
|
||||
1. If a `g__` segment exists, that segment *including the prefix* is the token — `g__Lactobacillus`.
|
||||
2. Otherwise the **most specific higher rank** present is used — a lineage stopping at
|
||||
`f__Lactobacillaceae` tokenises to `f__Lactobacillaceae`.
|
||||
3. If nothing matches, the token is `<unk>`.
|
||||
|
||||
Consequences that bite:
|
||||
|
||||
- **Any prefix outside that table is invisible.** QIIME 2 / SILVA / Greengenes2 write the domain as
|
||||
`d__Bacteria`; `d__` is not in the table, so such a segment is skipped entirely. A lineage
|
||||
truncated at domain becomes `<unk>`. Rewrite `d__` to `k__`.
|
||||
- **A `s__` species segment does not help by itself.** Species is *more* specific than genus, so
|
||||
fallback (which only goes up) cannot use it. A lineage with `s__` but no `g__` tokenises to
|
||||
whatever higher rank is present — or `<unk>` if none is. Keep the full lineage, not just the tip.
|
||||
- **Separator is `;`, not `|`.** A `|`-joined MetaPhlAn lineage is one unsplittable segment. Its
|
||||
first three characters are `k__`, so it matches at kingdom rank and the *entire pipe-joined
|
||||
string* is returned as a single token — which is not in the vocabulary, so it becomes `<unk>`.
|
||||
Verified against `TaxonomicTokenizer` 1.0.2:
|
||||
`k__Bacteria|p__Firmicutes|g__Lactobacillus` extracts to itself, while the `;`-separated form
|
||||
extracts to `g__Lactobacillus`.
|
||||
- **Bare names never tokenise.** `Lactobacillus` has no prefix. Use `prepare-dataset
|
||||
--taxonomy_format genus` to prefix them, accepting the loss of fallback.
|
||||
|
||||
## Token ordering and truncation
|
||||
|
||||
Samples are encoded as `[BOS] + ordered_token_ids + [EOS]`, padded to `max_length` (512).
|
||||
|
||||
Ordering is by **descending abundance z-score** — `(ra - mean) / std` per token, using
|
||||
`token_std_means.parquet` from the checkpoint. This puts taxa that are unusually abundant *for that
|
||||
taxon* first, rather than merely abundant. Without that file, ordering falls back to raw descending
|
||||
abundance.
|
||||
|
||||
Because truncation is applied after ordering, a sample with more than 510 in-vocabulary taxa loses
|
||||
its least distinctive ones. That is the intended behaviour, but it means `max_length` interacts with
|
||||
how deeply you profiled.
|
||||
|
||||
## Out-of-vocabulary taxa
|
||||
|
||||
The vocabulary is frozen at pretraining time from the Atlas corpus. During `waypoint embed`, tokens
|
||||
resolving to `<unk>` are dropped before ordering; during fine-tuning and benchmarking,
|
||||
`filter_unk_taxa: true` does the same. Neither warns you.
|
||||
|
||||
Every Compass dataset carries out-of-vocabulary taxa, and the paper names this the models' key
|
||||
limitation. Measure it before drawing conclusions:
|
||||
|
||||
```bash
|
||||
python scripts/vocab_coverage.py --model outpost-bio/Waypoint-6m --data dataset.parquet
|
||||
```
|
||||
|
||||
If coverage is poor, the usual causes are, in order: a different taxonomy database (SILVA vs. NCBI
|
||||
vs. GTDB naming), the `d__` prefix problem, `|` separators, and genuinely novel environments.
|
||||
|
||||
## Converting profiler output
|
||||
|
||||
`waypoint prepare-dataset` reads a plain abundance matrix whose labels are already `;`-separated
|
||||
lineages. `scripts/profiler_to_waypoint.py` handles the formats it cannot.
|
||||
|
||||
### MetaPhlAn
|
||||
|
||||
Merged tables from `merge_metaphlan_tables.py`: rows are clades with `|`-separated lineages, columns
|
||||
are samples, values are **percentages**, and the table is cumulative — every rank appears as its own
|
||||
row.
|
||||
|
||||
```bash
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input merged_abundance_table.txt --format metaphlan \
|
||||
--rank species --output dataset.parquet
|
||||
```
|
||||
|
||||
The converter drops `#` comment lines and the `NCBI_tax_id` / `clade_taxid` column, keeps only rows
|
||||
whose deepest rank equals `--rank` (default `species`, which avoids double-counting parents),
|
||||
rewrites `|` to `; `, and renormalises each sample to sum to 1. Strain rows (`t__`) are always
|
||||
excluded.
|
||||
|
||||
### Kraken2 / Bracken
|
||||
|
||||
Kraken2 reports are per-sample and encode the hierarchy as two-space indentation, with no lineage
|
||||
string. Pass one report per sample:
|
||||
|
||||
```bash
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input reports/*.kreport --format kraken \
|
||||
--rank species --output dataset.parquet
|
||||
```
|
||||
|
||||
The converter walks the indentation to rebuild each lineage, maps Kraken rank codes to prefixes
|
||||
(`D`/`K` → `k__`, `P` → `p__`, `C` → `c__`, `O` → `o__`, `F` → `f__`, `G` → `g__`, `S` → `s__`),
|
||||
skips sub-ranks (`D1`, `S1`, …) and unclassified rows, takes clade-level read counts at the target
|
||||
rank, and normalises. Sample IDs come from the filenames. Both the 6-column and the 8-column
|
||||
(`--report-minimizer-data`) layouts are handled.
|
||||
|
||||
Bracken's own `.bracken` output carries no lineage at all — use the Kraken-style report Bracken
|
||||
writes with `-o`/`--report`, not the tabular abundance file.
|
||||
|
||||
### QIIME 2 / biom TSV
|
||||
|
||||
Exported feature tables with a `taxonomy` column (or `#OTU ID` rows already labelled by lineage):
|
||||
|
||||
```bash
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input feature-table.tsv --format qiime2 \
|
||||
--taxonomy-column taxonomy --output dataset.parquet
|
||||
```
|
||||
|
||||
The converter strips the `# Constructed from biom file` banner, uses the taxonomy column as the
|
||||
lineage, rewrites `d__` to `k__`, and normalises counts to relative abundances. Features whose
|
||||
taxonomy is `Unassigned` are dropped.
|
||||
|
||||
### MGnify
|
||||
|
||||
MGnify amplicon abundance TSVs are taxa-as-rows with a `taxonomy` first column and `;`-separated
|
||||
lineages — the native layout. `waypoint prepare-dataset --orientation auto` reads them directly; no
|
||||
conversion needed. This is the format Atlas itself was built from.
|
||||
|
||||
### Anything else
|
||||
|
||||
If you already have a sample × taxa table with lineage labels, use `--format generic`, which applies
|
||||
only the separator and prefix normalisation:
|
||||
|
||||
```bash
|
||||
python scripts/profiler_to_waypoint.py \
|
||||
--input my_table.tsv --format generic --orientation taxa_as_rows \
|
||||
--output dataset.parquet
|
||||
```
|
||||
|
||||
## Attaching labels
|
||||
|
||||
Either merge them at conversion time —
|
||||
|
||||
```bash
|
||||
waypoint prepare-dataset --input matrix.tsv --metadata labels.csv --output dataset.parquet
|
||||
python scripts/profiler_to_waypoint.py --input ... --metadata labels.csv --output dataset.parquet
|
||||
```
|
||||
|
||||
— where `labels.csv` is indexed by sample ID, or join afterwards in pandas. Sample IDs must match
|
||||
exactly; the converters do an inner-style alignment and will silently produce `NaN` targets for
|
||||
unmatched rows, which then fail at fine-tuning time.
|
||||
|
||||
## Splits
|
||||
|
||||
`waypoint finetune` defaults to a random 80/10/10 split. Add a `Split` column and set
|
||||
`split_column: Split` in the config whenever samples are not independent:
|
||||
|
||||
- longitudinal cohorts (the Roswall infant data is exactly this shape),
|
||||
- technical or biological replicates,
|
||||
- multiple communities derived from one donor,
|
||||
- multiple drugs applied to the same starting community.
|
||||
|
||||
Grouping by subject or study when you build `Split` is the difference between a generalisation
|
||||
estimate and a memorisation estimate.
|
||||
219
skills/waypoint-bio/references/python-api.md
Normal file
219
skills/waypoint-bio/references/python-api.md
Normal file
@@ -0,0 +1,219 @@
|
||||
# Using Waypoint from Python
|
||||
|
||||
The CLI covers the standard paths. Drop to Python when you need a custom training loop, a different
|
||||
head, or embeddings inside a larger pipeline.
|
||||
|
||||
## Package surface
|
||||
|
||||
`waypoint_bio` lazily re-exports:
|
||||
|
||||
```python
|
||||
from waypoint_bio import (
|
||||
TaxonomicTokenizer, # the tokenizer class
|
||||
load_tokenizer, # load one from a Hub id or local dir
|
||||
MicrobiomePretrainingDataset, # causal-LM dataset
|
||||
MicrobiomeBenchmarkDataset, # supervised dataset with targets/covariates
|
||||
load_waypoint_dataframe, # read waypoint-format parquet/csv/tsv
|
||||
load_abundance_matrix, # read a sample x taxa matrix
|
||||
matrix_to_waypoint_df, # matrix -> waypoint format
|
||||
)
|
||||
```
|
||||
|
||||
Imports are deferred, so `import waypoint_bio` does not pull in torch.
|
||||
|
||||
## Loading a checkpoint directly with transformers
|
||||
|
||||
The tokenizer is custom and ships as remote code, so `trust_remote_code=True` is required for it.
|
||||
The model itself is a stock GPT-2 and does not need it.
|
||||
|
||||
```python
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
tok = AutoTokenizer.from_pretrained("outpost-bio/Waypoint-45m", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("outpost-bio/Waypoint-45m") # gated: needs HF_TOKEN
|
||||
```
|
||||
|
||||
`trust_remote_code=True` executes the tokenizer code stored in the repo. Pin a revision when that
|
||||
matters to you, so the code cannot change under a later run:
|
||||
|
||||
```python
|
||||
tok = AutoTokenizer.from_pretrained(
|
||||
"outpost-bio/Waypoint-45m", trust_remote_code=True, revision="1664ab5"
|
||||
)
|
||||
```
|
||||
|
||||
`AutoModelForCausalLM` also works if you want the LM head for likelihood scoring or generation —
|
||||
generation samples taxa, which is occasionally useful for probing what the model learned about
|
||||
co-occurrence, but is not a validated use.
|
||||
|
||||
## Tokenizing by hand
|
||||
|
||||
```python
|
||||
from waypoint_bio import load_tokenizer
|
||||
|
||||
tok = load_tokenizer("outpost-bio/Waypoint-6m")
|
||||
|
||||
lineage = "k__Bacteria; p__Firmicutes; c__Bacilli; o__Lactobacillales; f__Lactobacillaceae; g__Lactobacillus"
|
||||
print(tok.tokenize(lineage)) # ['g__Lactobacillus']
|
||||
print(tok.convert_tokens_to_ids(["g__Lactobacillus"]))
|
||||
|
||||
# One sample = newline-separated lineages
|
||||
sample = "\n".join([lineage, "k__Bacteria; p__Bacteroidota; g__Bacteroides"])
|
||||
print(tok(sample)["input_ids"])
|
||||
```
|
||||
|
||||
Checking whether a taxon is in vocabulary:
|
||||
|
||||
```python
|
||||
vocab = tok.get_vocab()
|
||||
"g__Lactobacillus" in vocab # True for anything seen in Atlas
|
||||
tok.convert_tokens_to_ids("g__Nonesuch") == tok.unk_token_id
|
||||
```
|
||||
|
||||
`tok._extract(lineage)` applies the rank extraction and higher-rank fallback and returns the token
|
||||
string, or `None`. It is private but stable across 1.0.x and is what the datasets and
|
||||
`scripts/vocab_coverage.py` use.
|
||||
|
||||
## Building a dataset
|
||||
|
||||
```python
|
||||
import pandas as pd
|
||||
from waypoint_bio import MicrobiomePretrainingDataset, load_tokenizer, load_waypoint_dataframe
|
||||
from waypoint_bio.dataset import try_load_token_std_means
|
||||
|
||||
df = load_waypoint_dataframe("dataset.parquet")
|
||||
tok = load_tokenizer("outpost-bio/Waypoint-6m")
|
||||
stats = try_load_token_std_means("outpost-bio/Waypoint-6m") # None if absent
|
||||
|
||||
ds = MicrobiomePretrainingDataset(df, tok, max_length=512, token_std_means=stats)
|
||||
ds[0]["input_ids"].shape # torch.Size([512])
|
||||
```
|
||||
|
||||
Each item is `[BOS] + z-score-ordered token ids + [EOS]`, right-padded.
|
||||
|
||||
Computing the ordering statistics for a corpus of your own:
|
||||
|
||||
```python
|
||||
from waypoint_bio.dataset import compute_token_std_means
|
||||
|
||||
stats = compute_token_std_means(df, tok, show_progress=True)
|
||||
stats.to_parquet("token_std_means.parquet") # index name "token", columns mean/std
|
||||
```
|
||||
|
||||
Drop that file next to a checkpoint and `embed`, `finetune`, and `benchmark` will pick it up.
|
||||
|
||||
## Embeddings without the CLI
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
from waypoint_bio.dataset import load_waypoint_dataframe, try_load_token_std_means
|
||||
from waypoint_bio.embed import tokenize_for_embedding
|
||||
from waypoint_bio.models import _pool
|
||||
from waypoint_bio.tokenizer import load_tokenizer
|
||||
|
||||
model_id = "outpost-bio/Waypoint-45m"
|
||||
df = load_waypoint_dataframe("dataset.parquet")
|
||||
tok = load_tokenizer(model_id)
|
||||
model = AutoModel.from_pretrained(model_id).eval()
|
||||
|
||||
samples = tokenize_for_embedding(df, tok, max_length=512,
|
||||
token_std_means=try_load_token_std_means(model_id))
|
||||
|
||||
input_ids = torch.stack([s["input_ids"] for s in samples])
|
||||
attn = torch.stack([s["attention_mask"] for s in samples])
|
||||
|
||||
with torch.no_grad():
|
||||
hidden = model(input_ids=input_ids, attention_mask=attn).last_hidden_state
|
||||
emb = _pool(hidden, attn, "last_token") # [n_samples, hidden_size]
|
||||
```
|
||||
|
||||
`tokenize_for_embedding` preserves one output row per input row even when a row has no
|
||||
in-vocabulary taxa, so `emb` stays aligned with `df.index`. Those rows encode as `[BOS][EOS]` and
|
||||
their embeddings should be discarded, not interpreted.
|
||||
|
||||
## Custom heads
|
||||
|
||||
`waypoint_bio.models` provides the two heads used by `finetune` and `benchmark`:
|
||||
|
||||
```python
|
||||
from waypoint_bio.models import ClassificationModel, RegressionModel
|
||||
|
||||
head = ClassificationModel(
|
||||
base_model=model,
|
||||
tokenizer=tok,
|
||||
label_dims=[3], # one entry per target column
|
||||
pooling_strategy="last_token",
|
||||
covariate_dim=0, # width of the one-hot covariate block
|
||||
class_weights=None, # list[torch.Tensor], one per target
|
||||
)
|
||||
```
|
||||
|
||||
Both pool `last_hidden_state`, concatenate the one-hot covariate block if present, and apply one
|
||||
`nn.Linear` per target column. Multi-target classification masks label `-100` per target, so targets
|
||||
with missing values in some rows are handled without dropping the row.
|
||||
|
||||
Pooling strategies: `mean` (mask-weighted average), `last_token` (last non-padding position — the
|
||||
default and what the checkpoints were tuned for), `first_token` / `cls_token` (position 0, the BOS
|
||||
token; weak in a causal LM).
|
||||
|
||||
## Loading Atlas and Compass
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
atlas = load_dataset("outpost-bio/Atlas", split="pretrain") # 485,377 rows
|
||||
atlas_bench = load_dataset("outpost-bio/Atlas", split="benchmark") # 53,931 held out
|
||||
|
||||
compass = load_dataset("outpost-bio/Compass", "mastrorilli")
|
||||
compass["train"], compass["validation"], compass["test"]
|
||||
```
|
||||
|
||||
Atlas is ~5.6 GB. Stream it if you are only inspecting:
|
||||
|
||||
```python
|
||||
atlas = load_dataset("outpost-bio/Atlas", split="pretrain", streaming=True)
|
||||
first = next(iter(atlas))
|
||||
```
|
||||
|
||||
Atlas rows carry `Taxa`, `Relative Abundances`, `Run Accession`, `Data Type`, `Sequencing Method`,
|
||||
`Pipeline Version`, `Study Accession`. Filtering by `Data Type` or `Sequencing Method` before
|
||||
pretraining is a reasonable way to build a modality-specific model; filtering by `Study Accession` is
|
||||
how you would hold out whole studies.
|
||||
|
||||
Provenance: scraped from MGnify across pipeline versions v1.0–v5.0 and four modalities (16S amplicon,
|
||||
whole-genome shotgun, metagenomic assembly, and metatranscriptomic), then filtered to a minimum
|
||||
relative abundance of 1e-4 and a minimum of 10 taxa per sample. The pretrain/benchmark split is
|
||||
random with `seed=42` — it is *not* a study-level holdout, so the Atlas `benchmark` split shares
|
||||
studies with `pretrain`.
|
||||
|
||||
## Fine-tuning programmatically
|
||||
|
||||
There is no stable public function for the whole loop; `waypoint_bio.finetune` is written as a CLI
|
||||
module. Two workable options:
|
||||
|
||||
1. Call the CLI with `subprocess` and read `finetune_results.json` — what the upstream webinar
|
||||
notebooks do.
|
||||
2. Assemble it yourself from `MicrobiomeBenchmarkDataset` + `ClassificationModel`/`RegressionModel`
|
||||
and a `transformers.Trainer`, mirroring `benchmark.py`. Reuse `waypoint_bio.scoring.score_task`
|
||||
and `predictions_to_arrays` so your metrics match the published definitions.
|
||||
|
||||
```python
|
||||
import json, subprocess
|
||||
|
||||
subprocess.run([
|
||||
"waypoint", "finetune",
|
||||
"--model", "outpost-bio/Waypoint-45m",
|
||||
"--data", "dataset.parquet",
|
||||
"--output_dir", "outputs/ft",
|
||||
"--task_type", "classification",
|
||||
"--target", "Group",
|
||||
], check=True)
|
||||
|
||||
results = json.loads(open("outputs/ft/finetune_results.json").read())
|
||||
print(results["test_score"], results["test_metrics"])
|
||||
```
|
||||
|
||||
The upstream repo's `examples/webinar/` carries two worked notebooks — a regression walkthrough on
|
||||
Compass task 6 and a classification walkthrough on task 8 that also plots PCA / t-SNE projections of
|
||||
the embeddings against a logistic-regression baseline. Shared helpers live in `webinar_utils.py`.
|
||||
481
skills/waypoint-bio/scripts/profiler_to_waypoint.py
Executable file
481
skills/waypoint-bio/scripts/profiler_to_waypoint.py
Executable file
@@ -0,0 +1,481 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert microbiome profiler output into waypoint format.
|
||||
|
||||
`waypoint prepare-dataset` reads a plain abundance matrix whose row or column
|
||||
labels are already ``;``-separated lineage strings. Real profilers rarely emit
|
||||
that: MetaPhlAn separates ranks with ``|``, Kraken2 encodes the hierarchy as
|
||||
indentation with no lineage string at all, and QIIME 2 / SILVA prefixes the
|
||||
domain with ``d__``, which the Waypoint tokenizer does not recognise and
|
||||
silently skips.
|
||||
|
||||
This script normalises those layouts into a waypoint-format table with the
|
||||
``Taxa`` / ``Relative Abundances`` list-columns the CLI expects.
|
||||
|
||||
Examples
|
||||
--------
|
||||
python profiler_to_waypoint.py --input merged_metaphlan.tsv \
|
||||
--format metaphlan --output dataset.parquet
|
||||
|
||||
python profiler_to_waypoint.py --input reports/*.kreport \
|
||||
--format kraken --rank species --output dataset.parquet
|
||||
|
||||
python profiler_to_waypoint.py --input feature-table.tsv \
|
||||
--format qiime2 --taxonomy-column taxonomy --output dataset.parquet
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# Ranks the Waypoint tokenizer understands, most specific first.
|
||||
RANK_PREFIXES: dict[str, str] = {
|
||||
"species": "s__",
|
||||
"genus": "g__",
|
||||
"family": "f__",
|
||||
"order": "o__",
|
||||
"class": "c__",
|
||||
"phylum": "p__",
|
||||
"kingdom": "k__",
|
||||
}
|
||||
|
||||
# Kraken2 single-letter rank codes -> tokenizer prefixes. Domain is folded into
|
||||
# kingdom because the tokenizer has no domain rank.
|
||||
KRAKEN_RANK_CODES: dict[str, str] = {
|
||||
"D": "k__",
|
||||
"K": "k__",
|
||||
"P": "p__",
|
||||
"C": "c__",
|
||||
"O": "o__",
|
||||
"F": "f__",
|
||||
"G": "g__",
|
||||
"S": "s__",
|
||||
}
|
||||
|
||||
# Prefixes emitted by other databases that mean the same rank as a tokenizer
|
||||
# prefix. SILVA and Greengenes2 write the domain as ``d__``.
|
||||
PREFIX_ALIASES: dict[str, str] = {"d__": "k__", "sk__": "k__"}
|
||||
|
||||
_VALID_PREFIXES = set(RANK_PREFIXES.values())
|
||||
_RANK_CODE_RE = re.compile(r"^[A-Z-]\d*$")
|
||||
_UNASSIGNED = {"unassigned", "unclassified", "", "na", "nan", "none"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lineage normalisation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def normalise_lineage(lineage: str) -> str:
|
||||
"""Rewrite a lineage string into the form the Waypoint tokenizer reads.
|
||||
|
||||
Accepts ``|`` or ``;`` separators, rewrites aliased rank prefixes
|
||||
(``d__`` -> ``k__``), drops strain (``t__``) and empty segments, and joins
|
||||
with ``"; "``.
|
||||
"""
|
||||
text = str(lineage).strip().strip('"')
|
||||
separator = "|" if "|" in text else ";"
|
||||
segments: list[str] = []
|
||||
for raw in text.split(separator):
|
||||
segment = raw.strip()
|
||||
if not segment:
|
||||
continue
|
||||
for alias, replacement in PREFIX_ALIASES.items():
|
||||
if segment.startswith(alias):
|
||||
segment = replacement + segment[len(alias) :]
|
||||
break
|
||||
if segment.startswith("t__"):
|
||||
continue # strain-level; below the tokenizer's deepest rank
|
||||
# Anything without a recognised rank prefix -- "Unassigned", "root",
|
||||
# "cellular organisms" -- is invisible to the tokenizer, so drop it
|
||||
# rather than let it become a bogus taxon label. Bare prefixes such as
|
||||
# "g__" carry no name and go too.
|
||||
if segment[:3] not in _VALID_PREFIXES or len(segment) <= 3:
|
||||
continue
|
||||
segments.append(segment)
|
||||
return "; ".join(segments)
|
||||
|
||||
|
||||
def deepest_rank(lineage: str) -> str | None:
|
||||
"""Return the most specific rank name present in a normalised lineage."""
|
||||
present = {
|
||||
rank
|
||||
for rank, prefix in RANK_PREFIXES.items()
|
||||
if any(seg.strip().startswith(prefix) for seg in lineage.split(";"))
|
||||
}
|
||||
for rank in RANK_PREFIXES: # dict order is most specific first
|
||||
if rank in present:
|
||||
return rank
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parsers: each returns a samples x lineage abundance matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_metaphlan(path: Path, rank: str) -> pd.DataFrame:
|
||||
"""Parse a merged MetaPhlAn table into a samples x lineage matrix."""
|
||||
header: list[str] | None = None
|
||||
rows: list[list[str]] = []
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
line = line.rstrip("\n")
|
||||
if not line.strip():
|
||||
continue
|
||||
fields = line.split("\t")
|
||||
if line.startswith("#"):
|
||||
# The real header is the comment line naming the clade column.
|
||||
if header is None and any(
|
||||
f.lstrip("#").strip().lower() in {"clade_name", "taxonomy"}
|
||||
for f in fields
|
||||
):
|
||||
header = [f.lstrip("#").strip() for f in fields]
|
||||
continue
|
||||
if header is None:
|
||||
header = [f.strip() for f in fields]
|
||||
continue
|
||||
rows.append(fields)
|
||||
|
||||
if header is None or not rows:
|
||||
raise ValueError(f"{path}: no MetaPhlAn table found")
|
||||
|
||||
frame = pd.DataFrame(rows, columns=header)
|
||||
clade_col = header[0]
|
||||
for dropped in ("NCBI_tax_id", "clade_taxid", "taxid"):
|
||||
if dropped in frame.columns:
|
||||
frame = frame.drop(columns=[dropped])
|
||||
|
||||
frame[clade_col] = frame[clade_col].map(normalise_lineage)
|
||||
frame = frame[frame[clade_col].map(deepest_rank) == rank]
|
||||
if frame.empty:
|
||||
raise ValueError(f"{path}: no rows at rank {rank!r}")
|
||||
|
||||
frame = frame.set_index(clade_col)
|
||||
matrix = frame.apply(pd.to_numeric, errors="coerce").fillna(0.0).T
|
||||
matrix.index.name = "sample_id"
|
||||
return matrix
|
||||
|
||||
|
||||
def parse_kraken_report(path: Path, rank: str) -> pd.Series:
|
||||
"""Parse one Kraken2 report into a lineage -> clade-read-count Series."""
|
||||
prefix = RANK_PREFIXES[rank]
|
||||
stack: list[tuple[int, str]] = [] # (indent depth, prefixed name)
|
||||
counts: dict[str, float] = {}
|
||||
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
if not line.strip() or line.startswith("#"):
|
||||
continue
|
||||
fields = line.rstrip("\n").split("\t")
|
||||
if len(fields) < 6:
|
||||
continue
|
||||
|
||||
# Locate the rank-code column: layouts differ between plain reports
|
||||
# (6 columns) and --report-minimizer-data (8 columns).
|
||||
code_idx = next(
|
||||
(
|
||||
i
|
||||
for i, f in enumerate(fields[:-2])
|
||||
if _RANK_CODE_RE.match(f.strip())
|
||||
),
|
||||
None,
|
||||
)
|
||||
if code_idx is None:
|
||||
continue
|
||||
|
||||
code = fields[code_idx].strip()
|
||||
name_field = "\t".join(fields[code_idx + 2 :])
|
||||
name = name_field.strip()
|
||||
if not name or name.lower() in _UNASSIGNED:
|
||||
continue
|
||||
|
||||
depth = (len(name_field) - len(name_field.lstrip(" "))) // 2
|
||||
try:
|
||||
clade_reads = float(fields[1])
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
stack = [entry for entry in stack if entry[0] < depth]
|
||||
|
||||
base_code = code[0]
|
||||
if len(code) > 1 or base_code not in KRAKEN_RANK_CODES:
|
||||
continue # sub-rank (D1, S1, ...) or U/R: keeps depth, no token
|
||||
stack.append((depth, KRAKEN_RANK_CODES[base_code] + name))
|
||||
|
||||
if KRAKEN_RANK_CODES[base_code] == prefix:
|
||||
lineage = "; ".join(entry[1] for entry in stack)
|
||||
counts[lineage] = counts.get(lineage, 0.0) + clade_reads
|
||||
|
||||
return pd.Series(counts, dtype=float)
|
||||
|
||||
|
||||
def parse_kraken(paths: list[Path], rank: str) -> pd.DataFrame:
|
||||
"""Parse many Kraken2 reports into a samples x lineage matrix."""
|
||||
per_sample = {}
|
||||
for path in paths:
|
||||
series = parse_kraken_report(path, rank)
|
||||
if series.empty:
|
||||
print(f"warning: {path} has no rows at rank {rank!r}", file=sys.stderr)
|
||||
per_sample[path.stem] = series
|
||||
matrix = pd.DataFrame(per_sample).T.fillna(0.0)
|
||||
matrix.index.name = "sample_id"
|
||||
return matrix
|
||||
|
||||
|
||||
def parse_table(
|
||||
path: Path,
|
||||
*,
|
||||
taxonomy_column: str | None,
|
||||
orientation: str,
|
||||
) -> pd.DataFrame:
|
||||
"""Parse a QIIME 2 / biom / generic delimited table into samples x lineage."""
|
||||
sep = "," if path.suffix.lower() == ".csv" else "\t"
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
first = handle.readline()
|
||||
skiprows = 1 if first.startswith("# Constructed from biom file") else 0
|
||||
|
||||
frame = pd.read_csv(path, sep=sep, skiprows=skiprows)
|
||||
frame.columns = [str(c).lstrip("#").strip() for c in frame.columns]
|
||||
|
||||
if taxonomy_column and taxonomy_column in frame.columns:
|
||||
lineages = frame[taxonomy_column].astype(str).map(normalise_lineage)
|
||||
feature_col = frame.columns[0]
|
||||
drop = {taxonomy_column, feature_col}
|
||||
values = frame.drop(columns=[c for c in frame.columns if c in drop])
|
||||
values = values.apply(pd.to_numeric, errors="coerce").fillna(0.0)
|
||||
values.index = lineages
|
||||
matrix = values.T
|
||||
else:
|
||||
label_col = frame.columns[0]
|
||||
indexed = frame.set_index(label_col)
|
||||
numeric = indexed.apply(pd.to_numeric, errors="coerce").fillna(0.0)
|
||||
taxa_as_rows = orientation == "taxa_as_rows" or (
|
||||
orientation == "auto"
|
||||
and str(label_col).lower()
|
||||
in {"taxonomy", "lineage", "taxon", "otu", "otu id", "feature id"}
|
||||
)
|
||||
if taxa_as_rows:
|
||||
numeric.index = [normalise_lineage(i) for i in numeric.index]
|
||||
matrix = numeric.T
|
||||
else:
|
||||
numeric.columns = [normalise_lineage(c) for c in numeric.columns]
|
||||
matrix = numeric
|
||||
|
||||
keep = [c for c in matrix.columns if str(c).strip()]
|
||||
matrix = matrix.loc[:, keep]
|
||||
matrix.index = matrix.index.astype(str)
|
||||
matrix.index.name = "sample_id"
|
||||
return matrix
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Matrix -> waypoint format
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def matrix_to_waypoint(
|
||||
matrix: pd.DataFrame,
|
||||
*,
|
||||
normalize: bool = True,
|
||||
drop_zeros: bool = True,
|
||||
min_abundance: float = 0.0,
|
||||
) -> pd.DataFrame:
|
||||
"""Convert a samples x lineage matrix into waypoint format.
|
||||
|
||||
Duplicate lineage columns are summed first: MetaPhlAn and Kraken can both
|
||||
produce the same normalised lineage from different rows.
|
||||
"""
|
||||
if matrix.empty:
|
||||
raise ValueError("abundance matrix is empty")
|
||||
|
||||
values = matrix.astype(float)
|
||||
if values.columns.duplicated().any():
|
||||
values = values.T.groupby(level=0).sum().T
|
||||
|
||||
if normalize:
|
||||
totals = values.sum(axis=1)
|
||||
empty = totals == 0
|
||||
if empty.any():
|
||||
names = ", ".join(map(str, values.index[empty][:5]))
|
||||
raise ValueError(f"samples with zero total abundance: {names}")
|
||||
values = values.div(totals, axis=0)
|
||||
|
||||
taxa: list[list[str]] = []
|
||||
abundances: list[list[float]] = []
|
||||
columns = list(values.columns)
|
||||
for _, row in values.iterrows():
|
||||
pairs = [
|
||||
(str(col), float(val))
|
||||
for col, val in zip(columns, row.to_numpy())
|
||||
if (not drop_zeros or val > 0) and val >= min_abundance
|
||||
]
|
||||
taxa.append([t for t, _ in pairs])
|
||||
abundances.append([a for _, a in pairs])
|
||||
|
||||
out = pd.DataFrame(
|
||||
{"Taxa": taxa, "Relative Abundances": abundances},
|
||||
index=values.index,
|
||||
)
|
||||
out.index.name = values.index.name or "sample_id"
|
||||
return out
|
||||
|
||||
|
||||
def attach_metadata(frame: pd.DataFrame, metadata_path: Path) -> pd.DataFrame:
|
||||
"""Join per-sample metadata, indexed by sample ID, onto a waypoint frame."""
|
||||
suffix = metadata_path.suffix.lower()
|
||||
if suffix == ".parquet":
|
||||
meta = pd.read_parquet(metadata_path)
|
||||
else:
|
||||
meta = pd.read_csv(metadata_path, sep="\t" if suffix in {".tsv", ".tab"} else ",")
|
||||
if meta.index.name is None or meta.index.dtype != object:
|
||||
meta = meta.set_index(meta.columns[0])
|
||||
meta.index = meta.index.astype(str)
|
||||
|
||||
joined = frame.join(meta, how="left")
|
||||
missing = int(joined[meta.columns[0]].isna().sum()) if len(meta.columns) else 0
|
||||
if missing:
|
||||
print(
|
||||
f"warning: {missing}/{len(joined)} samples had no metadata match; "
|
||||
"check that sample IDs agree",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return joined
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_matrix(args: argparse.Namespace, paths: list[Path]) -> pd.DataFrame:
|
||||
if args.format == "metaphlan":
|
||||
frames = [parse_metaphlan(p, args.rank) for p in paths]
|
||||
return pd.concat(frames) if len(frames) > 1 else frames[0]
|
||||
if args.format == "kraken":
|
||||
return parse_kraken(paths, args.rank)
|
||||
# qiime2 and generic share the delimited-table reader.
|
||||
frames = [
|
||||
parse_table(
|
||||
p,
|
||||
taxonomy_column=args.taxonomy_column,
|
||||
orientation=args.orientation,
|
||||
)
|
||||
for p in paths
|
||||
]
|
||||
return pd.concat(frames) if len(frames) > 1 else frames[0]
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert profiler output into waypoint format.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
nargs="+",
|
||||
required=True,
|
||||
help="Input file(s). Kraken takes one report per sample; the others take one table.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--format",
|
||||
required=True,
|
||||
choices=["metaphlan", "kraken", "qiime2", "generic"],
|
||||
help="Input layout.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=True,
|
||||
help="Output path (.parquet recommended; .csv and .tsv also supported).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rank",
|
||||
default="species",
|
||||
choices=list(RANK_PREFIXES),
|
||||
help="Deepest rank to keep. Used by metaphlan and kraken to pick leaf rows.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--taxonomy-column",
|
||||
default=None,
|
||||
help="For qiime2/generic: column holding the lineage string (e.g. 'taxonomy').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--orientation",
|
||||
default="auto",
|
||||
choices=["auto", "samples_as_rows", "taxa_as_rows"],
|
||||
help="For qiime2/generic tables without a taxonomy column.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--metadata",
|
||||
default=None,
|
||||
help="CSV/TSV/parquet of per-sample metadata, indexed by sample ID, merged as extra columns.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-abundance",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Drop taxa below this relative abundance (Atlas used 1e-4).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-normalize",
|
||||
action="store_true",
|
||||
help="Skip row-normalisation (use when values are already relative abundances summing to 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keep-zeros",
|
||||
action="store_true",
|
||||
help="Keep zero-abundance entries in each sample's lists.",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
paths = [Path(p) for p in args.input]
|
||||
for path in paths:
|
||||
if not path.exists():
|
||||
parser.error(f"input not found: {path}")
|
||||
if args.format != "kraken" and len(paths) > 1:
|
||||
print(
|
||||
f"note: concatenating {len(paths)} {args.format} tables by sample",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
matrix = build_matrix(args, paths)
|
||||
frame = matrix_to_waypoint(
|
||||
matrix,
|
||||
normalize=not args.no_normalize,
|
||||
drop_zeros=not args.keep_zeros,
|
||||
min_abundance=args.min_abundance,
|
||||
)
|
||||
|
||||
if args.metadata:
|
||||
frame = attach_metadata(frame, Path(args.metadata))
|
||||
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
suffix = out_path.suffix.lower()
|
||||
if suffix == ".parquet":
|
||||
frame.to_parquet(out_path)
|
||||
elif suffix in {".csv", ".tsv", ".tab"}:
|
||||
frame.to_csv(out_path, sep="\t" if suffix in {".tsv", ".tab"} else ",")
|
||||
else:
|
||||
parser.error(f"unsupported output format: {suffix!r}")
|
||||
|
||||
n_taxa = [len(t) for t in frame["Taxa"]]
|
||||
print(f"Wrote {len(frame)} samples to {out_path}")
|
||||
print(
|
||||
f"Taxa per sample: min {min(n_taxa)}, median {sorted(n_taxa)[len(n_taxa) // 2]}, "
|
||||
f"max {max(n_taxa)}"
|
||||
)
|
||||
if min(n_taxa) < 10:
|
||||
print(
|
||||
"warning: some samples have fewer than 10 taxa; Atlas filtered these out",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
235
skills/waypoint-bio/scripts/vocab_coverage.py
Executable file
235
skills/waypoint-bio/scripts/vocab_coverage.py
Executable file
@@ -0,0 +1,235 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Report how much of a dataset a Waypoint tokenizer can actually see.
|
||||
|
||||
The Waypoint vocabulary is frozen at pretraining time from the Atlas corpus.
|
||||
Taxa outside it resolve to ``<unk>`` and are dropped: ``waypoint embed`` skips
|
||||
them before ordering, and ``filter_unk_taxa: true`` does the same during
|
||||
fine-tuning and benchmarking. Neither warns you. The upstream paper names this
|
||||
as the models' main limitation, so measure it before trusting a downstream
|
||||
number.
|
||||
|
||||
Two coverage figures are reported per sample:
|
||||
|
||||
* **taxon coverage** -- fraction of a sample's taxa that map to a real token.
|
||||
* **abundance coverage** -- fraction of a sample's *relative abundance* carried
|
||||
by those taxa. This is the one that matters: losing 40% of taxa that together
|
||||
account for 2% of the community is fine; losing the dominant genus is not.
|
||||
|
||||
Examples
|
||||
--------
|
||||
python vocab_coverage.py --model outpost-bio/Waypoint-6m --data dataset.parquet
|
||||
python vocab_coverage.py --model outputs/pretrain/best_model --data dataset.parquet \
|
||||
--report-missing 20 --output coverage.csv
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def _parse_list_cell(value):
|
||||
"""Parse one CSV cell back into a list.
|
||||
|
||||
``ast.literal_eval`` only evaluates Python literals and cannot execute
|
||||
code, but it still raises on malformed or pathologically nested input.
|
||||
Turn that into a readable message instead of a traceback.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
try:
|
||||
return ast.literal_eval(value)
|
||||
except (ValueError, SyntaxError, MemoryError, RecursionError) as exc:
|
||||
raise ValueError(
|
||||
f"could not parse list cell {value[:60]!r}: {exc}. "
|
||||
"Waypoint-format CSV stores Taxa and Relative Abundances as Python "
|
||||
"list reprs -- prefer .parquet, which avoids this round-trip."
|
||||
) from exc
|
||||
|
||||
|
||||
def load_dataframe(path: Path) -> pd.DataFrame:
|
||||
"""Read a waypoint-format file, parsing list columns back from CSV/TSV."""
|
||||
suffix = path.suffix.lower()
|
||||
if suffix == ".parquet":
|
||||
frame = pd.read_parquet(path)
|
||||
elif suffix in {".csv", ".tsv", ".tab"}:
|
||||
sep = "\t" if suffix in {".tsv", ".tab"} else ","
|
||||
frame = pd.read_csv(path, sep=sep)
|
||||
# Lists round-trip through CSV as their repr. Check each value rather
|
||||
# than the column dtype: pandas 3 uses a dedicated string dtype, so a
|
||||
# ``dtype == object`` guard silently skips the parse.
|
||||
for column in ("Taxa", "Relative Abundances"):
|
||||
if column in frame.columns:
|
||||
frame[column] = frame[column].map(_parse_list_cell)
|
||||
else:
|
||||
raise ValueError(f"unsupported format {suffix!r}; use .parquet, .csv, or .tsv")
|
||||
|
||||
for column in ("Taxa", "Relative Abundances"):
|
||||
if column not in frame.columns:
|
||||
raise ValueError(f"{path} is not waypoint format: missing {column!r} column")
|
||||
return frame
|
||||
|
||||
|
||||
def load_tokenizer(model: str):
|
||||
"""Load a Waypoint tokenizer from a Hub id or local checkpoint directory.
|
||||
|
||||
Imported lazily so ``--help`` works without transformers installed.
|
||||
"""
|
||||
try:
|
||||
from waypoint_bio.tokenizer import load_tokenizer as _load
|
||||
except ImportError:
|
||||
pass
|
||||
else:
|
||||
return _load(model)
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer
|
||||
except ImportError as exc: # pragma: no cover - environment dependent
|
||||
raise SystemExit(
|
||||
"Neither waypoint-bio nor transformers is installed. "
|
||||
"Install with: pip install waypoint-bio"
|
||||
) from exc
|
||||
return AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
|
||||
|
||||
def coverage_report(
|
||||
frame: pd.DataFrame, tokenizer
|
||||
) -> tuple[pd.DataFrame, Counter[str]]:
|
||||
"""Per-sample taxon and abundance coverage, plus a missing-taxon counter."""
|
||||
unk_id = tokenizer.unk_token_id
|
||||
missing: Counter[str] = Counter()
|
||||
rows = []
|
||||
|
||||
for sample_id, row in frame.iterrows():
|
||||
taxa = row["Taxa"]
|
||||
abundances = row["Relative Abundances"]
|
||||
if not hasattr(taxa, "__iter__") or isinstance(taxa, str):
|
||||
taxa, abundances = [], []
|
||||
|
||||
n_total = len(taxa)
|
||||
n_known = 0
|
||||
abundance_total = 0.0
|
||||
abundance_known = 0.0
|
||||
|
||||
for taxon, abundance in zip(taxa, abundances):
|
||||
value = float(abundance)
|
||||
abundance_total += value
|
||||
if tokenizer.convert_tokens_to_ids(str(taxon)) == unk_id:
|
||||
missing[str(taxon)] += 1
|
||||
else:
|
||||
n_known += 1
|
||||
abundance_known += value
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"sample_id": sample_id,
|
||||
"n_taxa": n_total,
|
||||
"n_in_vocab": n_known,
|
||||
"taxon_coverage": n_known / n_total if n_total else 0.0,
|
||||
"abundance_coverage": (
|
||||
abundance_known / abundance_total if abundance_total else 0.0
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
return pd.DataFrame(rows).set_index("sample_id"), missing
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Report Waypoint tokenizer vocabulary coverage for a dataset.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="outpost-bio/Waypoint-6m",
|
||||
help="Hub id or local checkpoint directory (gated repos need HF_TOKEN).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data", required=True, help="Waypoint-format .parquet / .csv / .tsv."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--threshold",
|
||||
type=float,
|
||||
default=0.8,
|
||||
help="Flag samples whose abundance coverage falls below this.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report-missing",
|
||||
type=int,
|
||||
default=15,
|
||||
help="Show this many of the most frequent out-of-vocabulary taxa.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", default=None, help="Optional path to write the per-sample table."
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
data_path = Path(args.data)
|
||||
if not data_path.exists():
|
||||
parser.error(f"data not found: {data_path}")
|
||||
|
||||
frame = load_dataframe(data_path)
|
||||
tokenizer = load_tokenizer(args.model)
|
||||
report, missing = coverage_report(frame, tokenizer)
|
||||
|
||||
taxon = report["taxon_coverage"]
|
||||
abundance = report["abundance_coverage"]
|
||||
print(f"Model: {args.model} (vocab size {len(tokenizer.get_vocab())})")
|
||||
print(f"Dataset: {data_path} ({len(report)} samples)")
|
||||
print()
|
||||
print(f"{'':22}{'median':>10}{'mean':>10}{'min':>10}")
|
||||
print(
|
||||
f"{'taxon coverage':22}{taxon.median():>10.3f}"
|
||||
f"{taxon.mean():>10.3f}{taxon.min():>10.3f}"
|
||||
)
|
||||
print(
|
||||
f"{'abundance coverage':22}{abundance.median():>10.3f}"
|
||||
f"{abundance.mean():>10.3f}{abundance.min():>10.3f}"
|
||||
)
|
||||
|
||||
below = report[abundance < args.threshold]
|
||||
empty = report[report["n_in_vocab"] == 0]
|
||||
print()
|
||||
print(
|
||||
f"{len(below)}/{len(report)} samples below the "
|
||||
f"{args.threshold:.2f} abundance-coverage threshold"
|
||||
)
|
||||
if len(empty):
|
||||
print(
|
||||
f"{len(empty)} samples have NO in-vocabulary taxa -- these encode as "
|
||||
"[BOS][EOS] and their embeddings are meaningless",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if missing and args.report_missing > 0:
|
||||
print()
|
||||
print(f"Most frequent out-of-vocabulary taxa ({len(missing)} distinct):")
|
||||
for taxon_name, count in missing.most_common(args.report_missing):
|
||||
print(f" {count:>6} {taxon_name}")
|
||||
print()
|
||||
print(
|
||||
"Common causes: a different taxonomy database (SILVA/GTDB vs NCBI naming), "
|
||||
"a 'd__' domain prefix the tokenizer ignores, '|' separators instead of ';', "
|
||||
"or genuinely novel taxa."
|
||||
)
|
||||
|
||||
if args.output:
|
||||
out = Path(args.output)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
if out.suffix.lower() == ".parquet":
|
||||
report.to_parquet(out)
|
||||
else:
|
||||
report.to_csv(out)
|
||||
print(f"\nWrote per-sample coverage to {out}")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -483,6 +483,12 @@ packages = ["scikit-bio", "biom-format", "tables", "matplotlib", "numpy", "panda
|
||||
[skills.polars-bio]
|
||||
packages = ["polars-bio", "polars", "pandas", "bioframe"]
|
||||
|
||||
# waypoint-bio's own package pulls torch + transformers + datasets, which the
|
||||
# bundled scripts deliberately do not require: the converter is pandas-only and
|
||||
# vocab_coverage.py lazy-imports the tokenizer so --help works without it.
|
||||
[skills.waypoint-bio]
|
||||
packages = ["pandas", "pyarrow"]
|
||||
|
||||
[skills.gtars]
|
||||
packages = ["gtars"]
|
||||
|
||||
|
||||
3
tests/waypoint-bio/fixtures/labels.csv
Normal file
3
tests/waypoint-bio/fixtures/labels.csv
Normal file
@@ -0,0 +1,3 @@
|
||||
sample_id,Group
|
||||
sampleA,Case
|
||||
sampleB,Control
|
||||
|
9
tests/waypoint-bio/fixtures/metaphlan_merged.tsv
Normal file
9
tests/waypoint-bio/fixtures/metaphlan_merged.tsv
Normal file
@@ -0,0 +1,9 @@
|
||||
#mpa_vJan21_CHOCOPhlAnSGB_202103
|
||||
#clade_name NCBI_tax_id sampleA sampleB
|
||||
k__Bacteria 2 100.0 100.0
|
||||
k__Bacteria|p__Firmicutes 2|1239 60.0 30.0
|
||||
k__Bacteria|p__Firmicutes|c__Bacilli|o__Lactobacillales|f__Lactobacillaceae|g__Lactobacillus 2|1239|91061|186826|33958|1578 60.0 30.0
|
||||
k__Bacteria|p__Firmicutes|c__Bacilli|o__Lactobacillales|f__Lactobacillaceae|g__Lactobacillus|s__Lactobacillus_gasseri 2|1239|91061|186826|33958|1578|1596 60.0 30.0
|
||||
k__Bacteria|p__Bacteroidota 2|976 40.0 70.0
|
||||
k__Bacteria|p__Bacteroidota|c__Bacteroidia|o__Bacteroidales|f__Bacteroidaceae|g__Bacteroides 2|976|200643|171549|815|816 40.0 70.0
|
||||
k__Bacteria|p__Bacteroidota|c__Bacteroidia|o__Bacteroidales|f__Bacteroidaceae|g__Bacteroides|s__Bacteroides_uniformis 2|976|200643|171549|815|816|820 40.0 70.0
|
||||
|
Can't render this file because it has a wrong number of fields in line 2.
|
5
tests/waypoint-bio/fixtures/qiime2_table.tsv
Normal file
5
tests/waypoint-bio/fixtures/qiime2_table.tsv
Normal file
@@ -0,0 +1,5 @@
|
||||
# Constructed from biom file
|
||||
#OTU ID sampleA sampleB taxonomy
|
||||
f1 60 30 d__Bacteria; p__Firmicutes; c__Bacilli; o__Lactobacillales; f__Lactobacillaceae; g__Lactobacillus
|
||||
f2 40 70 d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Bacteroidales; f__Bacteroidaceae; g__Bacteroides
|
||||
f3 0 0 Unassigned
|
||||
|
Can't render this file because it has a wrong number of fields in line 2.
|
16
tests/waypoint-bio/fixtures/sampleA.kreport
Normal file
16
tests/waypoint-bio/fixtures/sampleA.kreport
Normal file
@@ -0,0 +1,16 @@
|
||||
5.00 500 500 U 0 unclassified
|
||||
95.00 9500 0 R 1 root
|
||||
95.00 9500 10 D 2 Bacteria
|
||||
60.00 6000 5 P 1239 Firmicutes
|
||||
60.00 6000 5 C 91061 Bacilli
|
||||
60.00 6000 5 O 186826 Lactobacillales
|
||||
60.00 6000 5 F 33958 Lactobacillaceae
|
||||
60.00 6000 10 G 1578 Lactobacillus
|
||||
55.00 5500 5500 S 1596 Lactobacillus gasseri
|
||||
5.00 500 500 S1 1597 Lactobacillus gasseri strain X
|
||||
35.00 3500 5 P 976 Bacteroidota
|
||||
35.00 3500 5 C 200643 Bacteroidia
|
||||
35.00 3500 5 O 171549 Bacteroidales
|
||||
35.00 3500 5 F 815 Bacteroidaceae
|
||||
35.00 3500 10 G 816 Bacteroides
|
||||
35.00 3500 3500 S 820 Bacteroides uniformis
|
||||
14
tests/waypoint-bio/fixtures/sampleB.kreport
Normal file
14
tests/waypoint-bio/fixtures/sampleB.kreport
Normal file
@@ -0,0 +1,14 @@
|
||||
100.00 8000 0 R 1 root
|
||||
100.00 8000 10 D 2 Bacteria
|
||||
25.00 2000 5 P 1239 Firmicutes
|
||||
25.00 2000 5 C 91061 Bacilli
|
||||
25.00 2000 5 O 186826 Lactobacillales
|
||||
25.00 2000 5 F 33958 Lactobacillaceae
|
||||
25.00 2000 10 G 1578 Lactobacillus
|
||||
25.00 2000 2000 S 1596 Lactobacillus gasseri
|
||||
75.00 6000 5 P 976 Bacteroidota
|
||||
75.00 6000 5 C 200643 Bacteroidia
|
||||
75.00 6000 5 O 171549 Bacteroidales
|
||||
75.00 6000 5 F 815 Bacteroidaceae
|
||||
75.00 6000 10 G 816 Bacteroides
|
||||
75.00 6000 6000 S 820 Bacteroides uniformis
|
||||
335
tests/waypoint-bio/test_scripts.py
Normal file
335
tests/waypoint-bio/test_scripts.py
Normal file
@@ -0,0 +1,335 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for the waypoint-bio skill's bundled scripts.
|
||||
|
||||
The conversion tests run against synthetic MetaPhlAn / Kraken2 / QIIME 2
|
||||
fixtures and need only pandas. The coverage tests use a stub tokenizer, so no
|
||||
model download and no Hugging Face access is required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
SKILL_ROOT = Path(__file__).resolve().parents[2] / "skills" / "waypoint-bio"
|
||||
SCRIPTS_DIR = SKILL_ROOT / "scripts"
|
||||
FIXTURES = Path(__file__).resolve().parent / "fixtures"
|
||||
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
import skill_contract # noqa: E402
|
||||
|
||||
pd = None
|
||||
try:
|
||||
import pandas as pd # noqa: E402
|
||||
except ImportError: # pragma: no cover - environment dependent
|
||||
pd = None
|
||||
|
||||
if pd is not None:
|
||||
import profiler_to_waypoint as converter # noqa: E402
|
||||
import vocab_coverage # noqa: E402
|
||||
|
||||
|
||||
CliHelpTests = skill_contract.cli.help_test_case(SKILL_ROOT)
|
||||
|
||||
requires_pandas = unittest.skipIf(pd is None, "pandas not installed")
|
||||
|
||||
|
||||
def run_script(name: str, *arguments: str) -> subprocess.CompletedProcess[str]:
|
||||
environment = os.environ.copy()
|
||||
environment["PYTHONDONTWRITEBYTECODE"] = "1"
|
||||
return subprocess.run(
|
||||
[sys.executable, str(SCRIPTS_DIR / name), *arguments],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
env=environment,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
|
||||
LACTOBACILLUS = (
|
||||
"k__Bacteria; p__Firmicutes; c__Bacilli; o__Lactobacillales; "
|
||||
"f__Lactobacillaceae; g__Lactobacillus"
|
||||
)
|
||||
BACTEROIDES = (
|
||||
"k__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Bacteroidales; "
|
||||
"f__Bacteroidaceae; g__Bacteroides"
|
||||
)
|
||||
|
||||
|
||||
class LineageNormalisationTests(unittest.TestCase):
|
||||
"""The tokenizer only reads k__/p__/c__/o__/f__/g__/s__ segments split on ';'."""
|
||||
|
||||
@requires_pandas
|
||||
def test_pipe_separator_becomes_semicolon(self) -> None:
|
||||
self.assertEqual(
|
||||
converter.normalise_lineage("k__Bacteria|p__Firmicutes|g__Lactobacillus"),
|
||||
"k__Bacteria; p__Firmicutes; g__Lactobacillus",
|
||||
)
|
||||
|
||||
@requires_pandas
|
||||
def test_domain_prefix_is_rewritten_to_kingdom(self) -> None:
|
||||
# SILVA/Greengenes2 write d__; the tokenizer would skip it entirely.
|
||||
self.assertEqual(
|
||||
converter.normalise_lineage("d__Bacteria; p__Firmicutes"),
|
||||
"k__Bacteria; p__Firmicutes",
|
||||
)
|
||||
|
||||
@requires_pandas
|
||||
def test_strain_and_empty_segments_are_dropped(self) -> None:
|
||||
self.assertEqual(
|
||||
converter.normalise_lineage("k__Bacteria; g__Lactobacillus; s__; t__SGB1"),
|
||||
"k__Bacteria; g__Lactobacillus",
|
||||
)
|
||||
|
||||
@requires_pandas
|
||||
def test_deepest_rank(self) -> None:
|
||||
self.assertEqual(converter.deepest_rank(LACTOBACILLUS), "genus")
|
||||
self.assertEqual(
|
||||
converter.deepest_rank(LACTOBACILLUS + "; s__Lactobacillus gasseri"),
|
||||
"species",
|
||||
)
|
||||
self.assertIsNone(converter.deepest_rank("root; cellular organisms"))
|
||||
|
||||
|
||||
class MetaphlanConversionTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_species_rows_only_and_renormalised(self) -> None:
|
||||
matrix = converter.parse_metaphlan(
|
||||
FIXTURES / "metaphlan_merged.tsv", rank="species"
|
||||
)
|
||||
# The cumulative parent rows must not be double counted.
|
||||
self.assertEqual(len(matrix.columns), 2)
|
||||
self.assertEqual(sorted(matrix.index), ["sampleA", "sampleB"])
|
||||
|
||||
frame = converter.matrix_to_waypoint(matrix)
|
||||
self.assertAlmostEqual(sum(frame.loc["sampleA", "Relative Abundances"]), 1.0)
|
||||
self.assertAlmostEqual(sum(frame.loc["sampleB", "Relative Abundances"]), 1.0)
|
||||
|
||||
taxa = frame.loc["sampleA", "Taxa"]
|
||||
abundances = dict(zip(taxa, frame.loc["sampleA", "Relative Abundances"]))
|
||||
gasseri = LACTOBACILLUS + "; s__Lactobacillus_gasseri"
|
||||
self.assertIn(gasseri, abundances)
|
||||
self.assertAlmostEqual(abundances[gasseri], 0.6)
|
||||
self.assertTrue(all("|" not in t for t in taxa))
|
||||
|
||||
@requires_pandas
|
||||
def test_genus_rank_selection(self) -> None:
|
||||
matrix = converter.parse_metaphlan(
|
||||
FIXTURES / "metaphlan_merged.tsv", rank="genus"
|
||||
)
|
||||
self.assertIn(LACTOBACILLUS, matrix.columns)
|
||||
self.assertIn(BACTEROIDES, matrix.columns)
|
||||
|
||||
|
||||
class KrakenConversionTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_lineage_rebuilt_from_indentation(self) -> None:
|
||||
matrix = converter.parse_kraken(
|
||||
[FIXTURES / "sampleA.kreport", FIXTURES / "sampleB.kreport"],
|
||||
rank="species",
|
||||
)
|
||||
self.assertEqual(sorted(matrix.index), ["sampleA", "sampleB"])
|
||||
gasseri = LACTOBACILLUS + "; s__Lactobacillus gasseri"
|
||||
uniformis = BACTEROIDES + "; s__Bacteroides uniformis"
|
||||
self.assertIn(gasseri, matrix.columns)
|
||||
self.assertIn(uniformis, matrix.columns)
|
||||
|
||||
@requires_pandas
|
||||
def test_subranks_and_unclassified_excluded(self) -> None:
|
||||
series = converter.parse_kraken_report(
|
||||
FIXTURES / "sampleA.kreport", rank="species"
|
||||
)
|
||||
# S1 strain row and the U unclassified row must not appear.
|
||||
self.assertTrue(all("strain X" not in name for name in series.index))
|
||||
self.assertTrue(all("unclassified" not in name for name in series.index))
|
||||
self.assertEqual(len(series), 2)
|
||||
|
||||
@requires_pandas
|
||||
def test_counts_normalise_to_relative_abundance(self) -> None:
|
||||
matrix = converter.parse_kraken(
|
||||
[FIXTURES / "sampleB.kreport"], rank="species"
|
||||
)
|
||||
frame = converter.matrix_to_waypoint(matrix)
|
||||
abundances = dict(
|
||||
zip(frame.loc["sampleB", "Taxa"], frame.loc["sampleB", "Relative Abundances"])
|
||||
)
|
||||
uniformis = BACTEROIDES + "; s__Bacteroides uniformis"
|
||||
self.assertAlmostEqual(abundances[uniformis], 6000 / 8000)
|
||||
|
||||
|
||||
class QiimeConversionTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_biom_banner_taxonomy_column_and_unassigned(self) -> None:
|
||||
matrix = converter.parse_table(
|
||||
FIXTURES / "qiime2_table.tsv",
|
||||
taxonomy_column="taxonomy",
|
||||
orientation="auto",
|
||||
)
|
||||
self.assertEqual(sorted(matrix.index), ["sampleA", "sampleB"])
|
||||
self.assertIn(LACTOBACILLUS, matrix.columns)
|
||||
self.assertIn(BACTEROIDES, matrix.columns)
|
||||
# "Unassigned" has no rank prefix, so it normalises to an empty string
|
||||
# and is dropped rather than becoming a bogus token.
|
||||
self.assertNotIn("Unassigned", matrix.columns)
|
||||
|
||||
frame = converter.matrix_to_waypoint(matrix)
|
||||
abundances = dict(
|
||||
zip(frame.loc["sampleA", "Taxa"], frame.loc["sampleA", "Relative Abundances"])
|
||||
)
|
||||
self.assertAlmostEqual(abundances[LACTOBACILLUS], 0.6)
|
||||
|
||||
|
||||
class WaypointFormatTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_duplicate_lineages_are_summed(self) -> None:
|
||||
matrix = pd.DataFrame(
|
||||
[[1.0, 2.0, 1.0]],
|
||||
index=pd.Index(["s1"], name="sample_id"),
|
||||
columns=[LACTOBACILLUS, BACTEROIDES, LACTOBACILLUS],
|
||||
)
|
||||
frame = converter.matrix_to_waypoint(matrix)
|
||||
abundances = dict(
|
||||
zip(frame.loc["s1", "Taxa"], frame.loc["s1", "Relative Abundances"])
|
||||
)
|
||||
self.assertEqual(len(abundances), 2)
|
||||
self.assertAlmostEqual(abundances[LACTOBACILLUS], 0.5)
|
||||
|
||||
@requires_pandas
|
||||
def test_zero_total_sample_is_rejected(self) -> None:
|
||||
matrix = pd.DataFrame(
|
||||
[[0.0, 0.0]], index=["s1"], columns=[LACTOBACILLUS, BACTEROIDES]
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
converter.matrix_to_waypoint(matrix)
|
||||
|
||||
@requires_pandas
|
||||
def test_min_abundance_filter(self) -> None:
|
||||
matrix = pd.DataFrame(
|
||||
[[0.99999, 0.00001]], index=["s1"], columns=[LACTOBACILLUS, BACTEROIDES]
|
||||
)
|
||||
frame = converter.matrix_to_waypoint(matrix, min_abundance=1e-4)
|
||||
self.assertEqual(frame.loc["s1", "Taxa"], [LACTOBACILLUS])
|
||||
|
||||
|
||||
class ConverterCliTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_metaphlan_end_to_end_with_metadata(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
output = Path(tmp) / "dataset.csv"
|
||||
result = run_script(
|
||||
"profiler_to_waypoint.py",
|
||||
"--input", str(FIXTURES / "metaphlan_merged.tsv"),
|
||||
"--format", "metaphlan",
|
||||
"--metadata", str(FIXTURES / "labels.csv"),
|
||||
"--output", str(output),
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertTrue(output.exists())
|
||||
|
||||
frame = vocab_coverage.load_dataframe(output)
|
||||
self.assertIn("Group", frame.columns)
|
||||
self.assertEqual(sorted(frame["Group"]), ["Case", "Control"])
|
||||
self.assertEqual(len(frame.loc[0, "Taxa"]), 2)
|
||||
|
||||
@requires_pandas
|
||||
def test_kraken_end_to_end_parquet(self) -> None:
|
||||
try:
|
||||
import pyarrow # noqa: F401
|
||||
except ImportError:
|
||||
self.skipTest("pyarrow not installed")
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
output = Path(tmp) / "dataset.parquet"
|
||||
result = run_script(
|
||||
"profiler_to_waypoint.py",
|
||||
"--input",
|
||||
str(FIXTURES / "sampleA.kreport"),
|
||||
str(FIXTURES / "sampleB.kreport"),
|
||||
"--format", "kraken",
|
||||
"--output", str(output),
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
frame = pd.read_parquet(output)
|
||||
self.assertEqual(sorted(frame.index), ["sampleA", "sampleB"])
|
||||
self.assertIn("Taxa", frame.columns)
|
||||
|
||||
def test_missing_input_is_an_error(self) -> None:
|
||||
result = run_script(
|
||||
"profiler_to_waypoint.py",
|
||||
"--input", "does_not_exist.tsv",
|
||||
"--format", "metaphlan",
|
||||
"--output", "out.parquet",
|
||||
)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn("input not found", result.stderr)
|
||||
|
||||
|
||||
class StubTokenizer:
|
||||
"""Minimal stand-in for TaxonomicTokenizer: genus tokens only, no fallback."""
|
||||
|
||||
unk_token_id = 1
|
||||
|
||||
def __init__(self, known: set[str]) -> None:
|
||||
self._vocab = {"<pad>": 0, "<unk>": 1}
|
||||
for index, token in enumerate(sorted(known), start=2):
|
||||
self._vocab[token] = index
|
||||
|
||||
def get_vocab(self) -> dict[str, int]:
|
||||
return dict(self._vocab)
|
||||
|
||||
def convert_tokens_to_ids(self, token: str) -> int:
|
||||
genus = next(
|
||||
(
|
||||
segment.strip()
|
||||
for segment in str(token).split(";")
|
||||
if segment.strip().startswith("g__")
|
||||
),
|
||||
None,
|
||||
)
|
||||
return self._vocab.get(genus, self.unk_token_id)
|
||||
|
||||
|
||||
class VocabCoverageTests(unittest.TestCase):
|
||||
@requires_pandas
|
||||
def test_coverage_weights_by_abundance(self) -> None:
|
||||
frame = pd.DataFrame(
|
||||
{
|
||||
"Taxa": [[LACTOBACILLUS, BACTEROIDES]],
|
||||
"Relative Abundances": [[0.9, 0.1]],
|
||||
},
|
||||
index=pd.Index(["s1"], name="sample_id"),
|
||||
)
|
||||
tokenizer = StubTokenizer({"g__Lactobacillus"})
|
||||
report, missing = vocab_coverage.coverage_report(frame, tokenizer)
|
||||
|
||||
# Half the taxa are known, but they carry 90% of the abundance.
|
||||
self.assertAlmostEqual(report.loc["s1", "taxon_coverage"], 0.5)
|
||||
self.assertAlmostEqual(report.loc["s1", "abundance_coverage"], 0.9)
|
||||
self.assertEqual(missing[BACTEROIDES], 1)
|
||||
|
||||
@requires_pandas
|
||||
def test_sample_with_no_known_taxa(self) -> None:
|
||||
frame = pd.DataFrame(
|
||||
{"Taxa": [[BACTEROIDES]], "Relative Abundances": [[1.0]]},
|
||||
index=pd.Index(["s1"], name="sample_id"),
|
||||
)
|
||||
report, _ = vocab_coverage.coverage_report(frame, StubTokenizer(set()))
|
||||
self.assertEqual(report.loc["s1", "n_in_vocab"], 0)
|
||||
self.assertAlmostEqual(report.loc["s1", "abundance_coverage"], 0.0)
|
||||
|
||||
@requires_pandas
|
||||
def test_non_waypoint_file_is_rejected(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = Path(tmp) / "bad.csv"
|
||||
path.write_text("a,b\n1,2\n", encoding="utf-8")
|
||||
with self.assertRaises(ValueError):
|
||||
vocab_coverage.load_dataframe(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user