# mgxLens A retrieval-based metagenomic taxonomic profiler. Given short DNA reads (~150 bp) from a sequencing run, it reports the estimated **relative abundance** of microbial taxa at the **genus** and **species** level. **Reference coverage.** The index covers **13,399 distinct species-level SGBs across 3,297 genera** — the full MetaPhlAn vJan25 SGB reference. Every indexed SGB resolves to a named Linnaean species (no unnamed placeholder bins among the *indexed* organisms; the naming rate over indexed SGBs is 100%). This means "closed-world" is about organisms *outside MetaPhlAn's reference entirely* — not gaps within it. An organism that is in MetaPhlAn's reference is representable; one that is not will be misassigned to its nearest reference genus (see Caveats → closed-world). --- ## 1. What's in the bundle | File | What it is | |---|---| | `encoder.pt` | The neural encoder (seqLens-89M backbone + attention pooling → 256-d embedding). Turns a DNA string into a vector. | | `index.faiss` | A FAISS vector index of ~2.32M reference marker embeddings. Cosine search via inner product. | | `index.clades.npy` | Row-aligned to the FAISS index: the SGB (clade) id for each indexed vector. | | `index.markers.npy` | Row-aligned marker source ids (provenance; not needed for inference). | | `index.config.json` | Index build parameters (window=150, nwin=5, emb_dim=256, max_length=128). | | `lineage.tsv` | `clade_id → genus, species, display_name`. Turns a retrieved SGB into a taxon name (display_name is MetaPhlAn-style). | | `mgx_encoder.py` | The encoder class + `load_encoder()` / `embed()` helpers. Import this; don't reimplement. | | `MANIFEST.json` | Machine-readable summary + file checksums + caveats. | **Encoder and index are from the same training run** (`index.config.json.source_ckpt` == `encoder.pt`). Do not mix this encoder with any other index or vice-versa — the vectors would be meaningless. --- ## 2. Environment ```bash pip install torch transformers faiss-cpu numpy # GPU optional: faiss-gpu + a CUDA torch build speed up large batches, but faiss-cpu is fine for # moderate throughput. The encoder runs on CPU or GPU; GPU is ~10-50x faster for embedding. ``` The encoder downloads its base model (`omicseye/seqLens_4096_512_89M-at-base-multi`) from HuggingFace on first load via `transformers`, so the serving host needs network access on first run (or pre-cache the HF model). `trust_remote_code=True` is required for that base model. --- ## 3. How it works (the inference recipe) For each input read: 1. **Window** the read to 150 bp (if longer, take a centered 150 bp window; if ≤150, use as-is). 2. **Embed** it with the encoder → a 256-d L2-normalized vector. 3. **Search** the FAISS index for the top-k nearest reference vectors (k=25). Because vectors are L2-normalized and the index is inner-product, the score **is** cosine similarity ∈ [-1, 1]. 4. **Gate**: if the top-1 cosine is below a threshold τ (start with **0.66**), the read is *abstained* (dropped, counted as "unclassified"). Otherwise it's assigned. 5. **Assign**: the read's taxon is the clade (`index.clades.npy`) of its top-1 hit. Map that clade to genus/species via `lineage.tsv`. 6. **Aggregate** across all assigned reads: count reads per genus (and per species), normalize to fractions → the abundance profile. Report the unclassified fraction separately. That's the whole pipeline. It's nearest-neighbor retrieval, not a trained classifier. --- ## 4. Reference implementation (copy-paste starting point) ```python import numpy as np, faiss, torch from collections import defaultdict from mgx_encoder import load_encoder, embed # ships in this bundle BUNDLE = "." # path to the bundle dir K = 25 # neighbors per read GATE = 0.66 # top-1 cosine threshold; below -> abstain WINDOW = 150 # ---- load once at startup ---- device = "cuda" if torch.cuda.is_available() else "cpu" model, tok, cfg = load_encoder(f"{BUNDLE}/encoder.pt", device) index = faiss.read_index(f"{BUNDLE}/index.faiss") clades = np.load(f"{BUNDLE}/index.clades.npy", allow_pickle=True) # (N,) SGB id per vector lineage = {} # clade_id -> (genus, species, display_name) with open(f"{BUNDLE}/lineage.tsv") as fh: next(fh) for line in fh: cid, g, s, disp = line.rstrip("\n").split("\t") lineage[cid] = (g, s, disp) def center_window(seq, w=WINDOW): if len(seq) <= w: return seq s = (len(seq) - w) // 2 return seq[s:s+w] def profile(reads, batch_size=4096): """reads: list of DNA strings. Returns genus/species abundance + unclassified fraction. Species keys are MetaPhlAn-style display names ('Escherichia coli', or 'Escherichia sp. (SGB123)' for unnamed genome bins).""" genus_counts, species_counts = defaultdict(float), defaultdict(float) n_total = len(reads); n_assigned = 0 for i in range(0, n_total, batch_size): batch = [center_window(r) for r in reads[i:i+batch_size]] ml = cfg["max_length"] if "max_length" in cfg else 128 Z = embed(model, tok, batch, device, max_length=ml).astype("float32") D, I = index.search(Z, K) # D=cosine sims, I=index rows for r in range(len(batch)): if D[r, 0] < GATE: # gate: abstain continue clade = str(clades[I[r, 0]]) # top-1 hit's clade g, s, disp = lineage.get(clade, ("", "", "")) if not g: continue n_assigned += 1 genus_counts[g] += 1 if disp: species_counts[disp] += 1 # display name, not raw s__ / SGB id def norm(d): tot = sum(d.values()) or 1.0 return {k: v/tot for k, v in sorted(d.items(), key=lambda x: -x[1])} return { "genus": norm(genus_counts), "species": norm(species_counts), "n_reads": n_total, "n_assigned": n_assigned, "unclassified_fraction": 1.0 - (n_assigned / n_total if n_total else 0.0), } # ---- example ---- if __name__ == "__main__": reads = ["ACGT..."] # your reads (from a FASTQ) print(profile(reads)) ``` **Reading FASTQ**: reads come from `.fastq`/`.fastq.gz`. Use `pysam`, `Bio.SeqIO`, or a 2nd/4th-line parser to get the sequence strings, then pass the list to `profile()`. Paired-end R1/R2 can both be fed as independent reads. --- ## 5. Serving it (FastAPI sketch) ```python from fastapi import FastAPI, UploadFile app = FastAPI() # model/index/lineage loaded once at module import (see section 4) @app.post("/profile") async def profile_endpoint(file: UploadFile): reads = parse_fastq(await file.read()) # your FASTQ parser -> list[str] return profile(reads) # JSON: {genus:{...}, species:{...}, ...} ``` Load the model, index, and lineage **once at startup** (they're large — encoder 342 MB, index 2.3 GB). Never reload per request. The index holds ~2.3 GB in RAM; size the host accordingly (≥8 GB RAM recommended, ≥16 GB comfortable). --- ## 6. Input / output contract **Input**: DNA reads as strings (from FASTQ), ~150 bp each. Non-ACGT characters are tolerated by the tokenizer but degrade the embedding; upstream QC/trimming (e.g. fastp) is assumed. **Output** (JSON): ```json { "genus": {"escherichia": 0.29, "pseudomonas": 0.24, "...": 0.0}, "species": {"Escherichia coli": 0.21, "Pseudomonas aeruginosa": 0.18, "...": 0.0}, "n_reads": 200000, "n_assigned": 94459, "unclassified_fraction": 0.528 } ``` Abundances are fractions summing to ~1.0 within each level (over assigned reads). All 13,399 indexed SGBs have real species names, so species keys are Linnaean names (the `Genus sp. (SGBxxxx)` fallback only appears if the index is rebuilt to include unnamed bins). --- ## 7. Tunable knobs | Knob | Default | Effect | |---|---|---| | `GATE` (top-1 cosine) | 0.66 | Higher = stricter, more reads abstained, fewer false assignments (but see Caveats — it does **not** reliably reject novel organisms). | | `K` (neighbors) | 25 | Only top-1 is used for assignment here; k>1 matters if you switch to k-NN vote aggregation. | | aggregation | top-1 vote | Simplest. Alternatives (k-NN vote, similarity-weighted) exist in the research code but top-1 is the documented default. | --- ## 8. Caveats — read before building product on this - **Proof of concept.** Not accuracy-tuned. The goal of this handoff is to enable infrastructure and integration work, not to ship a validated diagnostic. - **Genus is the reliable level; species is approximate.** At 150 bp, reads from different species in the same genus are near-indistinguishable to the encoder — the model reliably gets the *genus* right but frequently assigns the wrong *species within that genus*. Report species output, but do not treat it as trustworthy. - **Species naming.** All 13,399 indexed SGBs have real Linnaean species names, so normal output shows names like "Escherichia coli", not bin ids. (The `display_name` logic still falls back to `Genus sp. (SGBxxxx)` for unnamed bins — this only matters if the index is ever rebuilt to include the unnamed SGBs that exist elsewhere in the MetaPhlAn taxonomy; the shipped index has none.) - **Closed-world — this is the big one.** The model can only assign a read to a taxon that exists in its reference index. An organism that is **not** in the reference will still be confidently assigned to its nearest reference genus and reported as present. There is effectively **no novelty rejection** (abstain rate ~0% in evaluation), and this behavior was **never quantified**. **Do not deploy this against arbitrary/open real-world samples** and present the output as complete or trustworthy without adding a novelty-rejection layer. It is validated only in the closed-world setting where the sample's organisms are known to be in the reference. - **Evaluated on one mock community (Zymo D6300).** Broader validation was not done. --- ## 9. Provenance - Encoder: `enc_attention_bg1.best_genus.pt` — seqLens-89M backbone, attention pooling, 256-d projection, trained with contrastive InfoNCE + background negatives, checkpoint selected on real-Zymo genus accuracy. - Index: `index_bakeoff_mw_notest` — leakage-fixed (held-out test markers excluded). - Base model pulled at runtime: `omicseye/seqLens_4096_512_89M-at-base-multi`. See `MANIFEST.json` for file checksums and the machine-readable summary.