Unify README for the PhAI-IDE model family

Present PhAI-IDE-4B, PhAI-IDE-9B, and PhAI-IDE-72B together with a shared model table, verified ScienceAccelBench results, training description, and concise quick start. Retain repository-specific base-model and license metadata.

README.md

@@ -17,24 +17,46 @@ tags:
 - safetensors
 ---
 
-# PhAI-IDE-9B
+# PhAI-IDE
 
-**PhAI-IDE-9B** is a supervised fine-tune of [Qwen/Qwen3.5-9B](https://huggingface.co/Qwen/Qwen3.5-9B) for scientific coding and interaction with tools, trained using [ms-swift](https://github.com/modelscope/ms-swift). The release contains full BF16 weights with the final LoRA adapter merged.
+**PhAI-IDE** is a family of models for scientific coding and interaction with tools, available in **4B, 9B, and 72B** sizes. Each model is supervised fine-tuned with [ms-swift](https://github.com/modelscope/ms-swift) and released as full BF16 weights with the final LoRA adapter merged, together with its configuration and tokenizer.
 
-The training data is sourced from [AItonomy/ScienceIDE](https://huggingface.co/datasets/AItonomy/ScienceIDE).
+The **training dataset is Codex trajectories**, sourced from [AItonomy/ScienceIDE](https://huggingface.co/datasets/AItonomy/ScienceIDE).
+
+## Models
+
+| Model | Base model | BF16 weights | License |
+| --- | --- | ---: | --- |
+| [PhAI-IDE-4B](https://huggingface.co/AItonomy/PhAI-IDE-4B) | [Qwen3.5-4B](https://huggingface.co/Qwen/Qwen3.5-4B) | 9.08 GB | Apache-2.0 |
+| [PhAI-IDE-9B](https://huggingface.co/AItonomy/PhAI-IDE-9B) | [Qwen3.5-9B](https://huggingface.co/Qwen/Qwen3.5-9B) | 18.82 GB | Apache-2.0 |
+| [PhAI-IDE-72B](https://huggingface.co/AItonomy/PhAI-IDE-72B) | [Qwen2.5-72B-Instruct](https://huggingface.co/Qwen/Qwen2.5-72B-Instruct) | 145.41 GB | [Qwen](https://huggingface.co/Qwen/Qwen2.5-72B-Instruct/blob/495f39366efef23836d0cfae4fbe635880d2be31/LICENSE) |
+
+Weight sizes are approximate; inference also requires memory for runtime allocations and the KV cache.
+
+## ScienceAccelBench performance
+
+**PhAI-IDE-9B** improves the task-pooled pass rate from **17.09% to 20.89% (+3.80 percentage points)** across all **27 ScienceAccelBench environments**, using **158 valid paired tasks**.
+
+Task-held-out, localized scientific-code repair on familiar codebases, with original numerical verification. Each row compares the initial model with its corresponding fine-tuned model on identical tasks. Pass rates are percentages; gains are percentage points.
+
+| Model | Environment | Tasks | Initial model | PhAI-IDE | Gain (pp) |
+| --- | --- | ---: | ---: | ---: | ---: |
+| 4B | PLUTO-Particles-Dust | 3 | 0.00 | **33.33** | **+33.33** |
+| 9B | LAPS | 16 | 31.25 | **50.00** | **+18.75** |
+| 9B | MITgcm-biogeo | 8 | 0.00 | **12.50** | **+12.50** |
+| 9B | PLUTO-RMHD | 7 | 0.00 | **28.57** | **+28.57** |
 
 ## Quick start
 
-Use Transformers 5.16.1, PyTorch and Accelerate. The example loads the model, configuration and tokenizer directly from this repository.
+Use Transformers 5.16.1, PyTorch and Accelerate. Set `model_id` to any model in the table above; the example selects the matching model class.
 
 ```python
-from transformers import AutoTokenizer, AutoModelForImageTextToText
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText
 
-model_id = "AItonomy/PhAI-IDE-9B"
+model_id = "AItonomy/PhAI-IDE-4B"
+loader = AutoModelForCausalLM if model_id.endswith("72B") else AutoModelForImageTextToText
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForImageTextToText.from_pretrained(
-    model_id, dtype="bfloat16", device_map="auto",
-)
+model = loader.from_pretrained(model_id, dtype="bfloat16", device_map="auto")
 inputs = tokenizer.apply_chat_template(
     [{"role": "user", "content": "Explain how to verify a numerical simulation."}],
     add_generation_prompt=True, enable_thinking=False, return_dict=True, return_tensors="pt",
@@ -43,29 +65,13 @@ output = model.generate(**inputs, max_new_tokens=128, do_sample=False)
 print(tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
 ```
 
-BF16 weights occupy approximately 18.82 GB, plus runtime memory and KV cache.
-
-## ScienceAccelBench performance
-
-Across all 27 ScienceAccelBench environments, the task-pooled pass rate increases from **17.09% to 20.89% (+3.80 percentage points)** on **158 valid paired tasks**.
-
-Task-held-out, localized scientific-code repair on familiar codebases, with original numerical verification. Pass rates (%); gains in percentage points.
-
-| Environment | Tasks | Qwen3.5-9B | PhAI-IDE-9B | Gain (pp) |
-| --- | ---: | ---: | ---: | ---: |
-| laps | 16 | 31.25 | **50.00** | **+18.75** |
-| mitgcm-biogeo | 8 | 0.00 | **12.50** | **+12.50** |
-| pluto-rmhd-resrmhd | 7 | 0.00 | **28.57** | **+28.57** |
-
 ## Training procedure
 
 ScienceIDE demonstrations were collected with **GPT-5.6-sol** and filtered using a **numerical-equivalence verifier**. They capture code inspection, tool use, and responses to execution feedback.
 
-Training uses **ms-swift** supervised fine-tuning with **LoRA across trainable linear layers for three epochs**. The release merges the final checkpoint's adapter into the base model. Retained assistant targets provide the next-token training signal, while conversation history and tool observations provide context. The Codex-only data retains the native `exec` / `wait` interaction format. Heuristic target masking selects assistant actions for supervision while preserving the surrounding interaction history.
+All three models use **ms-swift** supervised fine-tuning with **LoRA across trainable linear layers for three epochs**. The release merges each final checkpoint's adapter into its base model. Retained assistant targets provide the next-token training signal, while conversation history and tool observations provide context. The trajectories retain the native `exec` / `wait` interaction format. Heuristic target masking selects assistant actions for supervision while preserving the surrounding interaction history.
 
-Fine-tuning uses text trajectories with the vision tower and aligner frozen. The merged export supports standard autoregressive Qwen3.5 generation.
-
-| Setting | Value |
+| Shared setting | Value |
 | --- | --- |
 | Training dataset | Codex trajectories |
 | Training examples / tasks | 4,567 segments / 564 tasks |
@@ -74,6 +80,13 @@ Fine-tuning uses text trajectories with the vision tower and aligner frozen. The
 | Training epochs | 3 |
 | LoRA rank / alpha / dropout | 32 / 64 / 0.05 |
 | Released weights | LoRA merged into BF16 Safetensors |
+
+Long trajectories are organized into segments. Source partition assignments are preserved, with no task identifiers shared between training and validation.
+
+For **4B and 9B**, fine-tuning uses text trajectories with the vision tower and aligner frozen. Their additional settings are:
+
+| Setting | 4B and 9B |
+| --- | --- |
 | Maximum training sequence length | 36,864 tokens |
 | Learning rate / schedule | 2e-5 / cosine |
 | Warmup ratio | 0.03 |
@@ -82,8 +95,6 @@ Fine-tuning uses text trajectories with the vision tower and aligner frozen. The
 | Optimizer steps | 216 |
 | Packing | Disabled |
 
-Long trajectories are organized into segments. Source partition assignments are preserved, with no task identifiers shared between training and validation.
-
 ### Framework versions
 
 The release was validated with the following environment.