ti_coding_agent_probe/scripts/download_dataset.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import gzip
import os
import shutil
from pathlib import Path
from urllib.parse import urlparse


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_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:
        raise SystemExit(
            "HF_DATASET_REPO_ID must be owner/name when HF_TOKEN is not set. "
            f"Got unqualified repo name: {raw}"
        )
    owner = HfApi(token=token, endpoint=endpoint).whoami()["name"]
    return f"{owner}/{raw}"


def gunzip_if_needed(src: Path, dst: Path) -> None:
    if dst.exists() and dst.stat().st_size > 0:
        return
    dst.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(src, "rb") as fin, dst.open("wb") as fout:
        shutil.copyfileobj(fin, fout, length=1024 * 1024)


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)


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)
    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


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,
        )
    )


def materialize_processed(raw_dir: Path, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    for name in ("train", "validation"):
        gz = raw_dir / f"{name}.jsonl.gz"
        if gz.exists():
            gunzip_if_needed(gz, out_dir / f"{name}.jsonl")
            link_or_copy(gz, out_dir / f"{name}.jsonl.gz")
        parquet = raw_dir / f"{name}.parquet"
        if parquet.exists():
            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


if __name__ == "__main__":
    raise SystemExit(main())