Spaces:
Running
Running
File size: 2,873 Bytes
e317359 | 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 | from __future__ import annotations
import json
from pathlib import Path
from typing import Any
import yaml
from tsfm_bench.data.interface import TimeSeriesDataSource
from tsfm_bench.data.open_meteo import OpenMeteoDataSource
from tsfm_bench.data.synthetic import DEFAULT_SYNTHETIC_SPECS, SyntheticDataSource, _SyntheticSpec
from tsfm_bench.data.ts_bench import TsBenchDataSource
def load_data_source(config_path: str | Path | None = None) -> TimeSeriesDataSource:
"""Load a data source from config. Defaults to built-in synthetic demo."""
if config_path is None:
return SyntheticDataSource()
path = Path(config_path)
if not path.exists():
raise FileNotFoundError(f"Dataset config not found: {path}")
if path.suffix in {".yaml", ".yml"}:
payload = yaml.safe_load(path.read_text())
elif path.suffix == ".json":
payload = json.loads(path.read_text())
else:
raise ValueError(f"Unsupported config format: {path.suffix}")
payload["__config_dir"] = str(path.resolve().parent)
source_type = payload.get("source_type", "synthetic")
if source_type == "synthetic":
specs = tuple(_parse_spec(entry) for entry in payload.get("datasets", []))
return SyntheticDataSource(specs=specs or DEFAULT_SYNTHETIC_SPECS)
if source_type == "open_meteo":
return OpenMeteoDataSource.from_config(payload)
if source_type == "ts_bench":
return TsBenchDataSource.from_config(payload)
raise NotImplementedError(
f"Source type '{source_type}' is not implemented yet. "
"Implement TimeSeriesDataSource and register it in registry.py."
)
def load_dataset_properties(config_path: str | Path | None = None) -> dict[str, dict[str, Any]]:
"""Build dataset_properties map keyed by dataset stem (before frequency suffix)."""
source = load_data_source(config_path)
properties: dict[str, dict[str, Any]] = {}
for name in source.list_datasets():
meta = source.get_metadata(name)
if isinstance(source, TsBenchDataSource):
ds_key = source.leaderboard_dataset_name(name).split("/")[0].lower()
else:
ds_key = name.split("/")[0].lower()
properties[ds_key] = {
"domain": meta.domain,
"frequency": meta.frequency,
"num_variates": meta.num_variates,
"terms": list(meta.terms),
}
return properties
def _parse_spec(entry: dict[str, Any]) -> _SyntheticSpec:
return _SyntheticSpec(
name=entry["name"],
domain=entry["domain"],
frequency=entry["frequency"],
num_variates=int(entry["num_variates"]),
terms=tuple(entry.get("terms", ["short"])),
n_series=int(entry.get("n_series", 30)),
series_length=int(entry.get("series_length", 500)),
seed=int(entry.get("seed", 0)),
)
|