TEDlm - TED-domain protein language models

TEDlm is a family of protein language models trained on TED (The Encyclopedia of Domains) structural domains from the AlphaFold Database. This repository provides a small, self-contained PyTorch implementation for running inference locally from the released checkpoints - no external services, no model hub, no registry.

Two families share one transformer backbone:

Model	Params	Width	Layers	Heads	LM head	Contact head
TEDlm-160M	168M	1024	16	16	20	–
TEDlm-650M	661M	1536	28	24	22	–
TEDlm-1.3B	1.3B	2048	32	32	22	–
TEDlm3D-160M	169M	1024	16	16	20	✓
TEDlm3D-650M	663M	1536	28	24	20	✓
TEDlm3D-1.3B	1.3B	2048	32	32	20	✓

TEDlm3D models add an auxiliary inter-residue contact head. The architecture and amino-acid ordering of any checkpoint are inferred from its tensors, so the loader needs no per-model configuration. (The 160M sequence model predates the 650M/1.3B ones and shares the 20-channel head and ESM ordering of the TEDlm3D family.)

Installation

pip install -r requirements.txt          # torch + numpy
# or:  conda env create -f environment.yml && conda activate tedlm

FlashAttention is optional. With it (CUDA, fp16/bf16) attention runs on the flash kernel; without it the models run anywhere on PyTorch's scaled_dot_product_attention.

Model weights

Checkpoints are distributed separately (see the release page) and are not in this repository. Place them under checkpoints/:

checkpoints/tedlm/tedlm_160M.pt        checkpoints/tedlm3D/tedlm3D_160M.pt
checkpoints/tedlm/tedlm_650M.pt        checkpoints/tedlm3D/tedlm3D_650M.pt
checkpoints/tedlm/tedlm_1.3B.pt        checkpoints/tedlm3D/tedlm3D_1.3B.pt

Quick start

from util import load_model, model_embed, masked_recovery, predict_contacts

# The amino-acid ordering and architecture are detected from the checkpoint.
model, tok = load_model("checkpoints/tedlm/tedlm_650M.pt", device="cuda")

seq = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVK"

emb = model_embed(model, tok, seq, layers=(1, 4))   # {layer_from_end: vector}
acc = masked_recovery(model, tok, seq)              # top-1 recovery of masked residues

# Contact map (TEDlm3D only):
m3d, tok3d = load_model("checkpoints/tedlm3D/tedlm3D_650M.pt", device="cuda")
cmap = predict_contacts(m3d, tok3d, seq, apply_apc=True)   # (L, L) probabilities

model_embed counts layers from the end (1 = last); intermediate layers usually carry the strongest homology signal.

Command-line tools

# per-sequence embeddings (.npz)
python scripts/run_embeddings.py --checkpoint checkpoints/tedlm/tedlm_650M.pt \
    --fasta testdata/homology/cath_s40.fasta --layers 1,4 --limit 100 --out embeddings.npz

# masked-residue recovery (.csv)
python scripts/run_recovery.py --checkpoint checkpoints/tedlm3D/tedlm3D_160M.pt \
    --fasta testdata/homology/cath_s40.fasta --limit 100 --out recovery.csv

# contact prediction + scoring (TEDlm3D)
bash scripts/download_cath.sh --pdb
python scripts/make_contact_gt.py --pdb-dir testdata/dompdb --out testdata/contact/ground_truth
python scripts/run_contacts.py --checkpoint checkpoints/tedlm3D/tedlm3D_650M.pt \
    --gt-dir testdata/contact/ground_truth --apc --out contacts_eval.csv

Repository layout

model/          pure-PyTorch model definition
  attention.py    RoPE multi-head attention (+ optional FlashAttention)
  nndef_tedlm.py  unified SeqTransformer (LM head + optional contact head)
util.py         tokeniser, checkpoint loading (architecture inference), inference ops
scripts/        command-line tools + CATH data download / contact ground truth
testdata/       CATH v4.4.0 S40 sequences + T-level representative IDs

Amino-acid ordering

Token ids 0–19 are the amino acids; 20 <unk>, 21 <eod>, 22 <mask>, 23 <pad>. The ordering follows the LM head size, and load_model selects it automatically (override with aa_order=):

• TEDlm-650M / 1.3B — 22-channel head, ACDEFGHIKLMNPQRSTVWY
• TEDlm3D (all sizes) and TEDlm-160M — 20-channel head, ESM ordering ARNDCQEGHILKMFPSTWYV

License

Apache-2.0 (see LICENSE).