Add ModelScope dataset download support

2026-06-24 23:01:53 +08:00
parent 4063246ac2
commit 533adfa256
4 changed files with 183 additions and 41 deletions
--- a/README.md
+++ b/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。

--- a/SKILL.md
+++ b/SKILL.md
@@ -1,6 +1,6 @@
 # TI Coding Agent Probe Skill

-Use this skill when the user wants to run or modify the TI coding-agent SFT probe experiments based on the Hugging Face dataset `ti_coding_agent_training_probe_20260624`.
+Use this skill when the user wants to run or modify the TI coding-agent SFT probe experiments based on the ModelScope dataset `eigentom/ti_coding_agent_training_probe_20260624`. Hugging Face remains available as a fallback backend.

 ## Repository Contract

@@ -25,9 +25,24 @@ git submodule update --init --recursive
 ./scripts/setup_env.sh
 ```

-Download the training probe dataset:
+Download the training/eval probe dataset from ModelScope:

 ```bash
+export MODELSCOPE_API_TOKEN=<modelscope-api-key>
+./scripts/download_dataset.py
+```
+
+The downloader writes:
+
+- `data/raw/training_probe/{train,validation}.jsonl.gz`
+- `data/raw/training_probe/{train,validation}.parquet`
+- `data/processed/training_probe/{train,validation}.jsonl`
+- `data/processed/training_probe/{train,validation}.parquet`
+
+Use Hugging Face fallback only when requested:
+
+```bash
+export DATASET_BACKEND=huggingface
 export HF_ENDPOINT=https://hf-mirror.com
 export HF_DATASET_REPO_ID=<owner>/ti_coding_agent_training_probe_20260624
 ./scripts/download_dataset.py
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,6 +9,7 @@ description = "Reproducible SWIFT SFT probe scripts for Open-SWE-Traces derived
 requires-python = ">=3.10"
 dependencies = [
  "huggingface_hub>=0.23",
+  "modelscope>=1.37",
  "datasets>=2.20",
  "pyarrow>=15",
  "tensorboard>=2.15",
--- a/scripts/download_dataset.py
+++ b/scripts/download_dataset.py
@@ -6,14 +6,26 @@ import gzip
 import os
 import shutil
 from pathlib import Path
+from urllib.parse import urlparse

+import requests
 from huggingface_hub import HfApi, snapshot_download


-DEFAULT_REPO_NAME = "ti_coding_agent_training_probe_20260624"
+DEFAULT_BACKEND = "modelscope"
+DEFAULT_MS_DATASET_ID = "eigentom/ti_coding_agent_training_probe_20260624"
+DEFAULT_HF_REPO_NAME = "ti_coding_agent_training_probe_20260624"
+DOWNLOAD_FILES = [
+    "README.md",
+    "metadata.json",
+    "train.parquet",
+    "validation.parquet",
+    "train.jsonl.gz",
+    "validation.jsonl.gz",
+]


-def resolve_dataset_id(raw: str, token: str | None, endpoint: str | None) -> str:
+def resolve_hf_dataset_id(raw: str, token: str | None, endpoint: str | None) -> str:
    if "/" in raw:
        return raw
    if not token:
@@ -33,50 +45,135 @@ def gunzip_if_needed(src: Path, dst: Path) -> None:
        shutil.copyfileobj(fin, fout, length=1024 * 1024)


-def main() -> int:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset-id", default=os.environ.get("HF_DATASET_REPO_ID", DEFAULT_REPO_NAME))
-    parser.add_argument("--raw-dir", default="data/raw/training_probe")
-    parser.add_argument("--out-dir", default="data/processed/training_probe")
-    args = parser.parse_args()
+def link_or_copy(src: Path, dst: Path) -> None:
+    if dst.exists():
+        return
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        dst.symlink_to(src.resolve())
+    except OSError:
+        shutil.copy2(src, dst)

-    token = os.environ.get("HF_TOKEN")
-    endpoint = os.environ.get("HF_ENDPOINT")
-    dataset_id = resolve_dataset_id(args.dataset_id, token, endpoint)

-    raw_dir = Path(args.raw_dir)
-    out_dir = Path(args.out_dir)
+def stream_download(url: str, dst: Path, token: str | None) -> None:
+    if dst.exists() and dst.stat().st_size > 0:
+        return
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    headers = {}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    with requests.get(url, headers=headers, stream=True, timeout=60) as response:
+        response.raise_for_status()
+        tmp = dst.with_suffix(dst.suffix + ".tmp")
+        with tmp.open("wb") as handle:
+            for chunk in response.iter_content(chunk_size=16 * 1024 * 1024):
+                if chunk:
+                    handle.write(chunk)
+        tmp.replace(dst)
+
+
+def download_from_modelscope(dataset_id: str, raw_dir: Path, token: str | None, files: list[str]) -> Path:
+    from modelscope.hub.api import HubApi
+
+    if "/" not in dataset_id:
+        raise SystemExit(f"ModelScope dataset id must be namespace/name. Got: {dataset_id}")
+    namespace, dataset_name = dataset_id.split("/", 1)
+    api = HubApi()
    raw_dir.mkdir(parents=True, exist_ok=True)
-    out_dir.mkdir(parents=True, exist_ok=True)
+    for file_name in files:
+        dst = raw_dir / file_name
+        if dst.exists() and dst.stat().st_size > 0:
+            continue
+        url = api.get_dataset_file_url(file_name=file_name, dataset_name=dataset_name, namespace=namespace)
+        if isinstance(url, dict):
+            url = url.get("url") or url.get("Url") or url.get("download_url")
+        if not isinstance(url, str) or not urlparse(url).scheme:
+            raise RuntimeError(f"Could not resolve ModelScope URL for {dataset_id}/{file_name}: {url!r}")
+        print(f"DOWNLOADING modelscope://{dataset_id}/{file_name} -> {dst}", flush=True)
+        stream_download(url, dst, token)
+    return raw_dir

-    local_path = snapshot_download(
-        repo_id=dataset_id,
-        repo_type="dataset",
-        local_dir=raw_dir,
-        token=token,
-        endpoint=endpoint,
-        allow_patterns=[
-            "README.md",
-            "metadata.json",
-            "train.parquet",
-            "validation.parquet",
-            "train.jsonl.gz",
-            "validation.jsonl.gz",
-        ],
+
+def download_from_huggingface(
+    dataset_id: str,
+    raw_dir: Path,
+    token: str | None,
+    endpoint: str | None,
+    files: list[str],
+) -> Path:
+    return Path(
+        snapshot_download(
+            repo_id=dataset_id,
+            repo_type="dataset",
+            local_dir=raw_dir,
+            token=token,
+            endpoint=endpoint,
+            allow_patterns=files,
+        )
    )
-    local = Path(local_path)
+
+
+def materialize_processed(raw_dir: Path, out_dir: Path) -> None:
+    out_dir.mkdir(parents=True, exist_ok=True)
    for name in ("train", "validation"):
-        gz = local / f"{name}.jsonl.gz"
+        gz = raw_dir / f"{name}.jsonl.gz"
        if gz.exists():
            gunzip_if_needed(gz, out_dir / f"{name}.jsonl")
-        parquet = local / f"{name}.parquet"
+            link_or_copy(gz, out_dir / f"{name}.jsonl.gz")
+        parquet = raw_dir / f"{name}.parquet"
        if parquet.exists():
-            target = out_dir / f"{name}.parquet"
-            if not target.exists():
-                target.symlink_to(parquet.resolve())
+            link_or_copy(parquet, out_dir / f"{name}.parquet")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Download training/eval data for TI coding-agent probe experiments.")
+    parser.add_argument(
+        "--backend",
+        choices=["modelscope", "huggingface"],
+        default=os.environ.get("DATASET_BACKEND", DEFAULT_BACKEND),
+        help="Dataset hosting backend. Defaults to ModelScope.",
+    )
+    parser.add_argument(
+        "--dataset-id",
+        default=os.environ.get("MS_DATASET_REPO_ID")
+        or os.environ.get("MODELSCOPE_DATASET_REPO_ID")
+        or os.environ.get("HF_DATASET_REPO_ID"),
+        help="Dataset repo id. Defaults depend on backend.",
+    )
+    parser.add_argument("--raw-dir", default="data/raw/training_probe")
+    parser.add_argument("--out-dir", default="data/processed/training_probe")
+    parser.add_argument(
+        "--files",
+        default=",".join(DOWNLOAD_FILES),
+        help="Comma-separated repo files to download. Defaults to the full train/eval probe set.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    raw_dir = Path(args.raw_dir)
+    out_dir = Path(args.out_dir)
+    files = [item.strip() for item in args.files.split(",") if item.strip()]
+
+    if args.backend == "modelscope":
+        dataset_id = args.dataset_id or DEFAULT_MS_DATASET_ID
+        token = os.environ.get("MODELSCOPE_API_TOKEN") or os.environ.get("MODELSCOPE_TOKEN") or os.environ.get("MS_TOKEN")
+        local = download_from_modelscope(dataset_id, raw_dir, token, files)
+    else:
+        endpoint = os.environ.get("HF_ENDPOINT")
+        token = os.environ.get("HF_TOKEN")
+        dataset_id = resolve_hf_dataset_id(args.dataset_id or DEFAULT_HF_REPO_NAME, token, endpoint)
+        local = download_from_huggingface(dataset_id, raw_dir, token, endpoint, files)
+
+    materialize_processed(local, out_dir)
+    print(f"DATASET_BACKEND={args.backend}")
    print(f"DATASET_ID={dataset_id}")
+    print(f"RAW_DIR={local}")
    print(f"TRAIN_JSONL={out_dir / 'train.jsonl'}")
    print(f"VALIDATION_JSONL={out_dir / 'validation.jsonl'}")
+    print(f"TRAIN_PARQUET={out_dir / 'train.parquet'}")
+    print(f"VALIDATION_PARQUET={out_dir / 'validation.parquet'}")
    return 0