Update README.md

README.md

@@ -10,8 +10,8 @@ license: apache-2.0
 
 # CSI-4CAST Organization
 
-Welcome to the CSI-4CAST organization on Hugging Face! This organization hosts datasets for CSI prediction research.
-This dataset is originally created for our research paper: [CSI-4CAST: A Hybrid Deep Learning Model for CSI Prediction with Comprehensive Robustness and Generalization Testing](https://arxiv.org/abs/2510.12996). The corresponding code and implementation are available in our [GitHub repo](https://github.com/AI4OPT/CSI-4CAST).
+Welcome to the CSI-4CAST organization on Hugging Face. This organization hosts both **datasets** and **model weights** for CSI prediction research.
+These resources are released as part of our paper [CSI-4CAST: A Hybrid Deep Learning Model for CSI Prediction with Comprehensive Robustness and Generalization Testing](https://arxiv.org/abs/2510.12996v2). The corresponding code and implementation are available in our [GitHub repo](https://github.com/AI4OPT/CSI-4CAST).
 
 ## TL;DR
 
@@ -19,6 +19,7 @@ This dataset is originally created for our research paper: [CSI-4CAST: A Hybrid
 
 - **For specific datasets**: Use the `snapshot_download` command to download individual datasets you need
 - **For all datasets with original structure**: Run [`download.py`](https://huggingface.co/spaces/CSI-4CAST/README/blob/main/download.py)  followed by [`reconstruction.py`](https://huggingface.co/spaces/CSI-4CAST/README/blob/main/reconstruction.py) to get the complete, well-structured dataset
+- **For trained checkpoints**: Download model weights from [`CSI-4CAST/weights`](https://huggingface.co/CSI-4CAST/weights)
 
 See the **Usage** section below for detailed instructions.
 
@@ -113,6 +114,69 @@ snapshot_download(repo_id="CSI-4CAST/stats", repo_type="dataset")
 snapshot_download(repo_id="CSI-4CAST/test_regular_cm_A_ds_030_ms_001", repo_type="dataset")
 ```
 
+### Downloading Model Weights
+
+All trained checkpoints are available in the model repository [`CSI-4CAST/weights`](https://huggingface.co/CSI-4CAST/weights).
+
+Download the full weights repository:
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(repo_id="CSI-4CAST/weights")
+```
+
+Download only one scenario or model:
+
+```python
+from huggingface_hub import snapshot_download
+
+# Download one scenario (e.g. FDD)
+snapshot_download(
+    repo_id="CSI-4CAST/weights",
+    allow_patterns=["fdd/*"],
+)
+
+# Download one checkpoint directory
+snapshot_download(
+    repo_id="CSI-4CAST/weights",
+    allow_patterns=["tdd/llm4cp/*"],
+)
+```
+
+The weights repository is organized as:
+
+```text
+fdd/
+  abl_no_arl/
+  abl_no_subcarrier_arl/
+  cnn/
+  llm4cp/
+  model/
+  rnn/
+  stemgnn/
+  wiener/
+tdd/
+  abl_add_subcarrier_arl/
+  abl_lstm_replace_pred/
+  abl_mlp_replace_embed/
+  abl_mlp_replace_pred/
+  abl_mobilenet_replace_embed/
+  abl_no_arl/
+  abl_no_denoiser/
+  abl_no_idft/
+  abl_norm_replace_arl/
+  ar/
+  cnn/
+  llm4cp/
+  model/
+  rnn/
+  stemgnn/
+  wiener/
+```
+
+Neural models store checkpoints as `model.ckpt`. Statistical baselines store parameters as `params.npz`.
+
 ### Downloading All Datasets
 
 To download all available datasets at once, use the provided [`download.py`](https://huggingface.co/spaces/CSI-4CAST/README/blob/main/download.py) script:
@@ -160,6 +224,10 @@ Each dataset folder contains:
 - `H_U_hist.pt`: Historical H_U values (PyTorch tensor)  
 - `H_U_pred.pt`: Predicted H_U values (PyTorch tensor)
 
+The weights repository contains:
+- `model.ckpt`: PyTorch Lightning checkpoint for trained neural models
+- `params.npz`: Saved parameter arrays for statistical baselines such as AR and Wiener filters
+
 ## Questions & Contributions
 
 For further questions or any contribution suggestions, you can create pull requests here or to the [GitHub homepage](https://github.com/AI4OPT/CSI-4CAST) of this organization.
@@ -167,13 +235,13 @@ For further questions or any contribution suggestions, you can create pull reque
 ## Citation
 
 ```bibtex
-@misc{cheng2025csi4casthybriddeeplearning,
+@misc{cheng2026csi4casthybriddeeplearning,
       title={CSI-4CAST: A Hybrid Deep Learning Model for CSI Prediction with Comprehensive Robustness and Generalization Testing}, 
       author={Sikai Cheng and Reza Zandehshahvar and Haoruo Zhao and Daniel A. Garcia-Ulloa and Alejandro Villena-Rodriguez and Carles Navarro Manchón and Pascal Van Hentenryck},
-      year={2025},
+      year={2026},
       eprint={2510.12996},
       archivePrefix={arXiv},
       primaryClass={cs.LG},
-      url={https://arxiv.org/abs/2510.12996}, 
+      url={https://arxiv.org/abs/2510.12996v2}, 
 }
-```
+```