| """Embed every train + validation protein with ESM-2 on an Apple-silicon GPU (MPS). |
| |
| Run on the Mac: |
| python3 esm_embed.py --model 35M (~10-20 min) |
| python3 esm_embed.py --model 650M (~2-4 h, optional) |
| |
| Writes to ~/enzyme-esm/: |
| ids_train.txt, ids_val.txt sequence ids in row order |
| emb_<model>_train.npy, emb_<model>_val.npy float16, mean-pooled final layer |
| The ESM-2 weights are sequence-only (masked-LM pretraining, no EC labels). |
| """ |
| import argparse, os, sys, time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoTokenizer, EsmModel |
|
|
| MODELS = {'8M': 'facebook/esm2_t6_8M_UR50D', '35M': 'facebook/esm2_t12_35M_UR50D', |
| '150M': 'facebook/esm2_t30_150M_UR50D', '650M': 'facebook/esm2_t33_650M_UR50D'} |
| OUT = Path.home() / 'enzyme-esm' |
| MAX_LEN = 1022 |
| PART = 10000 |
|
|
|
|
| def load_data(): |
| d = OUT / 'data' |
| d.mkdir(parents=True, exist_ok=True) |
| paths = {} |
| for f in ('train.parquet', 'val_inputs.parquet'): |
| paths[f] = hf_hub_download('starkAhmed43/leak-curber-enzyme-classification', f, |
| repo_type='dataset', local_dir=d) |
| tr = pd.read_parquet(paths['train.parquet']).drop_duplicates('sequence') |
| va = pd.read_parquet(paths['val_inputs.parquet']).drop_duplicates('sequence') |
| return tr[['id', 'sequence']], va[['id', 'sequence']] |
|
|
|
|
| @torch.no_grad() |
| def embed(model, tok, seqs, device, token_budget): |
| order = np.argsort([len(s) for s in seqs])[::-1] |
| out = np.zeros((len(seqs), model.config.hidden_size), dtype=np.float16) |
| i, t0, done = 0, time.time(), 0 |
| while i < len(order): |
| L = min(len(seqs[order[i]]), MAX_LEN) + 2 |
| bs = max(1, token_budget // L) |
| idx = order[i:i + bs] |
| batch = tok([seqs[j][:MAX_LEN] for j in idx], return_tensors='pt', padding=True) |
| batch = {k: v.to(device) for k, v in batch.items()} |
| h = model(**batch).last_hidden_state |
| m = batch['attention_mask'].clone() |
| m[:, 0] = 0 |
| m[torch.arange(len(idx)), batch['attention_mask'].sum(1) - 1] = 0 |
| e = (h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=1) |
| out[idx] = e.float().cpu().numpy().astype(np.float16) |
| i += bs |
| done += len(idx) |
| if done // 2000 != (done - len(idx)) // 2000: |
| rate = done / (time.time() - t0) |
| print(f' {done}/{len(seqs)} {rate:.0f} seq/s eta {(len(seqs) - done) / rate / 60:.1f} min', flush=True) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--model', default='35M', choices=MODELS) |
| ap.add_argument('--tokens', type=int, default=0, help='tokens per batch (default by model size)') |
| ap.add_argument('--max-train', type=int, default=0, help='embed only a random sample of N training proteins') |
| a = ap.parse_args() |
| device = 'mps' if torch.backends.mps.is_available() else 'cpu' |
| budget = a.tokens or {'8M': 20000, '35M': 8000, '150M': 4000, '650M': 2000}[a.model] |
| print(f'device={device} model={MODELS[a.model]} tokens/batch={budget}', flush=True) |
| tr, va = load_data() |
| sfx = '' |
| if a.max_train and a.max_train < len(tr): |
| tr = tr.sample(a.max_train, random_state=0) |
| sfx = f'_s{a.max_train}' |
| (OUT / f'ids_train{sfx}.txt').write_text('\n'.join(tr.id)) |
| (OUT / 'ids_val.txt').write_text('\n'.join(va.id)) |
| tok = AutoTokenizer.from_pretrained(MODELS[a.model]) |
| model = EsmModel.from_pretrained(MODELS[a.model], add_pooling_layer=False).to(device).eval() |
| for name, df in (('val', va), ('train' + sfx, tr)): |
| path = OUT / f'emb_{a.model}_{name}.npy' |
| if path.exists(): |
| print(f'skip {name}: {path} already exists', flush=True) |
| continue |
| |
| |
| seqs = list(df.sequence) |
| parts = [] |
| nparts = (len(seqs) + PART - 1) // PART |
| for k in range(nparts): |
| pp = OUT / f'emb_{a.model}_{name}.part{k:03d}.npy' |
| parts.append(pp) |
| if pp.exists(): |
| print(f' part {k + 1}/{nparts} already done', flush=True) |
| continue |
| print(f'embedding {name} part {k + 1}/{nparts} ({len(seqs[k * PART:(k + 1) * PART])} sequences)', flush=True) |
| e = embed(model, tok, seqs[k * PART:(k + 1) * PART], device, budget) |
| np.save(pp, e) |
| if device == 'mps': |
| torch.mps.empty_cache() |
| e = np.concatenate([np.load(pp) for pp in parts]) |
| np.save(path, e) |
| for pp in parts: |
| pp.unlink() |
| print(f'saved {path} {e.shape}', flush=True) |
| print('DONE', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| sys.exit(main()) |
|
|
