99 lines
4.3 KiB
Python
99 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
|
from airfrans_frontier.training.config import DataConfig
|
|
from airfrans_frontier.training.data_sources import publish_processed_dataset, resolve_training_data_root
|
|
|
|
remove_pythonpath_entries()
|
|
import numpy as np
|
|
|
|
|
|
class DataSourceTests(unittest.TestCase):
|
|
def test_huggingface_source_downloads_prefix_to_cache(self) -> None:
|
|
calls: list[dict[str, object]] = []
|
|
|
|
def fake_snapshot_download(**kwargs):
|
|
calls.append(kwargs)
|
|
local_dir = Path(str(kwargs["local_dir"]))
|
|
target = local_dir / "processed" / "full"
|
|
target.mkdir(parents=True)
|
|
return str(local_dir)
|
|
|
|
fake_module = types.SimpleNamespace(snapshot_download=fake_snapshot_download)
|
|
with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {}, clear=False):
|
|
tmp_path = Path(tmp)
|
|
config = DataConfig(
|
|
root=tmp_path / "configured-root",
|
|
train_cases=1,
|
|
val_cases=0,
|
|
test_cases=0,
|
|
points_per_case=1,
|
|
all_points_per_case=False,
|
|
batch_size=1,
|
|
source="huggingface",
|
|
hf_repo_id="owner/airfrans-processed",
|
|
hf_repo_type="dataset",
|
|
hf_path_prefix="processed/full",
|
|
cache_dir=tmp_path / "cache",
|
|
)
|
|
|
|
resolved = resolve_training_data_root(config)
|
|
|
|
self.assertEqual(resolved, tmp_path / "cache" / "processed" / "full")
|
|
self.assertEqual(calls[0]["repo_id"], "owner/airfrans-processed")
|
|
self.assertEqual(calls[0]["repo_type"], "dataset")
|
|
self.assertEqual(calls[0]["allow_patterns"], ["processed/full/**"])
|
|
|
|
def test_publish_processed_dataset_uploads_folder_and_manifest(self) -> None:
|
|
created: list[tuple[str, str, bool]] = []
|
|
uploaded_folders: list[tuple[str, str]] = []
|
|
uploaded_files: list[str] = []
|
|
|
|
class FakeApi:
|
|
def __init__(self, token: str) -> None:
|
|
self.token = token
|
|
|
|
def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None:
|
|
created.append((repo_id, repo_type, private))
|
|
|
|
def upload_folder(self, *, repo_id: str, repo_type: str, folder_path: str, path_in_repo: str, commit_message: str):
|
|
uploaded_folders.append((folder_path, path_in_repo))
|
|
return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/abc", oid="abc")
|
|
|
|
def upload_file(self, *, repo_id: str, repo_type: str, path_or_fileobj: str, path_in_repo: str, commit_message: str):
|
|
uploaded_files.append(path_in_repo)
|
|
return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/def", oid="def")
|
|
|
|
fake_module = types.SimpleNamespace(HfApi=FakeApi)
|
|
with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}):
|
|
root = Path(tmp) / "processed"
|
|
root.mkdir()
|
|
np.savez(root / "case_00.npz", features=np.zeros((2, 2), dtype=np.float32), targets=np.zeros((2, 1), dtype=np.float32))
|
|
manifest_path = Path(tmp) / "manifest.json"
|
|
|
|
manifest = publish_processed_dataset(
|
|
data_root=root,
|
|
repo_id="owner/repo",
|
|
path_in_repo="processed/full",
|
|
manifest_out=manifest_path,
|
|
)
|
|
|
|
self.assertEqual(created, [("owner/repo", "dataset", False)])
|
|
self.assertEqual(uploaded_folders, [(str(root), "processed/full")])
|
|
self.assertEqual(uploaded_files, ["processed/full/hf_dataset_manifest.json"])
|
|
self.assertEqual(manifest["npz_file_count"], 1)
|
|
self.assertTrue(manifest_path.is_file())
|
|
self.assertEqual(json.loads(manifest_path.read_text())["uploaded_manifest_path"], "processed/full/hf_dataset_manifest.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|