Update README.md

README.md

@@ -15,132 +15,91 @@ library_name: transformers
 
 # NeuroCLR
 
-**NeuroCLR** is a self-supervised learning (SSL) framework for learning **robust, disorder-agnostic neural representations** from raw, unlabeled resting-state fMRI (rs-fMRI) regional time series. NeuroCLR is designed for **multi-site generalization** and **transfer** to downstream disorder classification with limited labeled data.
+**NeuroCLR** is a self-supervised learning (SSL) framework for learning **robust, disorder-agnostic neural representations** from raw, unlabeled resting-state fMRI (rs-fMRI) regional time-series data. NeuroCLR is designed for **multi-site generalization** and **transfer** to downstream disorder classification tasks with limited labeled data.
 
-\[[GitHub Repo](https://github.com/pcdslab/NeuroCLR)\] | \[[Cite](#citation)\]
+[[GitHub Repo](https://github.com/pcdslab/NeuroCLR)] | [[Cite](#citation)]
 
 ---
 
 ## Abstract
 
-Self-supervised learning (SSL) is a powerful technique in computer vision for drastically reducing the dependency on large amounts of labeled training data. The availability of large-scale, unannotated, rs-fMRI data provides opportunities for the development of superior machine-learning models for classification of disorders across heterogeneous sites, and diverse subjects. In this paper, we propose NeuroCLR, a novel self-supervised learning (SSL) framework. NeuroCLR extracts robust and rich invariant neural representations - consistent across diverse experimental subjects and disorders - using contrastive principles, spatially constrained learning, and augmented views of unlabeled raw fMRI time series data. We pre-trained NeuroCLR using a combination of heterogeneous disorders from more than 3,600 participants across 44 different sites, and 720,000 region-specific time series fMRI data. The resultant disorder-agnostic pre-trained model is fine-tuned for downstream disorder-specific classification tasks on limited labelled data. We evaluate NeuroCLR on diverse disorder classification tasks and find that it outperforms both deep-learning, and SSL models that have been trained on a single disorder. Experiments also confirmed robust generalizability, consistently outperforming baselines across neuroimaging sites. This study is the first to present robust and reproducible self-supervised methodology with anatomically consistent contrastive objective that operates on raw unlabelled fMRI data, capable of reliable transfer across diagnostic categories. This will cultivate stronger participation by computational and clinical researchers, setting the stage for the development of sophisticated diagnostic models, for various neurodegenerative and neurodevelopmental disorders, leveraging NeuroCLR.
+Self-supervised learning (SSL) is a powerful technique for reducing dependence on large labeled datasets. The availability of large-scale, unannotated rs-fMRI data provides opportunities to develop robust machine-learning models for classification across heterogeneous sites and diverse cohorts.
+
+In this work, we propose **NeuroCLR**, a novel SSL framework that learns invariant neural representations using contrastive objectives, spatial constraints, and augmented views of raw fMRI time-series data. NeuroCLR is pre-trained on data from more than 3,600 participants across 44 sites, comprising over 720,000 region-specific fMRI time series.
+
+The resulting disorder-agnostic foundation model is fine-tuned for downstream classification tasks with limited labeled data and consistently outperforms both supervised deep-learning and SSL models trained on single disorders. NeuroCLR demonstrates strong cross-site generalizability and reliable transfer across diagnostic categories, enabling reproducible and scalable neuroimaging representation learning.
 
 ---
 
-## Model Contents (Subfolders)
+## Model Structure
 
-This repository provides two loadable model artifacts:
+This repository provides **two loadable model artifacts**:
 
-- **`pretraining/`**: SSL encoder + projector (contrastive pre-training)
-- **`classification/`**: Encoder + ResNet1D classification head (downstream classification)
+- **Root model (default)**  
+  Self-supervised **pretraining encoder + projector** (contrastive SSL)
 
-Both subfolders use **custom code** (Transformer encoder + projector; and a ResNet1D head), so you must set `trust_remote_code=True` when loading.
+- **`classification/` subfolder**  
+  Encoder + **ResNet1D classification head** for downstream tasks
+
+All models rely on **custom architectures**, so `trust_remote_code=True` is required.
 
 ---
 
 ## Model Details
 
-### 1) Pretraining model (`pretraining/`)
+### 1) Pretraining Model (Default, Loaded from Repo Root)
+
 - **Input**: region-wise rs-fMRI time series  
-  Shape: **`[B, 1, 128]`** (batch, sequence length, embedding dimension)
+  Shape: **`[B, 1, L]`**, where `L = 128` time points
 - **Output**:
   - `h`: pooled representation, shape **`[B, 128]`**
-  - `z`: projected representation, shape **`[B, projector_out2]`**
-
-### 2) Classification model (`classification/`)
-- **Input**: ROI-by-time representation (e.g., 200 ROIs)  
-  Shape: **`[B, 200, 128]`**
-- **Output**:
-  - `logits`: shape **`[B, num_labels]`**
-  - `loss`: returned when labels are provided
+  - `z`: projected representation, shape **`[B, projector_out_dim]`**
 
-**Note on freezing**: By default, the encoder is packaged with the classification model. Depending on the configuration used at training/export time, the encoder may be frozen for downstream training (recommended). See the GitHub repo for training scripts and details.
+This model is intended for:
+- representation learning
+- feature extraction
+- transfer learning
 
 ---
 
-## Usage (PyTorch)
+### 2) Classification Model (`classification/`)
 
-```python
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "SaeedLab/NeuroCLR",
-    subfolder="pretraining",
-    trust_remote_code=True
-)
-
-model.eval()
-
-x = torch.randn(4, 1, 128)  # [batch, seq_len, feature_dim]
-
-with torch.no_grad():
-    outputs = model(x)
-
-print(outputs["h"].shape)
-print(outputs["z"].shape)
-```
+- **Input**: ROI-by-time representation  
+  Shape: **`[B, N_ROIs, 128]`** (e.g., `N_ROIs = 200`)
+- **Output**:
+  - `logits`: shape **`[B, num_labels]`**
+  - `loss`: returned when labels are provided
 
-## Model Usage
+> **Note**  
+> The encoder is bundled with the classification model and may be frozen by default (recommended).  
+> See the GitHub repository for training and fine-tuning scripts.
 
-### Requirements
+---
 
-```bash
-pip install torch transformers huggingface_hub safetensors
-```
+## Quickstart (PyTorch)
 
 ### Load the Pretraining Model (SSL Encoder)
 
-```py
+```python
 import torch
 from transformers import AutoModel
 
 model = AutoModel.from_pretrained(
     "SaeedLab/NeuroCLR",
-    subfolder="pretraining",
     trust_remote_code=True
 )
 
 model.eval()
 
-x = torch.randn(4, 1, 128)  # [batch, sequence_length, embedding_dim]
+x = torch.randn(4, 1, 128)  # [batch, 1, time_points]
 
 with torch.no_grad():
     outputs = model(x)
 
 print(outputs["h"].shape)  # [4, 128]
-print(outputs["z"].shape)  # [4, projector_out2]
-```
-
-This repository is public, and no Hugging Face authentication is required to download or use the model.
-
-Users may see a warning when accessing the model without authentication. This warning is harmless and can be safely ignored.
-
-However, users may optionally authenticate using their own Hugging Face account by passing their access_token as follows:
-
-```py
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "SaeedLab/NeuroCLR",
-    subfolder="pretraining",
-    trust_remote_code=True,
-    access_token=your_huggingface_token
-)
-
-model.eval()
-
-x = torch.randn(4, 1, 128)  # [batch, sequence_length, embedding_dim]
-
-with torch.no_grad():
-    outputs = model(x)
-
-print(outputs["h"].shape)  # [4, 128]
-print(outputs["z"].shape)  # [4, projector_out2]
+print(outputs["z"].shape)
 ```
 
-
 ### Load the Downstream Classification Model
 ```py
 import torch