File size: 9,796 Bytes
7c5e40e | 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 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 | """Typed dataset registry for STRATA acquisition.
The registry is the single, version-controlled source of truth describing which
datasets STRATA uses, where they legitimately come from, their licensing tier,
and which languages they cover. It is intentionally declarative: acquisition
code in :mod:`strata.data.acquire` interprets these specs, but no download URLs
or credentials live in Python source.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
import yaml
# ---------------------------------------------------------------------------
# Controlled vocabularies
# ---------------------------------------------------------------------------
# Licensing / accessibility tiers. These drive default fetch behaviour.
TIER_OPEN = "open" # freely downloadable, no gate
TIER_GATED_AUTO = "gated_auto" # HF gate with automatic approval on request
TIER_GATED_MANUAL = "gated_manual" # HF gate requiring manual approval
TIER_LICENSE_REQUIRED = "license_required" # e.g. LDC / signed agreement; no auto DL
LICENSE_TIERS = frozenset(
{TIER_OPEN, TIER_GATED_AUTO, TIER_GATED_MANUAL, TIER_LICENSE_REQUIRED}
)
# Source kinds understood by the acquirer.
KIND_HF_SNAPSHOT = "hf_snapshot" # snapshot_download of a HF dataset repo
KIND_HF_FILES = "hf_files" # capped/prefix-filtered file selection from a HF repo
KIND_HTTP = "http" # direct file/tarball download(s)
KIND_GITHUB_ARCHIVE = "github_archive" # pinned GitHub repo tarball
KIND_LICENSE_REQUIRED = "license_required" # no download; write acquisition stub
SOURCE_KINDS = frozenset(
{
KIND_HF_SNAPSHOT,
KIND_HF_FILES,
KIND_HTTP,
KIND_GITHUB_ARCHIVE,
KIND_LICENSE_REQUIRED,
}
)
# Pipeline layers, matching the STRATA data stack. "multi" is used in a spec's
# languages list to mean "many languages, filtered downstream".
LANG_MULTI = "multi"
# ---------------------------------------------------------------------------
# Spec
# ---------------------------------------------------------------------------
@dataclass(frozen=True, slots=True)
class DatasetSpec:
"""Declarative description of one dataset."""
id: str
title: str
layer: str
languages: tuple[str, ...]
license: str
license_tier: str
source: dict[str, object]
homepage: str = ""
notes: str = ""
optional: bool = False
def __post_init__(self) -> None:
self.validate()
@property
def kind(self) -> str:
return str(self.source.get("kind", ""))
def covers_language(self, language: str) -> bool:
return language in self.languages or LANG_MULTI in self.languages
def validate(self) -> None:
if not self.id or not self.id.replace("_", "").replace("-", "").isalnum():
raise ValueError(
f"dataset id must be non-empty alphanumeric/_/-, got {self.id!r}"
)
if self.license_tier not in LICENSE_TIERS:
raise ValueError(
f"[{self.id}] invalid license_tier {self.license_tier!r}; "
f"expected one of {sorted(LICENSE_TIERS)}"
)
if not self.languages:
raise ValueError(f"[{self.id}] languages must be non-empty")
kind = self.kind
if kind not in SOURCE_KINDS:
raise ValueError(
f"[{self.id}] invalid source kind {kind!r}; "
f"expected one of {sorted(SOURCE_KINDS)}"
)
_validate_source(self.id, kind, self.source)
if kind == KIND_LICENSE_REQUIRED and self.license_tier != TIER_LICENSE_REQUIRED:
raise ValueError(
f"[{self.id}] source kind {KIND_LICENSE_REQUIRED} requires "
f"license_tier {TIER_LICENSE_REQUIRED}"
)
def _require(dataset_id: str, source: dict[str, object], *keys: str) -> None:
for key in keys:
if not source.get(key):
raise ValueError(
f"[{dataset_id}] source kind {source.get('kind')!r} requires "
f"non-empty {key!r}"
)
def _validate_source(dataset_id: str, kind: str, source: dict[str, object]) -> None:
if kind == KIND_HF_SNAPSHOT:
_require(dataset_id, source, "repo_id")
elif kind == KIND_HF_FILES:
_require(dataset_id, source, "repo_id", "prefixes")
if not isinstance(source["prefixes"], list):
raise ValueError(f"[{dataset_id}] hf_files.prefixes must be a list")
elif kind == KIND_HTTP:
urls = source.get("urls")
if not isinstance(urls, list) or not urls:
raise ValueError(f"[{dataset_id}] http.urls must be a non-empty list")
for item in urls:
if not isinstance(item, dict) or not item.get("url"):
raise ValueError(f"[{dataset_id}] each http url needs a 'url' field")
elif kind == KIND_GITHUB_ARCHIVE:
_require(dataset_id, source, "repo")
elif kind == KIND_LICENSE_REQUIRED:
_require(dataset_id, source, "acquisition")
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
@dataclass(frozen=True, slots=True)
class Registry:
version: int
languages_core: tuple[str, ...]
datasets: tuple[DatasetSpec, ...]
description: str = ""
def by_id(self, dataset_id: str) -> DatasetSpec:
for spec in self.datasets:
if spec.id == dataset_id:
return spec
raise KeyError(f"unknown dataset id: {dataset_id!r}")
def ids(self) -> list[str]:
return [spec.id for spec in self.datasets]
def layers(self) -> list[str]:
seen: dict[str, None] = {}
for spec in self.datasets:
seen.setdefault(spec.layer, None)
return list(seen)
def filter(
self,
*,
ids: list[str] | None = None,
layers: list[str] | None = None,
languages: list[str] | None = None,
tiers: list[str] | None = None,
include_optional: bool = True,
) -> list[DatasetSpec]:
"""Return specs matching all provided filters (AND semantics).
``ids`` matching is explicit and always includes optional datasets, so
an operator can request an optional dataset by name.
"""
id_set = set(ids) if ids else None
layer_set = set(layers) if layers else None
tier_set = set(tiers) if tiers else None
selected: list[DatasetSpec] = []
for spec in self.datasets:
if id_set is not None and spec.id not in id_set:
continue
if layer_set is not None and spec.layer not in layer_set:
continue
if tier_set is not None and spec.license_tier not in tier_set:
continue
if languages and not any(spec.covers_language(lang) for lang in languages):
continue
explicitly_named = id_set is not None and spec.id in id_set
if spec.optional and not include_optional and not explicitly_named:
continue
selected.append(spec)
return selected
def load_registry(path: Path) -> Registry:
"""Load and validate a registry YAML file."""
if not path.exists():
raise FileNotFoundError(f"registry file not found: {path}")
data = yaml.safe_load(path.read_text(encoding="utf-8"))
if not isinstance(data, dict):
raise ValueError(f"registry {path} must be a mapping at the top level")
raw_datasets = data.get("datasets")
if not isinstance(raw_datasets, list) or not raw_datasets:
raise ValueError(f"registry {path} must contain a non-empty 'datasets' list")
specs: list[DatasetSpec] = []
seen_ids: set[str] = set()
for raw in raw_datasets:
if not isinstance(raw, dict):
raise ValueError(f"registry {path}: each dataset entry must be a mapping")
spec = _spec_from_dict(raw)
if spec.id in seen_ids:
raise ValueError(f"registry {path}: duplicate dataset id {spec.id!r}")
seen_ids.add(spec.id)
specs.append(spec)
return Registry(
version=int(data.get("version", 1)),
languages_core=tuple(data.get("languages_core", []) or []),
datasets=tuple(specs),
description=str(data.get("description", "")),
)
def _spec_from_dict(raw: dict[str, object]) -> DatasetSpec:
missing = [key for key in ("id", "layer", "license_tier", "source") if key not in raw]
if missing:
raise ValueError(f"dataset entry missing required keys: {missing} in {raw!r}")
source = raw["source"]
if not isinstance(source, dict):
raise ValueError(f"[{raw.get('id')}] source must be a mapping")
languages = raw.get("languages", [])
if not isinstance(languages, list):
raise ValueError(f"[{raw.get('id')}] languages must be a list")
return DatasetSpec(
id=str(raw["id"]),
title=str(raw.get("title", raw["id"])),
layer=str(raw["layer"]),
languages=tuple(str(lang) for lang in languages),
license=str(raw.get("license", "unknown")),
license_tier=str(raw["license_tier"]),
source=dict(source),
homepage=str(raw.get("homepage", "")),
notes=str(raw.get("notes", "")),
optional=bool(raw.get("optional", False)),
)
# Backwards-friendly export list.
__all__ = [
"DatasetSpec",
"Registry",
"load_registry",
"LICENSE_TIERS",
"SOURCE_KINDS",
"TIER_OPEN",
"TIER_GATED_AUTO",
"TIER_GATED_MANUAL",
"TIER_LICENSE_REQUIRED",
"KIND_HF_SNAPSHOT",
"KIND_HF_FILES",
"KIND_HTTP",
"KIND_GITHUB_ARCHIVE",
"KIND_LICENSE_REQUIRED",
"LANG_MULTI",
]
|