Update README.md

README.md

@@ -2,242 +2,255 @@
 language:
 - en
 license: other
-library_name: transformers
 pipeline_tag: text-generation
 tags:
 - clinical-nlp
-- medical-reasoning
+- medical-coding
+- icd10
 - icd-10-cm
+- reasoning
 - reinforcement-learning
 - grpo
-- qwen2.5
-- diagnosis-prediction
-- chain-of-thought
-- research
+- healthcare
 base_model:
 - Qwen/Qwen2.5-32B-Instruct
-model-index:
-- name: DeepICD-R1-zero-32B
-  results: []
 ---
 
 # DeepICD-R1-zero-32B
 
-DeepICD-R1-zero-32B is a clinical reasoning model for **ICD-10-CM diagnosis outcome prediction from admission notes**, obtained by applying **Group Relative Policy Optimization (GRPO)** to **Qwen2.5-32B-Instruct**.
+## Model Summary
 
-This model corresponds to the **GRPO-only large model** described in the DeepICD-R1 paper, where it serves as the first-stage reasoning model before the creation of the distilled supervised fine-tuning dataset used for smaller downstream models. In the paper, this model is referred to as **DeepICD-R1-zero-32B**. It is part of the broader **DeepICD-R1** framework for hierarchical medical reasoning with verifiable rewards.  
+**DeepICD-R1-zero-32B** is a clinical reasoning model designed for **ICD-10-CM diagnosis outcome prediction from admission notes**.  
+It follows the **DeepICD-R1 framework**, which treats diagnosis prediction as a reasoning task optimized with reinforcement learning and structured reward signals.
 
-## Relation to the paper
+This checkpoint corresponds to a **“R1-Zero” style model**, meaning it was trained primarily through **reinforcement learning without a supervised fine-tuning (SFT) initialization**, allowing reasoning behaviors to emerge directly from reward optimization.
 
-This repository contains the model corresponding to the **“zero” GRPO-trained large model** in:
+The approach is inspired by reasoning-focused training pipelines where reinforcement learning alone can induce structured reasoning behaviors and self-verification in large language models.
 
-**DeepICD-R1: Medical Reasoning through Hierarchical Rewards and Unsupervised Distillation**
+---
 
-In the paper, the overall framework has two stages:
+# Model Details
 
-1. A large instruction-tuned model is optimized with **GRPO** using structured clinical rewards.  
-   This yields **DeepICD-R1-zero-32B**.
-2. That model is then used to generate reasoning traces which are distilled into a large supervised dataset for smaller models such as DeepICD-R1-7B.
+- **Model name:** DeepICD-R1-zero-32B  
+- **Organization:** DATEXIS  
+- **Model size:** ~32B parameters  
+- **Task:** Single ICD-10-CM diagnosis prediction from clinical text  
+- **Training paradigm:** Reinforcement learning (GRPO-style)  
+- **Framework:** VERL reinforcement learning trainer  
+- **Domain:** Clinical NLP / medical reasoning  
 
-The paper describes this role as follows: the large base LLM is trained with GRPO and dedicated reward functions to produce **DeepICD-R1-zero-32B**, which is then used in dataset construction and later fine-tuning stages. :contentReference[oaicite:1]{index=1}
+### Related Research
 
-## Model description
+This model follows the **DeepICD-R1** framework introduced in:
 
-- **Model name:** DeepICD-R1-zero-32B
-- **Base model:** `Qwen/Qwen2.5-32B-Instruct`
-- **Training method:** Reinforcement learning with **GRPO**
-- **Domain:** Clinical NLP
-- **Task:** Predicting the first annotated **ICD-10-CM diagnosis code** from admission notes
-- **Input:** Admission note text
-- **Output:** Structured reasoning plus a predicted ICD-10-CM code
+> *DeepICD-R1: Medical Reasoning through Hierarchical Rewards and Unsupervised Distillation*
 
-The model is trained to generate outputs in the following structure:
+The paper proposes a system for diagnosis prediction that combines:
 
-```xml
-<think>
-...
-</think>
-<diagnosis>
-...
-</diagnosis>
-```
-The paper uses this structured format together with **hierarchical ICD-aware rewards** and an **LLM-based reasoning reward**.
+- structured reasoning traces  
+- hierarchical reward signals aligned with ICD code structure  
+- reinforcement learning for reasoning optimization  
 
 ---
 
 # Intended Use
 
-This model is intended for:
+This model is intended for **research purposes**, including:
+
+- clinical reasoning experiments  
+- ICD-10-CM code prediction research  
+- reinforcement learning for language models  
+- reasoning trace generation  
+- structured prediction from clinical notes  
 
-- research on clinical reasoning with language models
-- ICD-10-CM outcome prediction from admission notes
-- studying reinforcement learning with verifiable hierarchical rewards
-- generating reasoning traces for analysis or data distillation
-- reproducing or extending the DeepICD-R1 framework
+### Out-of-Scope Use
+
+This model **must not** be used for:
+
+- medical diagnosis  
+- clinical decision making  
+- patient triage  
+- automated medical coding without expert supervision  
+- billing or compliance workflows  
 
 ---
 
-# Out-of-Scope Use
+# Training Methodology
+
+## R1-Zero Training Paradigm
 
-This model is **not intended for**:
+The model follows a **Zero-stage reasoning training approach**, where reinforcement learning is applied directly to a base language model without prior supervised instruction tuning.
 
-- real-world diagnosis
-- clinical decision support in production
-- autonomous medical coding in care settings
-- unsupervised deployment on patient data
-- use without human oversight
+This method encourages models to discover reasoning strategies autonomously during training, allowing behaviors such as:
 
-As emphasized in the paper, this is a **research prototype** and must not be used for real-world diagnosis or clinical decision-making. Generated reasoning may appear plausible while still being clinically incorrect.
+- chain-of-thought reasoning  
+- self-verification  
+- iterative reasoning refinement  
+
+to emerge naturally from the reward signal.
+
+However, purely RL-trained models may also exhibit issues such as:
+
+- repetitive reasoning patterns  
+- readability problems  
+- mixed language outputs  
 
 ---
 
 # Training Data
 
-The model was trained on **MIMIC-IV admission notes** for single-label prospective ICD-10-CM outcome prediction.
-According to the paper, the task is formulated as predicting the **first annotated diagnosis code from admission-time information**, using MIMIC-IV admission notes and excluding leakage-prone diagnostic and treatment sections.
-PhysioNet Link soon!
+The training task uses **clinical admission notes paired with ICD-10-CM diagnoses**, derived from de-identified electronic health record datasets such as **MIMIC-IV**.
 
----
+Task formulation:
+
+- **Input:** admission note describing a patient case  
+- **Output:** reasoning trace and predicted ICD-10-CM code  
 
-# Training Procedure
-
-This model was trained with the **verl PPO trainer** using **GRPO** as the advantage estimator.
-
-## Core Setup
-
-- **Trainer:** `verl.trainer.main_ppo`
-- **Advantage estimator:** `grpo`
-- **Base model:** `Qwen/Qwen2.5-32B-Instruct`
-- **Epochs:** 1
-- **Effective train batch size:** 64
-- **Rollouts per prompt:** 8
-- **Max prompt length:** 2048
-- **Max response length:** 1024
-- **Sampling temperature:** 0.9
-- **Learning rate:** 1e-6
-- **Warmup steps:** 80
-- **Entropy coefficient:** 0.001
-- **KL loss:** disabled
-- **Actor torch compile:** enabled
-- **Gradient checkpointing:** enabled
-- **Rollout engine:** vLLM
-- **Rollout dtype:** bfloat16
+The model learns to infer diagnostic outcomes based on the textual description of the patient presentation.
 
 ---
 
-# Hardware
+# Output Format
 
-- **GPUs:** 8
-- **Nodes:** 1
-- **GPU type:** not explicitly specified in the config
-- **Memory limit:** 512 GiB
+The model is trained to produce structured outputs separating reasoning from the final diagnosis.
 
----
+### Example
 
-# Reward Setup
+```text
+<think>
+The patient presents with ...
+Symptoms and history suggest ...
+...
+</think>
 
-Training used a **custom batched reward function** with the following active components:
+<diagnosis>
+M5116
+</diagnosis>
+```
+The reasoning trace allows the model to explain how the diagnosis is derived from the clinical note.
 
-- Outcome reward: enabled
-- Format reward: enabled
-- LLM-as-a-judge reward: enabled
-- Judge RAG: enabled
-- Guidelines file: ICD-10-CM chapter guidelines JSON
-- Judge model: `meta-llama/Llama-3.1-8B-Instruct`
+---
 
-## Selected Reward Environment Settings
+## Evaluation
 
-ACTIVATE_OUTCOME_REWARD=True
-ACTIVATE_FORMAT_REWARD=True
-JUDGE_RAG_ENABLED=True
-NO_MATCH_MALUS=-1
-THINK_TRACE_REWARD=1
-MATCH_REWARD=15
-LLM_REWARD_SCALING=0.8
+Evaluation follows the methodology described in the **DeepICD-R1 paper**.
 
+Performance is typically measured using **macro-averaged F1 scores** at multiple levels of the ICD hierarchy.
 
-This aligns with the paper’s training design, which combines:
+| Level | Description |
+|------|-------------|
+| Chapter | Broad ICD category |
+| Category | First three digits |
+| Full code | Complete ICD-10 code |
 
-- a **format reward** for `<think>` and `<diagnosis>` structure
-- a **hierarchical ICD outcome reward**
-- an **LLM-as-a-judge reward** to improve reasoning clarity and consistency
+Hierarchical evaluation allows partial credit when the model predicts the correct high-level diagnostic category even if the full code is incorrect.
 
 ---
 
-# Prompt / Output Format
+## Limitations
 
-The model expects an **admission note** and is trained to return:
+Models following the DeepICD-R1 framework share several limitations.
 
-1. a reasoning trace inside `<think>...</think>`
-2. a predicted ICD-10-CM code inside `<diagnosis>...</diagnosis>`
+### Dataset limitations
 
-## Example Schema
+- Training data consists primarily of **English clinical notes**
+- Distribution reflects **hospital-specific patient populations**
+- ICD labels are **highly imbalanced**, affecting rare diagnoses
 
-```xml
-<think>
-Reasoning over presenting symptoms, history, and admission note evidence.
-</think>
+### Model limitations
 
-<diagnosis>
-M5116
-</diagnosis>
-```
-Users should validate that generated outputs conform to this format before downstream evaluation.
+- Reasoning traces may appear convincing while being incorrect
+- Predictions may fail for rare or long-tail diagnoses
+- Models may demonstrate **premature diagnostic closure**
+- Reinforcement learning signals are only proxies for expert feedback
 
 ---
 
-# Evaluation
+## Ethical Considerations
 
-In the paper, evaluation is performed on **hierarchical ICD-10-CM prediction tasks** at three levels:
+This model is trained on **de-identified clinical data** and intended strictly for research.
 
-- **Chapter**
-- **Category**
-- **Full diagnosis code**
+Potential risks include:
 
-The paper reports that the **GRPO-only 32B model improves over the instruction-tuned baseline**, but remains weaker than models trained with **both supervised fine-tuning and GRPO**, especially for **fine-grained full-code prediction**.
+- propagation of dataset biases  
+- overconfidence in generated reasoning  
+- misuse in clinical decision making  
 
-This repository does **not claim any additional benchmark results** beyond those reported in the paper unless explicitly added later.
+Appropriate safeguards include:
 
----
+- expert oversight  
+- dataset bias evaluation  
+- fairness audits  
+- controlled deployment environments  
 
-# Limitations
+---
 
-Important limitations discussed in the paper include:
+## Hardware and Training Setup
 
-- reasoning traces may be coherent but **not clinically correct**
-- the model can exhibit **premature diagnostic closure**
-- performance drops on **fine-grained and rare ICD codes**
-- the underlying data reflects **institutional and demographic bias**
-- the model may fail to capture the **severity or clinical significance of diagnoses**
-- reinforcement signals based on **automatic rewards and LLM judging are only proxies for expert review**
+Typical training configuration for models in this family includes:
 
-The paper also notes that clinicians often preferred **concise reasoning**, and that plausible-looking outputs may still omit important differential diagnoses.
+- **GPUs:** multi-GPU training (4–8 GPUs)  
+- **Precision:** bfloat16  
+- **Rollout engine:** vLLM  
+- **Training framework:** VERL PPO/GRPO trainer  
+- **Sampling:** multiple rollouts per prompt  
 
 ---
 
-# Ethical Considerations
+## Usage
+
+### Transformers Example
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-- Trained on **de-identified MIMIC-IV data** under the applicable data-use framework
-- **Research-only release**
-- Not suitable for **patient-facing or clinician-facing decision support** without substantial additional validation
-- May propagate **dataset bias and disease-frequency imbalance**
-- Outputs should **not be interpreted as medical advice**
+model_id = "DATEXIS/DeepICD-R1-zero-32B"
 
-Please read the paper’s **Ethical Considerations** and **Limitations** sections before using this model.
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype="auto"
+)
+
+prompt = """
+You are a clinical reasoning model.
+
+Given the following admission note,
+produce reasoning in <think> tags
+and a final ICD-10 diagnosis in <diagnosis> tags.
+
+[ADMISSION NOTE]
+"""
+
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=512,
+)
+
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+## Recommended Inference Practices
+
+- Use prompts consistent with the training format.
+- Validate predicted ICD-10 codes against official code formats.
+- Always review predictions with medical experts.
+- Avoid exposing reasoning traces in safety-critical settings without verification.
 
 ---
 
-# Citation
+## Citation
 
-If you use this model, please cite the paper:
+If you use this model, please cite:
 
 ```bibtex
-@article{roehr2026deepicdr1,
+@inproceedings{roehr2026deepicdr1,
   title={DeepICD-R1: Medical Reasoning through Hierarchical Rewards and Unsupervised Distillation},
-  author={R{\"o}hr, Tom and Steffek, Thomas and Teucher, Roman and Bressem, Keno and Figueroa, Alexei and Grundmann, Paul and Troeger, Peter and Gers, Felix and L{\"o}ser, Alexander},
-  year={2026},
-  journal={Proceedings of LREC-COLING 2026}
+  author={R{\"o}hr, Tom and Steffek, Thomas and Teucher, Roman and Bressem, Keno and others},
+  booktitle={Proceedings of LREC-COLING},
+  year={2026}
 }
 
-
-