File size: 1,497 Bytes
b142f2b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---
library_name: tedbench
tags:
  - protein
  - sequence
  - fold-classification
  - tedbench
  - esm2
pipeline_tag: other
license: bsd-3-clause
---

# TEDBench — ESM2-650M fine-tuned on TEDBench

Backbone: [ESM2-650M](https://huggingface.co/facebook/esm2_t33_650M_UR50D) (33 layers, hidden dim 1280).

Fine-tuned on [TEDBench](https://github.com/BorgwardtLab/TEDBench) for protein
fold classification into 965 CATH topology (T-level) classes (ICML 2026).

## Usage

```python
import sys
sys.path.insert(0, "baselines")  # from repo root

from pathlib import Path
import torch
from models.esm2_classifier import ESM2Classifier
from omegaconf import OmegaConf
from huggingface_hub import snapshot_download

local_dir = Path(snapshot_download("TEDBench/esm2-650M-ft"))
with open(local_dir / "config.json") as f:
    import json
    cfg = OmegaConf.create(json.load(f))

model = ESM2Classifier(cfg)
sd = torch.load(local_dir / "pytorch_model.bin", map_location="cpu", weights_only=False)
model.load_state_dict(sd)
model.eval()
```

Or pass the repo ID directly to the test script:

```bash
python baselines/esm2_test_ted.py train.ckpt_path=TEDBench/esm2-650M-ft
```

## Citation

```bibtex
@inproceedings{chen2026tedbench,
  title={Protein Fold Classification at Scale: Benchmarking and Pretraining},
  author={Chen, Dexiong and Manolache, Andrei and Niepert, Mathias and Borgwardt, Karsten},
  booktitle={Proceedings of the 43rd International Conference on Machine Learning},
  year={2026}
}
```