File size: 2,729 Bytes
bd3493c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """Datasets for training Aardvark from official-schema pickle tasks."""
from __future__ import annotations
import copy
import pickle
import random
from pathlib import Path
from typing import Any
import numpy as np
import torch
from torch.utils.data import Dataset
from .aardvark_adapter import validate_sample
def discover_samples(path: Path) -> list[Path]:
path = path.resolve()
samples = sorted(path.glob("*.pkl")) if path.is_dir() else [path]
if not samples or any(not sample.is_file() for sample in samples):
raise FileNotFoundError(f"No Aardvark sample pickle found at {path}")
for sample in samples:
validate_sample(sample)
return samples
def split_samples(samples: list[Path], validation_fraction: float, seed: int) -> tuple[list[Path], list[Path]]:
if not 0.0 < validation_fraction < 1.0:
raise ValueError("validation_fraction must be between 0 and 1")
if len(samples) == 1:
return samples, samples
shuffled = samples.copy()
random.Random(seed).shuffle(shuffled)
validation_count = max(1, round(len(shuffled) * validation_fraction))
validation_count = min(validation_count, len(shuffled) - 1)
return shuffled[validation_count:], shuffled[:validation_count]
class AardvarkTaskDataset(Dataset):
"""Repeat one or more already-batched official tasks for a fixed number of steps."""
def __init__(self, samples: list[Path], steps: int) -> None:
if steps < 1:
raise ValueError("steps must be at least 1")
self.tasks = [self._load(path) for path in samples]
self.steps = steps
@staticmethod
def _load(path: Path) -> dict[str, Any]:
with path.open("rb") as handle:
return pickle.load(handle)
def __len__(self) -> int:
return self.steps
def __getitem__(self, index: int) -> dict[str, Any]:
return copy.deepcopy(self.tasks[index % len(self.tasks)])
def collate_tasks(items: list[Any]) -> Any:
"""Concatenate the batch dimension already present in official task tensors."""
first = items[0]
if isinstance(first, torch.Tensor):
return torch.cat(items, dim=0)
if isinstance(first, np.ndarray):
return np.concatenate(items, axis=0)
if isinstance(first, dict):
return {key: collate_tasks([item[key] for item in items]) for key in first}
if isinstance(first, tuple):
return tuple(collate_tasks(list(values)) for values in zip(*items))
if isinstance(first, list):
return [collate_tasks(list(values)) for values in zip(*items)]
if all(item == first for item in items):
return first
raise TypeError(f"Cannot collate Aardvark values of type {type(first).__name__}")
|