Update the model card for the revised paper

README.md

@@ -12,13 +12,44 @@ tags:
 
 # TRIAGE-4B-P12-SFT
 
-SFT checkpoints released as part of the \\( \textbf{TRIAGE}_{\text{SFT}} \\) work, applied to the P12 dataset. 
+SFT checkpoints of **TRIAGE** for in-hospital mortality (P12), one per cross-validation split. Part of the [TRIAGE collection](https://huggingface.co/collections/Hyeongwon/triage).
 
-The model was presented in the paper [TRIAGE: Dialectical Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series with LLMs](https://huggingface.co/papers/2606.09030).
+- Paper: [TRIAGE: Dialectical LLM Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series](https://huggingface.co/papers/2606.09030)
+- Code, data and training pipeline: [github.com/HyeongWon-Jang/TRIAGE](https://github.com/HyeongWon-Jang/TRIAGE)
+- Project page: [hyeongwon-jang.github.io/triage-project-page](https://hyeongwon-jang.github.io/triage-project-page/)
 
-Each split was fine-tuned from [Qwen/Qwen3-4B-Base](https://huggingface.co/Qwen/Qwen3-4B-Base) and placed in its own `split_N/` subfolder; per-split checkpoints were selected by **Validation AUPRC** over a 3-epoch SFT run.
+TRIAGE trains an LLM to reason dialectically over the competing clinical outcomes. Given a patient's irregularly sampled medical time series, the model writes a rationale for survival and a rationale for in-hospital death before it commits to a final decision, and the risk score is read from the model's probability over that decision. This mitigates *risk polarization*, the collapse of graded risk into overconfident predictions under conventional LLM reasoning, so that one model gives both a clinician-readable rationale and a risk score comparable across patients.
 
-**Code, data, & training pipeline**: <https://github.com/HyeongWon-Jang/TRIAGE>
+The SFT + RL checkpoints post-trained from this stage are released as [TRIAGE-4B-P12](https://huggingface.co/Hyeongwon/TRIAGE-4B-P12).
+
+## Model description
+
+| | |
+|---|---|
+| Base model | [Qwen/Qwen3-4B-Base](https://huggingface.co/Qwen/Qwen3-4B-Base) |
+| Task | In-hospital mortality prediction on P12 (PhysioNet Challenge 2012) |
+| Training | supervised fine-tuning on dialectical-reasoning data with [TRL](https://github.com/huggingface/trl), 3 epochs ([`triage_sft/`](https://github.com/HyeongWon-Jang/TRIAGE/tree/main/triage_sft)) |
+| Layout | `split_1` … `split_5`, one checkpoint per split, each in its own subfolder |
+| Selection | per split, the checkpoint with the best validation AUPRC over the SFT run |
+
+## Input and output format
+
+The model expects the chat-style prompt the training data uses: the patient's static features and per-variable `(time, value)` series, the question, and the reasoning order (which outcome to argue for first). The exact prompt and how it is built from the raw data are in [`triage_sft/data/P12/`](https://github.com/HyeongWon-Jang/TRIAGE/tree/main/triage_sft/data/P12).
+
+The answer has three headed sections, in the order the prompt asks for:
+
+```
+## Rationale for survival
+...
+
+## Rationale for in-hospital death
+...
+
+## Final Decision
+1
+```
+
+`0` means survival and `1` means in-hospital death. The risk score is the next-token probability of `"1"` at the decision position, renormalised over the two tokens `"0"` and `"1"`, averaged over the two rationale orderings.
 
 ## Quick start
 
@@ -29,17 +60,33 @@ split = "split_1"  # one of split_1 ... split_5
 repo  = "Hyeongwon/TRIAGE-4B-P12-SFT"
 
 tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=split)
-model     = AutoModelForCausalLM.from_pretrained(repo, subfolder=split, device_map="auto")
+model     = AutoModelForCausalLM.from_pretrained(repo, subfolder=split, torch_dtype="auto", device_map="auto")
 ```
 
-The model expects a task-specific input/output template; for the full inference pipeline, see the linked GitHub repository.
+Scoring is two-step: sample the dialectical reasoning at temperature 0.7 up to the `## Final Decision` header, then read the logits of `"0"` and `"1"` at the next position. A complete `predict_risk` function is in the [GitHub README](https://github.com/HyeongWon-Jang/TRIAGE#quick-start).
 
 ## Data
 
-- **Raw**: [PhysioNet Mortality Prediction Challenge 2012](https://physionet.org/content/challenge-2012/1.0.0/)
-- **Processed** (Raindrop, CC BY 4.0): [figshare DOI](https://doi.org/10.6084/m9.figshare.19514341.v1)
+- Raw: [PhysioNet Challenge 2012](https://physionet.org/content/challenge-2012/1.0.0/), in-hospital mortality (Open Data Commons Attribution License v1.0).
+- Processed: [Raindrop-preprocessed release on figshare](https://doi.org/10.6084/m9.figshare.19514341.v1).
+- Splits: the five cross-validation index files are in the GitHub repository ([`triage_sft/data/P12/splits/`](https://github.com/HyeongWon-Jang/TRIAGE/tree/main/triage_sft/data/P12/splits), `phy12_split{1..5}.npy`); each split's checkpoint here was trained on the training portion of the file with the same number.
+
+No data is included in this repository.
+
+## License
+
+Model checkpoints are released under **CC BY-NC 4.0** (non-commercial). The base model and the datasets remain under their own licences.
 
-Further preprocessing and split-construction details are in the linked GitHub repository.
+## Citation
+
+```bibtex
+@article{jang2026triage,
+  title={TRIAGE: Dialectical LLM Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series},
+  author={Jang, Hyeongwon and Chu, Gyouk and Kim, Changhun and Yoon, Hangyul and Lee, Jeonguk and Yang, Eunho and Park, Joonhyung},
+  journal={arXiv preprint arXiv:2606.09030},
+  year={2026}
+}
+```
 
 ## Framework versions
 
@@ -47,7 +94,3 @@ Further preprocessing and split-construction details are in the linked GitHub re
 - PyTorch: 2.6.0
 - Datasets: 3.6.0
 - Tokenizers: 0.22.2
-
-## License
-
-Model checkpoints are released under **CC BY-NC 4.0** (non-commercial). Datasets remain under their respective licenses.