| --- |
| license: mit |
| library_name: pytorch |
| tags: |
| - biology |
| - genomics |
| - dna |
| - mitochondrial-dna |
| - haplogroups |
| - dna-language-model |
| --- |
| |
| # HaploLM |
|
|
| HaploLM and HaploLM-VL are the two DNA language models of the paper **Human Maternal Family Tree |
| from DNA Language Models** by Faias Satter and Aminul Islam. |
| They are trained only to predict human mitochondrial DNA and never see a haplogroup label, yet |
| they learn the human maternal family tree (Mitotree). |
|
|
| Code, data preparation and every test: [github.com/FaiasPromit/HaploLM](https://github.com/FaiasPromit/HaploLM) |
|
|
| | Folder | Model | Reads | Parameters | Training | |
| |---|---|---|---|---| |
| | `SC30/` | **HaploLM** | the whole mitochondrial sequence, base by base | 20.6M | causal (each base predicted from the bases before it), 30 epochs over 49,310 sequences | |
| | `VC/` | **HaploLM-VL** | only the person's variants (differences from the reference sequence, rCRS), about 37 per person | 11.2M | causal, 20,000 steps of 128 people (about 52 epochs over 49,259 variant lists) | |
|
|
| Each folder holds `checkpoint.pt`, a PyTorch file with the parameters (`model`), the settings |
| (`args`) and the training step (`step`). For HaploLM-VL, `vocab` maps every distinct variant |
| seen in training to its token. The folder names are the run names used in the GitHub repository. |
|
|
| ## Results |
|
|
| | Model | Marker test | Frequency-matched | Carrier vs. relative | Held-out branches | |
| |---|---|---|---|---| |
| | Counting (no model) | .890 | .533 | .500 | .475 | |
| | HaploLM | .947 | .872 | .808 | .467 | |
| | HaploLM-VL | .955 | .939 | .965 | .502 | |
|
|
| - **Marker test**: AUROC for telling a person's branch markers from their private variants, by |
| the probability the model gives each variant. |
| - **Frequency-matched**: the same, with markers compared only with private variants carried by |
| about as many training people. |
| - **Carrier vs. relative** (the main score): the share of pairs in which the model finds a |
| branch marker more probable in a person who carries it than in a close relative who does not; |
| 0.5 is random guessing. |
| - **Held-out branches** (negative control): the marker test on six branches excluded from |
| training. A model that relies on relationships learned during training should score near 0.5 |
| here, and all models do. |
|
|
| Counting scores a variant by how many training people carry it. The published Evo 2, NTv3 and |
| mtDNA-FM score close to random guessing on all four tests. |
|
|
| ## How to use |
|
|
| The model code is in the GitHub repository. To score the paper's tests with these checkpoints, |
| first build the data (step 1 of the repository's README), then: |
|
|
| ```bash |
| python -c "from huggingface_hub import snapshot_download; snapshot_download('FaiasPromit/HaploLM', local_dir='runs', allow_patterns=['SC30/*', 'VC/*'])" |
| torchrun --standalone --nproc_per_node=4 scripts/score_causal_sequence_model.py --run runs/SC30 |
| python scripts/train_causal_variant_list_model.py --name VC --score-only |
| ``` |
|
|
| To load HaploLM in Python and score a variant in a sequence: |
|
|
| ```python |
| import sys |
| |
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
| |
| sys.path.insert(0, "HaploLM/scripts") # the scripts folder of the GitHub repository |
| from train_causal_sequence_model import CausalSequenceModel, next_base_logp |
| from train_masked_sequence_model import BASE_CODE |
| |
| ck = torch.load(hf_hub_download("FaiasPromit/HaploLM", "SC30/checkpoint.pt"), |
| map_location="cpu", weights_only=False) |
| a = ck["args"] |
| model = CausalSequenceModel(a["d"], a["layers"], a["heads"], dropout=0.0) |
| model.load_state_dict(ck["model"]) |
| model.eval() |
| |
| |
| def log_prob(seq): |
| """log P(seq): the sum over bases of log P(base | the bases before it).""" |
| ids = torch.from_numpy(BASE_CODE[np.frombuffer(seq.encode("ascii"), np.uint8)]).long()[None] |
| with torch.no_grad(): |
| lp, _ = next_base_logp(model(ids), ids) |
| return lp.double().sum().item() |
| |
| # score of a variant, as in the paper: |
| # log_prob(sequence with the variant) - log_prob(the same sequence with the rCRS base there) |
| ``` |
|
|
| The sequence must be in upper case, start at position 1 of the rCRS (the repository's |
| `clean_sequences.py` rotates sequences this way) and be at most 16,768 bases long. A GPU is much |
| faster: move `model` and `ids` to `"cuda"`. |
|
|
| ## Training data |
|
|
| 49,310 human mitochondrial genomes from GenBank, on MITOMAP's list of full-length sequences: |
| the training set of our split. Identical sequences are always in the same set, and six |
| branches (L0a1, L3f, M9a, D4b, H13 and U2e), with everything below them, are held out of |
| training. The haplogroup labels, from HaploGrep 3 with Mitotree 2026.05.06, are used only to |
| evaluate the models. |
|
|
| ## Limitations |
|
|
| - The models are trained on human mitochondrial DNA only. |
| - The answer key depends on HaploGrep and Mitotree. The tests limit the effects of errors by |
| excluding poorly assigned sequences, hotspots and the control region. |
| - Public mitochondrial data overrepresent some populations, leaving lineages common in |
| undersampled populations with fewer training examples. |
| - They are research models for studying what DNA language models learn. They are not |
| haplogroup classifiers and are not validated for medical or forensic use. |
|
|
