Improve model card and add paper/code links

README.md

@@ -1,24 +1,24 @@
 ---
 library_name: tedbench
-tags:
-  - protein
-  - structure
-  - fold-classification
-  - tedbench
-pipeline_tag: feature-extraction
 license: bsd-3-clause
+pipeline_tag: graph-ml
+tags:
+- protein
+- structure
+- fold-classification
+- tedbench
 ---
 
 # TEDBench — Pretrained autoencoder (structure only)
 
 **Variant:** `miae_s`  |  **Parameters:** 29M  |  **Layers:** 6  |  **Hidden dim:** 512  |  **Attn heads:** 8
 
-This is a **pretrained MiAE** checkpoint. Use it as a feature extractor or as the starting point for fine-tuning.
+This repository contains a **pretrained MiAE** checkpoint. Use it as a feature extractor or as the starting point for fine-tuning on protein structure tasks.
+
+MiAE (Masked Invariant Autoencoder) is an $\mathrm{SE(3)}$-invariant masked autoencoder that masks up to 90% of backbone frames and reconstructs the full structure with a lightweight decoder. It was introduced in the paper [Protein Fold Classification at Scale: Benchmarking and Pretraining](https://arxiv.org/abs/2605.18552).
 
-Part of the [TEDBench](https://github.com/BorgwardtLab/TEDBench) benchmark for
-protein fold classification (ICML 2026). MiAE is an SE(3)-invariant masked
-autoencoder that masks up to 90% of backbone frames and reconstructs the full
-structure with a lightweight decoder.
+- **Paper:** [Protein Fold Classification at Scale: Benchmarking and Pretraining](https://arxiv.org/abs/2605.18552)
+- **Repository:** [BorgwardtLab/TEDBench](https://github.com/BorgwardtLab/TEDBench)
 
 ## Architecture sizes
 
@@ -34,6 +34,8 @@ Append `+model.use_seq_input=true` to `miae_b` for the **+seq** variant.
 
 ### Load from the HuggingFace Hub
 
+You can load the model using the `tedbench` library:
+
 ```python
 from tedbench.utils.io import load_from_hf
 
@@ -41,6 +43,15 @@ model = load_from_hf("TEDBench/miae-s")
 model.eval()
 ```
 
+Alternatively, using the library's high-level API:
+
+```python
+import tedbench
+
+model = tedbench.load_model("miae-s")
+model.eval()
+```
+
 ### From a Lightning checkpoint
 
 ```python
@@ -56,7 +67,7 @@ model.eval()
 @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},
+  booktitle={Proceedings of the 43rd International Conference on Machine Learning (ICML)},
   year={2026}
 }
-```
+```