Add ModelScope dataset download support

scripts/download_dataset.py

@@ -6,14 +6,25 @@ import gzip
 import os
 import shutil
 from pathlib import Path
-
-from huggingface_hub import HfApi, snapshot_download
+from urllib.parse import urlparse
 
 
-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:
+    from huggingface_hub import HfApi
+
     if "/" in raw:
         return raw
     if not token:
@@ -33,50 +44,139 @@ 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:
+    import requests
+
+    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:
+    from huggingface_hub import snapshot_download
+
+    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