File size: 9,478 Bytes
35d483e | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 | """Lazy local-Parquet and Hugging Face dataset ingestion."""
from __future__ import annotations
import csv
import gzip
import json
from pathlib import Path
from typing import Any, Iterable, Iterator, Mapping, Sequence
class OptionalDependencyError(ImportError):
"""An optional data backend is required for the requested source."""
class DatasetReadError(RuntimeError):
"""A source exists but could not be decoded into records."""
def _open_text(path: Path):
if path.suffix == ".gz":
return gzip.open(path, "rt", encoding="utf-8", newline="")
return path.open("r", encoding="utf-8", newline="")
def _with_provenance(row: Mapping[str, Any], source: str, index: int) -> dict[str, Any]:
result = dict(row)
result.setdefault("__source_file", source)
result.setdefault("__source_row", index)
return result
def _iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
with _open_text(path) as handle:
for index, line in enumerate(handle):
if not line.strip():
continue
try:
row = json.loads(line)
except json.JSONDecodeError as exc:
raise DatasetReadError(f"invalid JSON at {path}:{index + 1}: {exc}") from exc
if not isinstance(row, Mapping):
raise DatasetReadError(f"expected an object at {path}:{index + 1}")
yield _with_provenance(row, str(path), index)
def _iter_json(path: Path) -> Iterator[dict[str, Any]]:
try:
with _open_text(path) as handle:
payload = json.load(handle)
except (OSError, json.JSONDecodeError) as exc:
raise DatasetReadError(f"cannot read JSON source {path}: {exc}") from exc
if isinstance(payload, Mapping):
for key in ("records", "data", "examples", "rows"):
if key in payload:
payload = payload[key]
break
if not isinstance(payload, list):
raise DatasetReadError(f"JSON source {path} must contain a list of records")
for index, row in enumerate(payload):
if not isinstance(row, Mapping):
raise DatasetReadError(f"expected an object at {path} record {index}")
yield _with_provenance(row, str(path), index)
def _iter_csv(path: Path) -> Iterator[dict[str, Any]]:
with _open_text(path) as handle:
reader = csv.DictReader(handle)
for index, row in enumerate(reader):
yield _with_provenance(row, str(path), index)
def _import_pyarrow_parquet():
try:
import pyarrow.parquet as parquet # type: ignore[import-not-found]
except ImportError as exc:
raise OptionalDependencyError(
"reading local Parquet requires pyarrow; install the project's data dependencies"
) from exc
return parquet
def iter_parquet_records(
paths: Sequence[str | Path],
*,
batch_size: int = 256,
columns: Sequence[str] | None = None,
) -> Iterator[dict[str, Any]]:
"""Stream rows from Parquet files without decoding Hugging Face audio."""
if batch_size <= 0:
raise ValueError("batch_size must be positive")
parquet = _import_pyarrow_parquet()
for raw_path in paths:
path = Path(raw_path)
try:
parquet_file = parquet.ParquetFile(str(path))
row_index = 0
for batch in parquet_file.iter_batches(batch_size=batch_size, columns=columns):
for row in batch.to_pylist():
if not isinstance(row, Mapping):
raise DatasetReadError(f"non-object Parquet row in {path}")
yield _with_provenance(row, str(path), row_index)
row_index += 1
except OptionalDependencyError:
raise
except Exception as exc:
raise DatasetReadError(f"cannot read Parquet source {path}: {exc}") from exc
def _local_files(source: Path) -> list[Path]:
if source.is_file():
return [source]
if not source.exists():
raise FileNotFoundError(f"dataset source does not exist: {source}")
if not source.is_dir():
raise DatasetReadError(f"dataset source is not a file or directory: {source}")
supported = {".parquet", ".jsonl", ".ndjson", ".json", ".csv"}
files = [path for path in source.rglob("*") if path.is_file() and path.suffix.lower() in supported]
files.extend(
path
for path in source.rglob("*.jsonl.gz")
if path.is_file() and path not in files
)
return sorted(files, key=lambda item: item.as_posix())
def iter_local_records(
source: str | Path | Sequence[str | Path],
*,
batch_size: int = 256,
columns: Sequence[str] | None = None,
) -> Iterator[dict[str, Any]]:
"""Iterate records from local Parquet, JSON(L), or CSV sources."""
roots = [source] if isinstance(source, (str, Path)) else list(source)
files: list[Path] = []
for root in roots:
files.extend(_local_files(Path(root)))
files = sorted(set(files), key=lambda item: item.as_posix())
if not files:
raise DatasetReadError("no supported dataset files were found")
parquet_paths = [path for path in files if path.suffix.lower() == ".parquet"]
other_paths = [path for path in files if path.suffix.lower() != ".parquet"]
if parquet_paths:
yield from iter_parquet_records(parquet_paths, batch_size=batch_size, columns=columns)
for path in other_paths:
name = path.name.lower()
if name.endswith((".jsonl", ".ndjson", ".jsonl.gz")):
yield from _iter_jsonl(path)
elif name.endswith(".json"):
yield from _iter_json(path)
elif name.endswith(".csv"):
yield from _iter_csv(path)
def _normalize_hf_id(source: str) -> str:
value = source
if value.startswith("hf://datasets/"):
value = value[len("hf://datasets/") :]
elif value.startswith("hf://"):
value = value[len("hf://") :]
return value.strip("/")
def iter_hf_records(
dataset_id: str,
*,
split: str = "train",
revision: str | None = None,
token: str | None = None,
streaming: bool = True,
) -> Iterator[dict[str, Any]]:
"""Stream raw rows from a Hugging Face dataset with audio decoding disabled."""
try:
import datasets # type: ignore[import-not-found]
except ImportError as exc:
raise OptionalDependencyError(
"reading a Hugging Face dataset requires the 'datasets' package"
) from exc
dataset_id = _normalize_hf_id(dataset_id)
kwargs: dict[str, Any] = {"split": split, "streaming": streaming}
if revision is not None:
kwargs["revision"] = revision
if token is not None:
kwargs["token"] = token
try:
dataset = datasets.load_dataset(dataset_id, **kwargs)
except TypeError as exc:
# Compatibility with older datasets releases which used use_auth_token.
if token is None or "token" not in str(exc):
raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {exc}") from exc
kwargs.pop("token", None)
kwargs["use_auth_token"] = token
try:
dataset = datasets.load_dataset(dataset_id, **kwargs)
except Exception as retry_exc:
raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {retry_exc}") from retry_exc
except Exception as exc:
raise DatasetReadError(f"cannot load Hugging Face dataset {dataset_id}: {exc}") from exc
# ``decode=False`` keeps encoded bytes available and avoids torchcodec/ffmpeg.
try:
if "audio" in dataset.features:
dataset = dataset.cast_column("audio", datasets.Audio(decode=False))
except (AttributeError, TypeError, ValueError):
# Some custom builders already return undecoded mappings or do not expose
# cast_column in streaming mode. Defer representation validation to audio.py.
pass
provenance = f"hf://datasets/{dataset_id}@{revision or 'default'}:{split}"
for index, row in enumerate(dataset):
if not isinstance(row, Mapping):
raise DatasetReadError(f"Hugging Face dataset {dataset_id} yielded a non-object row")
yield _with_provenance(row, provenance, index)
def _looks_like_hf_id(source: str) -> bool:
if source.startswith("hf://"):
return True
return source.count("/") == 1 and not source.startswith(("./", "../", "/"))
def iter_records(
source: str | Path | Sequence[str | Path],
*,
split: str = "train",
revision: str | None = None,
token: str | None = None,
streaming: bool = True,
batch_size: int = 256,
columns: Sequence[str] | None = None,
limit: int | None = None,
) -> Iterator[dict[str, Any]]:
"""Iterate local or Hugging Face rows through one stable entry point."""
if limit is not None and limit < 0:
raise ValueError("limit cannot be negative")
if isinstance(source, str) and _looks_like_hf_id(source) and not Path(source).exists():
iterator: Iterable[dict[str, Any]] = iter_hf_records(
source,
split=split,
revision=revision,
token=token,
streaming=streaming,
)
else:
iterator = iter_local_records(source, batch_size=batch_size, columns=columns)
for index, row in enumerate(iterator):
if limit is not None and index >= limit:
break
yield row
|