| """Base data loader definitions for AEGIS.""" | |
| from abc import ABC, abstractmethod | |
| from dataclasses import dataclass, field | |
| from typing import Any, Dict, List, Optional | |
| import torch | |
| from torch.utils.data import Dataset | |
| class DataSample: | |
| """Standardized data container for judging tasks.""" | |
| input: str | |
| label: Any | |
| metadata: Dict[str, Any] = field(default_factory=dict) | |
| class TorchDatasetWrapper(Dataset): | |
| """PyTorch Dataset wrapper around a list of DataSample objects.""" | |
| def __init__(self, samples: List[DataSample]): | |
| self.samples = samples | |
| def __len__(self) -> int: | |
| return len(self.samples) | |
| def __getitem__(self, idx: int) -> Dict[str, Any]: | |
| sample = self.samples[idx] | |
| return { | |
| "input": sample.input, | |
| "label": sample.label, | |
| "metadata": sample.metadata, | |
| } | |
| class BaseDataLoader(ABC): | |
| """Abstract base class for all AEGIS data loaders.""" | |
| def __init__(self, split: str = "train", max_samples: Optional[int] = None): | |
| self.split = split | |
| self.max_samples = max_samples | |
| self._samples: Optional[List[DataSample]] = None | |
| def load(self) -> List[DataSample]: | |
| """Loads and converts raw dataset to a list of DataSample objects.""" | |
| pass | |
| def get_samples(self) -> List[DataSample]: | |
| if self._samples is None: | |
| samples = self.load() | |
| if self.max_samples is not None: | |
| samples = samples[: self.max_samples] | |
| self._samples = samples | |
| return self._samples | |
| def to_torch_dataset(self) -> Dataset: | |
| """Converts loaded samples into a PyTorch Dataset.""" | |
| return TorchDatasetWrapper(self.get_samples()) | |
| def __len__(self) -> int: | |
| return len(self.get_samples()) | |
| def __getitem__(self, idx: int) -> DataSample: | |
| return self.get_samples()[idx] | |
Xet Storage Details
- Size:
- 1.93 kB
- Xet hash:
- 1908ddeffc2f5c3287d9349e1e330e28818f4cdf5a00c568873f6856ee8e58b5
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.