From 533adfa256cb46ee0e78364fc6cba74ef65c998e Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 23:01:53 +0800 Subject: [PATCH] Add ModelScope dataset download support --- README.md | 39 +++++++-- SKILL.md | 19 ++++- pyproject.toml | 1 + scripts/download_dataset.py | 165 ++++++++++++++++++++++++++++-------- 4 files changed, 183 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index fe792da..124ca67 100644 --- 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= +# 或 +export MS_TOKEN= +``` + +下载训练和验证数据： + +```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=/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。 diff --git a/SKILL.md b/SKILL.md index 84860b6..c5389c2 100644 --- 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= +./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=/ti_coding_agent_training_probe_20260624 ./scripts/download_dataset.py diff --git a/pyproject.toml b/pyproject.toml index 8319c0d..efc2acf 100644 --- 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", diff --git a/scripts/download_dataset.py b/scripts/download_dataset.py index ee0863a..f43a865 100755 --- 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