Hugging Face's logo

FaiasPromit
/
HaploLM

PyTorch
biology
genomics
dna
mitochondrial-dna
haplogroups
dna-language-model
Model card Files
xet
 Community
HaploLM
File size: 5,207 Bytes
54bab0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
---
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.