Add ModelScope dataset download support

README.md

@@ -1,12 +1,12 @@
 # TI Coding Agent Training Probe
 
-这个仓库用于复现一组 coding-agent SFT probing 实验：从 Hugging Face 下载已经构造好的 Open-SWE-Traces probe 数据集，下载 Qwen3.5-9B 和 Qwen3.6-27B，然后用 ModelScope-SWIFT 依次跑四个 1 epoch 训练任务。
+这个仓库用于复现一组 coding-agent SFT probing 实验：从 ModelScope 下载已经构造好的 Open-SWE-Traces probe 训练/验证数据集，下载 Qwen3.5-9B 和 Qwen3.6-27B，然后用 ModelScope-SWIFT 依次跑四个 1 epoch 训练任务。Hugging Face 数据源仍保留为 fallback。
 
 ## 目录
 
 - `third_party/modelscope-swift/`: ModelScope-SWIFT submodule。
 - `scripts/setup_env.sh`: 一键创建 repo 内 `.venv` 并安装本项目和 SWIFT。
-- `scripts/download_dataset.py`: 下载 Hugging Face 数据集并解压 `train.jsonl`、`validation.jsonl`。
+- `scripts/download_dataset.py`: 下载 ModelScope 或 Hugging Face 数据集并解压 `train.jsonl`、`validation.jsonl`。
 - `scripts/download_models.sh`: 下载 Qwen3.5-9B 和 Qwen3.6-27B 到 `models/`。
 - `scripts/train_qwen35_9b_lora.sh`: Qwen3.5-9B rank=32 LoRA。
 - `scripts/train_qwen35_9b_full.sh`: Qwen3.5-9B bf16 full SFT。
@@ -63,19 +63,48 @@ EXTRA_UV_PACKAGES="flash-attn==2.8.3" ./scripts/setup_env.sh
 
 ## 下载数据集
 
-数据集默认名是 `ti_coding_agent_training_probe_20260624`。如果环境里设置了 `HF_TOKEN`，脚本会用 token owner 自动拼成 `owner/ti_coding_agent_training_probe_20260624`。也可以显式指定：
+默认数据源是 ModelScope：
+
+```text
+eigentom/ti_coding_agent_training_probe_20260624
+```
+
+如果数据集需要鉴权，设置任一 token 环境变量即可：
+
+```bash
+export MODELSCOPE_API_TOKEN=<modelscope-api-key>
+# 或
+export MS_TOKEN=<modelscope-api-key>
+```
+
+下载训练和验证数据：
+
+```bash
+./scripts/download_dataset.py
+```
+
+也可以显式指定 ModelScope repo：
+
+```bash
+export MS_DATASET_REPO_ID=eigentom/ti_coding_agent_training_probe_20260624
+./scripts/download_dataset.py --backend modelscope
+```
+
+Hugging Face 后端仍保留为 fallback。使用 HF 时，如果环境里设置了 `HF_TOKEN`，脚本会用 token owner 自动拼成 `owner/ti_coding_agent_training_probe_20260624`。也可以显式指定：
 
 ```bash
 export HF_ENDPOINT=https://hf-mirror.com
 export HF_DATASET_REPO_ID=<owner>/ti_coding_agent_training_probe_20260624
-./scripts/download_dataset.py
+./scripts/download_dataset.py --backend huggingface
 ```
 
 输出：
 
-- `data/raw/training_probe/`: Hugging Face snapshot。
+- `data/raw/training_probe/`: 原始下载文件。
 - `data/processed/training_probe/train.jsonl`
 - `data/processed/training_probe/validation.jsonl`
+- `data/processed/training_probe/train.parquet`
+- `data/processed/training_probe/validation.parquet`
 
 训练数据里 `system`、`user`、`tool` 消息带 `loss=false`，只有 assistant 轨迹带 `loss=true`。system prompt 会作为上下文参与 attention，但不作为预测目标计算 loss。