code/local_eval/download.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path

from .constants import (
    DEFAULT_DATA_ROOT,
    DEFAULT_MODEL_DIR,
    GENESIS_REPO,
    GENESIS_REVISION,
    GENESIS_WEIGHT_HASH,
    REQUIRED_FILES,
)

LITE_MINI_CODER_REPO = "ricdomolm/mini-coder-trajs-400k"
LITE_MINI_CODER_SHARDS = (
    "data/train-00000-of-00060.parquet",
    "data/train-00001-of-00060.parquet",
)


def download_genesis(dest: Path = DEFAULT_MODEL_DIR, *, revision: str = GENESIS_REVISION) -> Path:
    from huggingface_hub import snapshot_download

    dest.mkdir(parents=True, exist_ok=True)
    print(f"downloading {GENESIS_REPO}@{revision} -> {dest}", flush=True)
    snapshot_download(
        repo_id=GENESIS_REPO,
        revision=revision,
        local_dir=str(dest),
        max_workers=8,
    )
    pin = {
        "repo": GENESIS_REPO,
        "revision": revision,
        "weight_hash": GENESIS_WEIGHT_HASH,
    }
    (dest / ".albedo-genesis-pin.json").write_text(json.dumps(pin, indent=2) + "\n")
    return dest


def download_lite_data(root: Path = DEFAULT_DATA_ROOT) -> Path:
    from huggingface_hub import hf_hub_download

    dest = root / "mini-coder"
    dest.mkdir(parents=True, exist_ok=True)
    print(f"downloading lite mini-coder shards -> {dest}", flush=True)
    for rel in LITE_MINI_CODER_SHARDS:
        path = hf_hub_download(
            LITE_MINI_CODER_REPO,
            rel,
            repo_type="dataset",
            local_dir=str(dest),
        )
        print(f"  {rel} -> {path}", flush=True)
    return root


def verify_model_dir(path: Path) -> dict:
    path = Path(path)
    if not path.is_dir():
        return {"ok": False, "reason": f"not a directory: {path}"}

    files = {p.name for p in path.iterdir() if p.is_file()}
    missing = [name for name in REQUIRED_FILES if name not in files]
    extras = sorted(
        name
        for name in files
        if name not in REQUIRED_FILES
        and name not in {"model.safetensors.index.json", ".gitattributes", "LICENSE", "README.md"}
        and not name.endswith(".safetensors")
        and not name.startswith(".")
    )
    shards = sorted(name for name in files if name.endswith(".safetensors"))
    shard_bytes = sum((path / name).stat().st_size for name in shards)
    report = {
        "ok": not missing and bool(shards),
        "path": str(path),
        "missing_required": missing,
        "unexpected_extras": extras,
        "shard_count": len(shards),
        "shard_gib": round(shard_bytes / 1024**3, 2),
        "has_index": "model.safetensors.index.json" in files,
    }
    if missing:
        report["reason"] = f"missing required files: {missing}"
    elif not shards:
        report["reason"] = "no safetensors shards"
        report["ok"] = False
    return report


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()