Upload README.md with huggingface_hub

README.md

@@ -18,15 +18,15 @@ pipeline_tag: other
 
 Pre-trained checkpoints for **QHFlow2** on the **QH9** dataset (DFT Hamiltonian prediction).
 
-**Paper:** [High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction](https://arxiv.org/abs/2602.16897)
-**Authors:** Seongsu Kim, Nayoung Kim, Dongwoo Kim, Sungsoo Ahn (KAIST)
-**Venue:** NeurIPS 2025
-**Code:** [github.com/seongsukim-ml/QHFlow2](https://github.com/seongsukim-ml/QHFlow2)
+> **Paper:** [High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction](https://arxiv.org/abs/2602.16897)
+> **Authors:** Seongsu Kim, Nayoung Kim, Dongwoo Kim, Sungsoo Ahn (KAIST)
+> **Venue:** NeurIPS 2025
+> **Code:** [github.com/seongsukim-ml/QHFlow2](https://github.com/seongsukim-ml/QHFlow2)
 
 ## Model Variants
 
 | Size | hidden_size | num_gnn_layers | Params | Checkpoint Size |
-|------|-------------|---------------|--------|----------------|
+|------|-------------|----------------|--------|-----------------|
 | small | 64 | 3 | ~35M | 139 MB |
 | middle | 128 | 3 | ~125M | 494 MB |
 | large | 256 | 4 | ~530M | 2.0 GB |
@@ -34,26 +34,121 @@ Pre-trained checkpoints for **QHFlow2** on the **QH9** dataset (DFT Hamiltonian
 
 ## Dataset Splits
 
-- **QH9Stable-random** — Random train/test split
-- **QH9Stable-size_ood** — Size out-of-distribution split
-- **QH9Dynamic-300k-geometry** — Geometry OOD split (300k conformers)
-- **QH9Dynamic-300k-mol** — Molecular OOD split (300k conformers)
+| Split | Description | Epochs |
+|-------|-------------|--------|
+| QH9Stable-random | Random train/test split | 78 |
+| QH9Stable-size_ood | Size out-of-distribution | 78 |
+| QH9Dynamic-300k-geometry | Geometry OOD (300k conformers) | 33 |
+| QH9Dynamic-300k-mol | Molecular OOD (300k conformers) | 33 |
 
 ## File Structure
 
-
-
-## Usage
-
-
-
-Or load manually:
-
-
+```
+{split}/
+  QHFlow_so2_v5_1_{size}-{split}/
+    weights-epoch=XX-val_loss=0.0000000.ckpt
+```
+
+## Quick Start
+
+### 1. Install QHFlow2
+
+```bash
+git clone https://github.com/seongsukim-ml/QHFlow2.git
+cd QHFlow2
+pip install -e ".[fairchem]"
+```
+
+### 2. Download Checkpoints
+
+```bash
+pip install huggingface_hub
+
+# Download a single checkpoint
+huggingface-cli download ksusu/QHFlow2-QH9 \
+  "QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt" \
+  --local-dir ckpt/QH9
+
+# Download all checkpoints for a specific split
+huggingface-cli download ksusu/QHFlow2-QH9 \
+  --include "QH9Stable-random/*" \
+  --local-dir ckpt/QH9
+
+# Download everything (~48 GB)
+huggingface-cli download ksusu/QHFlow2-QH9 --local-dir ckpt/QH9
+```
+
+Or in Python:
+
+```python
+from huggingface_hub import hf_hub_download
+
+path = hf_hub_download(
+    repo_id="ksusu/QHFlow2-QH9",
+    filename="QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt",
+)
+```
+
+### 3. Run Prediction
+
+```bash
+cd QHFlow2
+
+# Predict on QH9Stable test set
+python -m qhflow2.experiment.train_qh9 \
+  mode=predict \
+  dataset=QH9Stable \
+  dataset.split=random \
+  model=QHFlow_so2_v5_1_small \
+  ckpt=ckpt/QH9/QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt
+```
+
+### 4. Python API
+
+```python
+import torch
+from qhflow2.models import get_model, get_default_model_args
+
+# Build model
+args = get_default_model_args("qh9")
+args["version"] = "QHFlow_so2_v5_1"
+args["hidden_size"] = 64         # small
+args["num_gnn_layers"] = 3
+model = get_model(args)
+
+# Load checkpoint
+ckpt = torch.load("weights-epoch=78-val_loss=0.0000000.ckpt", map_location="cpu")
+state_dict = {k.replace("model.", ""): v for k, v in ckpt["state_dict"].items()}
+model.load_state_dict(state_dict, strict=False)
+model.eval()
+```
+
+### 5. Load via Lightning (full pipeline)
+
+```python
+from qhflow2.pl_module import _get_model_by_pl_type
+from omegaconf import OmegaConf
+
+conf = OmegaConf.load("configs/qh9/config_flow_v2_simple.yaml")
+LitModel = _get_model_by_pl_type("flow")
+lit_model = LitModel.load_from_checkpoint(
+    "weights-epoch=78-val_loss=0.0000000.ckpt",
+    conf=conf,
+    strict=False,
+)
+lit_model.eval()
+```
 
 ## Citation
 
-
+```bibtex
+@inproceedings{kim2025high,
+  title={High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction},
+  author={Kim, Seongsu and Kim, Nayoung and Kim, Dongwoo and Ahn, Sungsoo},
+  booktitle={Advances in Neural Information Processing Systems},
+  year={2025}
+}
+```
 
 ## License