File size: 3,606 Bytes
aa45001 | 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 | ---
title: Metalorian / MetaLATTE
---
# Repo structure
```text
train.py Train either stage with --stage
sample_direct.py Direct sampling
sample_prototype.py Prototype-seeded sampling
metalorian/ Models and support functions
data/metalatte/ Training and validation split
data/metalorian/ Diffusion sequence-label training and test split
data/MbPA/ Original MbPA CSV exports
checkpoints/ Predictor and diffusion checkpoints
assets/ Local ESM-2 backbone and tokenizer
requirements.txt
```
## Install
Use Python 3.10. Run from this directory:
```bash
python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
```
Weights and tokenizer files are included for offline use after installation.
Use a CUDA-capable PyTorch installation. For CPU checks, pass `--device cpu`.
## Data
| Split | Records |
| --- | ---: |
| `data/metalatte/train` | 50,694 groups |
| `data/metalatte/validation` | 12,729 groups |
| `data/metalorian/train` | 38,123 sequence-label records |
| `data/metalorian/test` | 9,531 sequence-label records |
MetaLATTE groups are stored in 1,852 training pre-batches and 460 validation
pre-batches.
Metalorian uses the separate sequence-label splits listed above.
`data/MbPA/` contains eight original CSV exports, aggregated from 14 different metal labels. File names, row counts and
hashes are listed in `provenance/mbpa_inventory.csv`.
## Train either stage
MetaLATTE, initialized from pretrained ESM-2:
```bash
python train.py --stage metalatte --epochs 30 --device cuda:0 \
--output runs/metalatte
```
To start from the supplied predictor, add
`--init-checkpoint checkpoints/metalatte_predictor.ckpt`.
To resume training and optimizer state, use `--resume PATH`.
Metalorian, initialized from the supplied MetaLATTE predictor:
```bash
python train.py --stage metalorian --epochs 1 --batch-size 2 --device cuda:0 \
--output runs/metalorian
```
Use `--init-checkpoint PATH` for another MetaLATTE checkpoint, or
`--resume checkpoints/metalorian_diffusion.pt` to resume diffusion training.
Here, `--epochs` counts additional passes through the training set.
Outputs are `metalorian_latest.pt` in the output directory.
Short forward/backward checks for both stages:
```bash
python train.py --stage metalatte --smoke-test --device cpu \
--init-checkpoint checkpoints/metalatte_predictor.ckpt --output runs/check_metalatte
python train.py --stage metalorian --smoke-test --batch-size 2 --device cpu \
--output runs/check_metalorian
```
`--smoke-test` uses at most two records and 64 tokens for one update and writes
no model checkpoint. Training uses one device; see `VALIDATION.md` for checks.
## Direct sampling
```bash
python sample_direct.py --target_idx 3 --seq_len 30 --batch 1 \
--desired 1 --max_rounds 20 --output outputs/Cu_L30.csv
```
## Prototype-seeded sampling
```bash
python sample_prototype.py --target_idx 1 --seq_len 30 --batch 1 \
--proto MDPETCPCPSGGSCTCADSCKCEGCKCTSCKKSCCSCCPAECEKCAKDCVCKGGEAAEAEAEKCSCCQ \
--t_start 15 --desired 1 --max_rounds 20 --output outputs/Cd_L30.csv
```
```text
Metal indices: 0 Ag, 1 Cd, 2 Co, 3 Cu, 4 Fe, 5 Hg, 6 Mn, 7 Mo,
8 Ni, 9 Pb, 10 Pt, 11 V, 12 W, 13 Zn, 14 NA
```
Both scripts load `checkpoints/metalatte_predictor.ckpt` and
`checkpoints/metalorian_diffusion.pt`. Use `--help` for sampling options.
The CSV contains unique accepted sequences; the adjacent JSON records settings
and counts. A run may return fewer than `--desired` at `--max_rounds`. |