scripts/download_split.py

import argparse
import tarfile
from pathlib import Path
from huggingface_hub import hf_hub_download
import os
import tempfile
import shutil


def download_and_unpack(repo_id: str, split_name: str, output_dir: Path):
    """Downloads and cleanly untars a specified data split."""
    final_path = output_dir / split_name
    output_dir.mkdir(exist_ok=True, parents=True)
    archive_name = f"{split_name}.tar.gz"

    archive_path = hf_hub_download(
        repo_id=repo_id,
        filename=archive_name,
        repo_type="dataset",
        local_dir=output_dir,
        local_dir_use_symlinks=False,
    )
    print(f"Successfully downloaded {archive_name}")

    print(f"Uncompressing {archive_name}...")

    with tempfile.TemporaryDirectory(dir=output_dir) as temp_dir:
        temp_extract_path = Path(temp_dir)
        print(f"Uncompressing {archive_name} into temporary directory...")
        with tarfile.open(archive_path, "r:gz") as tar:
            tar.extractall(path=temp_extract_path)

        extracted_items = list(temp_extract_path.iterdir())
        
        if len(extracted_items) == 1 and extracted_items[0].is_dir():
            source_dir = extracted_items[0]
            print(f"Found nested directory: {source_dir.name}")
        else:
            source_dir = temp_extract_path

        shutil.move(str(source_dir), str(final_path))

    Path(archive_path).unlink()
    print(f"Data for '{split_name}' ready in: {final_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Download and unpack a split from the AQcat25 dataset."
    )
    parser.add_argument(
        "--repo-id",
        default="SandboxAQ/aqcat25",
        help="The Hugging Face repository ID.",
    )
    parser.add_argument(
        "--split",
        required=True,
        help="The name of the data split to download (e.g., 'train_id').",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("./aqcat_data"),
        help="Directory to save all downloaded data.",
    )
    args = parser.parse_args()

    download_and_unpack(args.repo_id, args.split, args.output_dir)