| from __future__ import annotations |
|
|
| import random |
| from collections.abc import Iterator, Sequence |
| from pathlib import Path |
| from typing import Any |
|
|
| from .common import write_jsonl |
|
|
|
|
| class DataObject(Sequence[dict[str, str]]): |
| def __init__( |
| self, |
| records: Sequence[dict[str, str]] | None = None, |
| *, |
| name: str = "dataset", |
| warnings: Sequence[str] | None = None, |
| ) -> None: |
| self.records = list(records or []) |
| self.name = name |
| self.warnings = list(warnings or []) |
|
|
| def __len__(self) -> int: |
| return len(self.records) |
|
|
| def __iter__(self) -> Iterator[dict[str, str]]: |
| return iter(self.records) |
|
|
| def __getitem__(self, index: int | slice) -> dict[str, str] | list[dict[str, str]]: |
| return self.records[index] |
|
|
| def take(self, amount: int) -> list[dict[str, str]]: |
| if amount < 0: |
| raise ValueError("amount must be non-negative") |
|
|
| return self.records[:amount] |
|
|
| def shuffle(self, *, seed: int | None = 42) -> "DataObject": |
| shuffled = list(self.records) |
| rng = random.Random(seed) |
| rng.shuffle(shuffled) |
|
|
| return DataObject( |
| shuffled, |
| name=self.name, |
| warnings=self.warnings, |
| ) |
|
|
| def to_jsonl(self, path: str | Path) -> None: |
| write_jsonl(self.records, path) |
|
|
| def extend(self, other: "DataObject") -> None: |
| self.records.extend(other.records) |
| self.warnings.extend(other.warnings) |
|
|
| @classmethod |
| def concat(cls, datasets: Sequence["DataObject"], *, name: str = "all") -> "DataObject": |
| records: list[dict[str, str]] = [] |
| warnings: list[str] = [] |
|
|
| for dataset in datasets: |
| records.extend(dataset.records) |
| warnings.extend(dataset.warnings) |
|
|
| return cls(records, name=name, warnings=warnings) |
|
|
| def summary(self) -> dict[str, Any]: |
| return { |
| "name": self.name, |
| "rows": len(self.records), |
| "warnings": self.warnings, |
| } |
|
|
