Upload model card draft

README.md

@@ -1,3 +1,98 @@
----
-license: mit
----
+---
+tags:
+- knowledge graph
+- graph neural networks
+- graph ai
+- neuroscience
+- neurology
+- drug repurposing
+- alzheimer
+- parkinson
+- bipolar
+license: mit
+---
+
+# PROTON Model Card
+
+[![Website](https://img.shields.io/badge/Website-000000?style=for-the-badge)](https://zitniklab.hms.harvard.edu/PROTON)
+[![Paper](https://img.shields.io/badge/ArXiv-B31B1B?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org)
+[![Code](https://img.shields.io/badge/Code-181717?style=for-the-badge&logo=github&logoColor=white)](https://github.com/mims-harvard/PROTON)
+[![Model](https://img.shields.io/badge/Model-FFCC00?style=for-the-badge&logo=huggingface&logoColor=black)](https://huggingface.co/ayushnoori/PROTON)
+
+
+## Introduction
+
+Neurological diseases are the leading global cause of disability, yet most lack disease-modifying treatments. To help address this gap, we developed PROTON, a graph AI model that generates hypotheses for neurological disease. PROTON uses a heterogeneous graph transformer contextualized to the adult human brain. PROTON generated predictions across Parkinson's disease (PD), bipolar disorder (BD), and Alzheimer's disease (AD), which we validated using three independent biological systems. In PD, PROTON linked genetic risk loci to genes essential for dopaminergic neuron survival and identified pesticides toxic to patient-derived neurons, including the insecticide Naled ranked within the top 6.75% of predictions. *In silico* PROTON screens reproduced six genome-wide alpha-synuclein experiments, including a split-ubiquitin yeast two-hybrid system (normalized enrichment score NES = 2.27, FDR-adjusted *p* < 1E-4), an ascorbate peroxidase proximity labeling assay (NES = 2.22, FDR < 1E-4), and a high-depth targeted deep sequencing study in 496 synucleinopathy patients (NES = 1.73, FDR < 1.9E-3). In BD, PROTON nominated calcitriol as a candidate drug that reversed proteomic alterations observed in cortical organoids derived from BD patients. In AD, we evaluated PROTON predictions in electronic health records from *n* = 610,524 patients at Mass General Brigham, confirming that five PROTON-predicted drugs were associated with reduced seven-year dementia risk (minimum hazard ratio = 0.63, 95% CI: 0.53–0.75, *p* < 1E-7). PROTON generated and validated mechanistic hypotheses across molecular, organoid, and clinical systems, defining a path for AI-driven discovery in neurological disease.
+
+## Training Data
+
+PROTON was trained on NeuroKG, a heterogeneous, undirected biomedical knowledge graph contextualized to the human brain. NeuroKG unifies 36 human datasets and ontologies, and integrates single-nucleus RNA-sequencing atlases comprising 3,756,702 cells from the adult human brain. The knowledge graph contains 147,020 nodes across 16 entity types and 7,366,745 edges across 47 relation types. NeuroKG is available via Harvard Dataverse at DOI: [10.7910/DVN/ZDLS3K](https://doi.org/10.7910/DVN/ZDLS3K). For more details, please refer to our [project website](https://zitniklab.hms.harvard.edu/PROTON).
+
+## Model Architecture
+
+PROTON is a a 578-million-parameter heterogeneous graph transformer for neurological disease. It was trained on NeuroKG using a self-supervised link prediction objective. Through Bayesian hyperparameter optimization, we selected a model architecture that achieved high link-prediction performance (AUROC $=0.9145$; accuracy $=82.23\%$) on an independent test set.
+
+### Model Hyperparameters
+- `num_feat`: `1024`
+- `num_heads`: `4`
+- `hidden_dim`: `256`
+- `output_dim`: `128`
+- `num_layers`: `3`
+- `dropout_prob`: `0.4546844003628963`
+- `pred_threshold`: `0.5`
+
+### Files Included
+- `model.ckpt`: PyTorch Lightning checkpoint containing model weights.
+- `decoder.pt`: Decoder weights for link prediction (shape `[94, 512]`).
+- `edge_types.pt`: Ordered list of 47 edge types in NeuroKG to create edge type IDs.
+- `embeddings.pt`: Store of learned embeddings for all 147,020 nodes in NeuroKG (shape `[147020, 512]`).
+- `embeddings.csv`: Embedding store as a CSV file.
+
+For more details, please refer to our [project website](https://zitniklab.hms.harvard.edu/PROTON).
+
+
+## Usage Instructions
+
+To use PROTON, please clone the [GitHub repository](https://github.com/mims-harvard/PROTON) and follow the instructions in the [README.md](https://github.com/mims-harvard/PROTON/blob/main/README.md). For example, after downloading the model weights and modifying the `conf/default.config.yaml` file appropriately, you can load the model with the following code:
+```python
+import pytorch_lightning as pl
+from src.config import conf
+from src.constants import TORCH_DEVICE
+from src.dataloaders import load_graph
+from src.models import HGT
+
+pl.seed_everything(conf.seed, workers=True)
+kg = load_graph(nodes, edges)
+pretrain_model = HGT.load_from_checkpoint(
+    checkpoint_path=str(conf.hgt.checkpoint_path),
+    kg=kg,
+    strict=False,
+)
+pretrain_model.eval()
+pretrain_model.cache_graph(kg, overwrite=False, degree_threshold=conf.neurokg.hparams.degree_threshold)
+pretrain_model = pretrain_model.to(TORCH_DEVICE)
+```
+
+
+## License
+
+PROTON is released under the [MIT License](https://github.com/mims-harvard/PROTON/blob/main/LICENSE).
+
+
+## Citation
+
+If you use PROTON, please cite:
+```
+@article{noori_graph_2025,
+  title={Graph AI generates neurological hypotheses validated in molecular, organoid, and clinical systems},
+  author={Noori, Ayush and Polonuer, Joaquin and Meyer, Katharina and Budnik, Bogdan and Morton, Shad and Wang, Xinyuan and Nazeem, Sumaiya and He, Yingnan and Arango, Iñaki and Vittor, Lucas and Woodworth, Matthew and Krolewski, Richard C. and Li, Michelle M. and Liu, Ninning and Kamath, Tushar and Macosko, Evan and Ritter, Dylan and Afroz, Jalwa and Henderson, Alexander B. H. and Studer, Lorenz and Rodriques, Samuel G. and White, Andrew and Dagan, Noa and Clifton, David A. and Church, George M. and Das, Sudeshna and Tam, Jenny M. and Khurana, Vikram and Zitnik, Marinka},
+  journal={arXiv preprint},
+  note={arXiv:XXXX.XXXXX (placeholder)},
+  year={2025}
+}
+```
+
+
+## Contact
+
+For any questions or feedback, please open an issue in the [GitHub repository](https://github.com/mims-harvard/PROTON/issues/new) or contact [Ayush Noori](mailto:ayush.noori@sjc.ox.ac.uk) and [Marinka Zitnik](mailto:marinka@hms.harvard.edu).