Spaces:
Sleeping
Sleeping
Download src/tsfm_bench/data/ts_bench.py from ThinkcatLab/LiveHouse-TS: direct link, hf CLI and curl.
- Browser
- Download file 30.8 kB
-
https://huggingface.co/spaces/ThinkcatLab/LiveHouse-TS/resolve/main/src/tsfm_bench/data/ts_bench.py
- Command line
-
hf download hf://spaces/ThinkcatLab/LiveHouse-TS/src/tsfm_bench/data/ts_bench.py
-
curl -L -o ts_bench.py https://huggingface.co/spaces/ThinkcatLab/LiveHouse-TS/resolve/main/src/tsfm_bench/data/ts_bench.py
30.8 kB
| from __future__ import annotations | |
| import csv | |
| import json | |
| import logging | |
| import subprocess | |
| import sys | |
| import time | |
| from dataclasses import dataclass | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| from typing import Any, Iterator | |
| import numpy as np | |
| import pandas as pd | |
| from tsfm_bench.data.interface import DatasetMetadata, TimeSeriesRecord | |
| logger = logging.getLogger(__name__) | |
| def preserve_missing_datasets(normalized_dir: Path, backup_dir: Path) -> list[str]: | |
| """Keep the last observations/tasks of sources absent from a partial refresh. | |
| Original timestamps remain unchanged, so stale contexts are still rejected by | |
| the prequential gate. This does not present cached data as freshly collected. | |
| """ | |
| def read_rows(path): | |
| if not path.exists(): | |
| return [], [] | |
| with path.open(newline="", encoding="utf-8") as handle: | |
| reader = csv.DictReader(handle) | |
| return list(reader.fieldnames or []), list(reader) | |
| _, old = read_rows(backup_dir / "observations.csv") | |
| _, new = read_rows(normalized_dir / "observations.csv") | |
| missing = {row["dataset_id"] for row in old} - {row["dataset_id"] for row in new} | |
| if not missing: | |
| return [] | |
| for name in ("observations", "entities", "variables", "forecast_tasks", "quality_reports", "raw_snapshots"): | |
| fields, current = read_rows(normalized_dir / f"{name}.csv") | |
| old_fields, cached = read_rows(backup_dir / f"{name}.csv") | |
| fields = fields or old_fields | |
| retained = [row for row in cached if row.get("dataset_id") in missing] | |
| if not retained: | |
| continue | |
| current = [row for row in current if row.get("dataset_id") not in missing] | |
| path = normalized_dir / f"{name}.csv" | |
| temporary = path.with_suffix(".tmp") | |
| with temporary.open("w", newline="", encoding="utf-8") as handle: | |
| writer = csv.DictWriter(handle, fieldnames=fields) | |
| writer.writeheader() | |
| writer.writerows([*current, *retained]) | |
| temporary.replace(path) | |
| logger.warning("Retained cached datasets absent from refresh: %s", sorted(missing)) | |
| return sorted(missing) | |
| DEFAULT_SOURCE_DOMAINS: dict[str, str] = { | |
| "open_meteo_weather": "Weather", | |
| "open_meteo_air_quality": "Air Quality", | |
| "nasa_power": "Weather", | |
| "usgs_water": "Hydrology", | |
| "noaa_coops": "Ocean", | |
| "noaa_ndbc": "Ocean", | |
| "gbfs_citibike": "Transport", | |
| "binance": "Finance", | |
| "coingecko": "Finance", | |
| "wikimedia_pageviews": "Web", | |
| "nws_observations": "Weather", | |
| "usgs_earthquake": "Hazards", | |
| "noaa_ncei": "Weather", | |
| "data_gov_sg_rainfall": "Climate", | |
| "data_gov_sg_taxi_availability": "Transportation", | |
| "data_gov_sg_sgx_turnover": "Finance", | |
| } | |
| # Explicit dataset_id → domain mapping derived from all known TS-Bench datasets. | |
| # This is the primary lookup and never relies on entities.csv being present. | |
| DATASET_ID_DOMAINS: dict[str, str] = { | |
| "binance_btcusdt_1h": "Finance", | |
| "coingecko_bitcoin_7d": "Finance", | |
| "gbfs_citibike_current": "Transport", | |
| "nasa_power_shanghai_hourly": "Weather", | |
| "noaa_coops_sf_water_level": "Ocean", | |
| "noaa_ndbc_46013_realtime": "Ocean", | |
| "nws_ksfo_observations": "Weather", | |
| "open_meteo_shanghai_air_quality_hourly": "Air Quality", | |
| "open_meteo_shanghai_hourly": "Weather", | |
| "usgs_earthquake_all_week_hourly": "Hazards", | |
| "usgs_potomac_iv": "Hydrology", | |
| "wikimedia_time_series_hourly": "Web", | |
| "data_gov_sg_singapore_rainfall_5min": "Climate", | |
| "data_gov_sg_singapore_taxi_availability": "Transportation", | |
| "data_gov_sg_sgx_equities_turnover_monthly": "Finance", | |
| } | |
| def resolve_domain( | |
| dataset_id: str, | |
| source_id: str, | |
| source_domains: dict[str, str], | |
| ) -> str: | |
| """Determine the leaderboard domain for a dataset. | |
| Resolution order (most-specific first): | |
| 1. Exact match in DATASET_ID_DOMAINS. | |
| 2. Longest-prefix match in DATASET_ID_DOMAINS. | |
| 3. Exact match of source_id in source_domains. | |
| 4. Longest-prefix match of dataset_id against source_domains keys. | |
| 5. Fallback to "Other". | |
| """ | |
| if dataset_id in DATASET_ID_DOMAINS: | |
| return DATASET_ID_DOMAINS[dataset_id] | |
| # Longest prefix match against DATASET_ID_DOMAINS keys. | |
| best = max( | |
| (k for k in DATASET_ID_DOMAINS if dataset_id.startswith(k)), | |
| key=len, | |
| default=None, | |
| ) | |
| if best: | |
| return DATASET_ID_DOMAINS[best] | |
| if source_id in source_domains: | |
| return source_domains[source_id] | |
| # Longest prefix match against source_domains keys. | |
| best_src = max( | |
| (k for k in source_domains if dataset_id.startswith(k)), | |
| key=len, | |
| default=None, | |
| ) | |
| if best_src: | |
| return source_domains[best_src] | |
| return "Other" | |
| FREQ_ALIASES = { | |
| "1h": "1h", | |
| "6min": "6min", | |
| "10min": "10min", | |
| "15min": "15min", | |
| "5min": "5min", | |
| "30s": "30s", | |
| "1mo": "M", | |
| "1d": "1D", | |
| "snapshot": "1h", | |
| } | |
| def _parse_duration(value: str | int | float | None) -> timedelta | None: | |
| if value in (None, ""): | |
| return None | |
| if isinstance(value, (int, float)): | |
| return timedelta(seconds=float(value)) | |
| text = str(value).strip().lower() | |
| units = { | |
| "s": "seconds", | |
| "sec": "seconds", | |
| "secs": "seconds", | |
| "second": "seconds", | |
| "seconds": "seconds", | |
| "m": "minutes", | |
| "min": "minutes", | |
| "mins": "minutes", | |
| "minute": "minutes", | |
| "minutes": "minutes", | |
| "h": "hours", | |
| "hr": "hours", | |
| "hrs": "hours", | |
| "hour": "hours", | |
| "hours": "hours", | |
| "d": "days", | |
| "day": "days", | |
| "days": "days", | |
| "w": "weeks", | |
| "week": "weeks", | |
| "weeks": "weeks", | |
| } | |
| import re | |
| match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*([a-z]+)\s*", text) | |
| if not match: | |
| raise ValueError(f"Invalid duration: {value!r}") | |
| amount = float(match.group(1)) | |
| unit = units.get(match.group(2)) | |
| if unit is None: | |
| raise ValueError(f"Unsupported duration unit in {value!r}") | |
| return timedelta(**{unit: amount}) | |
| def _duration_to_points(value: str | int | float | None, freq: str) -> int | None: | |
| duration = _parse_duration(value) | |
| if duration is None: | |
| return None | |
| offset = pd.tseries.frequencies.to_offset(FREQ_ALIASES.get(freq, freq)) | |
| delta = pd.Timedelta(offset.nanos, unit="ns") | |
| points = int(duration / delta) | |
| return max(1, points) | |
| class _TaskWindowPolicy: | |
| context_window: str | None = None | |
| prediction_window: str | None = None | |
| context_length_steps: int | None = None | |
| prediction_length_steps: int | None = None | |
| min_context_points: int = 32 | |
| max_context_points: int | None = None | |
| max_prediction_points: int | None = None | |
| def from_dict(cls, payload: dict[str, Any] | None) -> "_TaskWindowPolicy": | |
| payload = payload or {} | |
| max_context = payload.get("max_context_points") | |
| max_prediction = payload.get("max_prediction_points") | |
| context_steps = payload.get("context_length_steps", payload.get("history_length_steps")) | |
| prediction_steps = payload.get("prediction_length_steps", payload.get("forecast_horizon_steps")) | |
| return cls( | |
| context_window=payload.get("context_window"), | |
| prediction_window=payload.get("prediction_window"), | |
| context_length_steps=int(context_steps) if context_steps is not None else None, | |
| prediction_length_steps=int(prediction_steps) if prediction_steps is not None else None, | |
| min_context_points=int(payload.get("min_context_points", 32)), | |
| max_context_points=int(max_context) if max_context is not None else None, | |
| max_prediction_points=int(max_prediction) if max_prediction is not None else None, | |
| ) | |
| def context_points(self, freq: str, available: int) -> int: | |
| desired = self.context_length_steps or _duration_to_points(self.context_window, freq) or available | |
| if self.max_context_points is not None: | |
| desired = min(desired, self.max_context_points) | |
| return max(1, min(int(desired), available)) | |
| def prediction_points(self, freq: str, available: int) -> int: | |
| desired = self.prediction_length_steps or _duration_to_points(self.prediction_window, freq) or available | |
| if self.max_prediction_points is not None: | |
| desired = min(desired, self.max_prediction_points) | |
| return max(1, min(int(desired), available)) | |
| def _resolve_dataset_registry_path(raw_settings: dict[str, Any], payload: dict[str, Any]) -> Path | None: | |
| raw_path = raw_settings.get("dataset_registry_path") or payload.get("dataset_registry_path") | |
| if not raw_path: | |
| return None | |
| path = Path(raw_path).expanduser() | |
| if path.is_absolute() or path.exists(): | |
| return path.resolve() | |
| config_dir = payload.get("__config_dir") | |
| if config_dir: | |
| candidate = Path(config_dir) / path | |
| if candidate.exists(): | |
| return candidate.resolve() | |
| return path.resolve() | |
| def _load_dataset_registry_task_windows(path: Path | None) -> dict[str, _TaskWindowPolicy]: | |
| if path is None: | |
| return {} | |
| if not path.exists(): | |
| raise FileNotFoundError(f"Dataset registry not found: {path}") | |
| registry = json.loads(path.read_text(encoding="utf-8")) | |
| windows: dict[str, _TaskWindowPolicy] = {} | |
| for entry in registry.get("datasets", []): | |
| dataset_id = str(entry.get("dataset_id", "")) | |
| if not dataset_id: | |
| continue | |
| history_steps = entry.get("history_length_steps") | |
| horizon_steps = entry.get("forecast_horizon_steps") | |
| if history_steps is None and horizon_steps is None: | |
| continue | |
| windows[dataset_id] = _TaskWindowPolicy.from_dict( | |
| { | |
| "context_length_steps": history_steps, | |
| "prediction_length_steps": horizon_steps, | |
| "min_context_points": min(int(history_steps), 32) if history_steps is not None else 32, | |
| } | |
| ) | |
| return windows | |
| class _TsBenchSettings: | |
| ts_bench_root: Path | |
| dataset_registry_path: Path | None = None | |
| output_subdir: str = "data/sample_week" | |
| collect_days: int = 7 | |
| collect_timeout: int = 45 | |
| limit_entities: int = 10 | |
| task_limit: int = 20 | |
| max_tasks_per_dataset: int = 1 | |
| min_future_steps: int = 1 | |
| skip_frequencies: tuple[str, ...] = ("snapshot",) | |
| auto_refresh: bool = True | |
| collect_retries: int = 3 | |
| collect_retry_delay: int = 15 | |
| source_intervals: dict[str, int] = None # type: ignore[assignment] | |
| default_task_window: _TaskWindowPolicy = None # type: ignore[assignment] | |
| source_task_windows: dict[str, _TaskWindowPolicy] = None # type: ignore[assignment] | |
| dataset_task_windows: dict[str, _TaskWindowPolicy] = None # type: ignore[assignment] | |
| def __post_init__(self): | |
| if self.source_intervals is None: | |
| object.__setattr__(self, "source_intervals", {}) | |
| if self.default_task_window is None: | |
| object.__setattr__(self, "default_task_window", _TaskWindowPolicy()) | |
| if self.source_task_windows is None: | |
| object.__setattr__(self, "source_task_windows", {}) | |
| if self.dataset_task_windows is None: | |
| object.__setattr__(self, "dataset_task_windows", {}) | |
| class _TaskSpec: | |
| task_id: str | |
| dataset_id: str | |
| source_id: str | |
| domain: str | |
| target_variable_id: str | |
| frequency: str | |
| gluonts_freq: str | |
| prediction_length: int | |
| context_length: int | None | |
| task_dir: Path | |
| leaderboard_name: str | |
| class TsBenchDataSource: | |
| """Live TS-Bench data: refresh via collect_weekly_data.py, read exported tasks.""" | |
| def __init__( | |
| self, | |
| settings: _TsBenchSettings, | |
| source_domains: dict[str, str], | |
| tasks: tuple[_TaskSpec, ...], | |
| fetched_at: str = "", | |
| ): | |
| self._settings = settings | |
| self._source_domains = source_domains | |
| self._tasks = {task.task_id: task for task in tasks} | |
| self._fetched_at = fetched_at | |
| def from_config(cls, payload: dict[str, Any]) -> TsBenchDataSource: | |
| raw = payload.get("settings", {}) | |
| ts_bench_root = Path(raw.get("ts_bench_root", "../TS-Bench")).expanduser().resolve() | |
| raw_intervals = payload.get("source_intervals", {}) | |
| source_intervals = {k: int(v) for k, v in raw_intervals.items()} | |
| raw_windows = payload.get("task_windows", {}) | |
| default_task_window = _TaskWindowPolicy.from_dict(raw_windows.get("default", {})) | |
| source_task_windows = { | |
| str(k): _TaskWindowPolicy.from_dict(v) | |
| for k, v in raw_windows.get("sources", {}).items() | |
| } | |
| registry_path = _resolve_dataset_registry_path(raw, payload) | |
| registry_task_windows = _load_dataset_registry_task_windows(registry_path) | |
| dataset_task_windows = { | |
| str(k): _TaskWindowPolicy.from_dict(v) | |
| for k, v in raw_windows.get("datasets", {}).items() | |
| } | |
| dataset_task_windows = {**registry_task_windows, **dataset_task_windows} | |
| settings = _TsBenchSettings( | |
| ts_bench_root=ts_bench_root, | |
| dataset_registry_path=registry_path, | |
| output_subdir=str(raw.get("output_subdir", "data/sample_week")), | |
| collect_days=int(raw.get("collect_days", 7)), | |
| collect_timeout=int(raw.get("collect_timeout", 45)), | |
| limit_entities=int(raw.get("limit_entities", 10)), | |
| task_limit=int(raw.get("task_limit", 20)), | |
| max_tasks_per_dataset=int(raw.get("max_tasks_per_dataset", 1)), | |
| min_future_steps=int(raw.get("min_future_steps", 1)), | |
| skip_frequencies=tuple(raw.get("skip_frequencies", ["snapshot"])), | |
| auto_refresh=bool(raw.get("auto_refresh", True)), | |
| collect_retries=int(raw.get("collect_retries", 3)), | |
| collect_retry_delay=int(raw.get("collect_retry_delay", 15)), | |
| source_intervals=source_intervals, | |
| default_task_window=default_task_window, | |
| source_task_windows=source_task_windows, | |
| dataset_task_windows=dataset_task_windows, | |
| ) | |
| source_domains = {**DEFAULT_SOURCE_DOMAINS, **payload.get("source_domains", {})} | |
| source = cls(settings=settings, source_domains=source_domains, tasks=()) | |
| source._load_tasks() | |
| return source | |
| def data_root(self) -> Path: | |
| return self._settings.ts_bench_root / self._settings.output_subdir | |
| def _task_window_policy(self, dataset_id: str, source_id: str) -> _TaskWindowPolicy: | |
| if dataset_id in self._settings.dataset_task_windows: | |
| return self._settings.dataset_task_windows[dataset_id] | |
| best_dataset = max( | |
| (k for k in self._settings.dataset_task_windows if dataset_id.startswith(k)), | |
| key=len, | |
| default=None, | |
| ) | |
| if best_dataset: | |
| return self._settings.dataset_task_windows[best_dataset] | |
| if source_id in self._settings.source_task_windows: | |
| return self._settings.source_task_windows[source_id] | |
| best_source = max( | |
| (k for k in self._settings.source_task_windows if dataset_id.startswith(k)), | |
| key=len, | |
| default=None, | |
| ) | |
| if best_source: | |
| return self._settings.source_task_windows[best_source] | |
| return self._settings.default_task_window | |
| def _task_policy_payload(self) -> dict[str, Any]: | |
| def to_dict(policy: _TaskWindowPolicy) -> dict[str, Any]: | |
| payload: dict[str, Any] = { | |
| "min_context_points": policy.min_context_points, | |
| } | |
| if policy.context_window: | |
| payload["context_window"] = policy.context_window | |
| if policy.prediction_window: | |
| payload["prediction_window"] = policy.prediction_window | |
| if policy.context_length_steps is not None: | |
| payload["context_length_steps"] = policy.context_length_steps | |
| if policy.prediction_length_steps is not None: | |
| payload["prediction_length_steps"] = policy.prediction_length_steps | |
| if policy.max_context_points is not None: | |
| payload["max_context_points"] = policy.max_context_points | |
| if policy.max_prediction_points is not None: | |
| payload["max_prediction_points"] = policy.max_prediction_points | |
| return payload | |
| return { | |
| "default": to_dict(self._settings.default_task_window), | |
| "sources": { | |
| key: to_dict(value) | |
| for key, value in sorted(self._settings.source_task_windows.items()) | |
| }, | |
| "datasets": { | |
| key: to_dict(value) | |
| for key, value in sorted(self._settings.dataset_task_windows.items()) | |
| }, | |
| } | |
| def refresh(self) -> None: | |
| import shutil | |
| root = self._settings.ts_bench_root | |
| output_dir = self.data_root | |
| collect_script = root / "collect_weekly_data.py" | |
| export_script = root / "export_model_inputs.py" | |
| if not collect_script.exists(): | |
| raise FileNotFoundError(f"TS-Bench collect script not found: {collect_script}") | |
| # Create backups of existing data if they exist to prevent partial or empty runs from breaking the cache | |
| normalized_dir = self.data_root / "normalized" | |
| model_inputs_dir = self.data_root / "model_inputs_latest" | |
| backup_normalized = self.data_root / "normalized.bak" | |
| backup_model_inputs = self.data_root / "model_inputs_latest.bak" | |
| if normalized_dir.exists(): | |
| shutil.rmtree(backup_normalized, ignore_errors=True) | |
| shutil.copytree(normalized_dir, backup_normalized) | |
| if model_inputs_dir.exists(): | |
| shutil.rmtree(backup_model_inputs, ignore_errors=True) | |
| shutil.copytree(model_inputs_dir, backup_model_inputs) | |
| collect_cmd = [ | |
| sys.executable, | |
| str(collect_script), | |
| "--days", | |
| str(self._settings.collect_days), | |
| "--output", | |
| self._settings.output_subdir, | |
| "--timeout", | |
| str(self._settings.collect_timeout), | |
| "--limit-entities", | |
| str(self._settings.limit_entities), | |
| "--max-tasks-per-dataset", | |
| str(self._settings.max_tasks_per_dataset), | |
| "--task-window-policy", | |
| json.dumps(self._task_policy_payload()), | |
| ] | |
| retries = max(1, self._settings.collect_retries) | |
| delay = self._settings.collect_retry_delay | |
| collect_ok = False | |
| for attempt in range(1, retries + 1): | |
| logger.info( | |
| "Collecting TS-Bench data → %s (attempt %d/%d)", | |
| output_dir, | |
| attempt, | |
| retries, | |
| ) | |
| result = subprocess.run(collect_cmd, cwd=root, check=False) | |
| if result.returncode == 0: | |
| collect_ok = True | |
| break | |
| if attempt < retries: | |
| wait = delay * attempt | |
| logger.warning( | |
| "collect_weekly_data.py failed (exit %d); retrying in %ds", | |
| result.returncode, | |
| wait, | |
| ) | |
| time.sleep(wait) | |
| else: | |
| logger.warning( | |
| "collect_weekly_data.py failed after %d attempts (exit %d); " | |
| "will use existing cached data if available", | |
| retries, | |
| result.returncode, | |
| ) | |
| export_ok = False | |
| if collect_ok: | |
| preserve_missing_datasets(normalized_dir, backup_normalized) | |
| normalized_rel = self._settings.output_subdir + "/normalized" | |
| model_inputs_rel = self._settings.output_subdir + "/model_inputs_latest" | |
| if export_script.exists(): | |
| logger.info("Exporting TS-Bench model inputs → %s", self.data_root / "model_inputs_latest") | |
| export_cmd = [ | |
| sys.executable, | |
| str(export_script), | |
| "--input", | |
| normalized_rel, | |
| "--output", | |
| model_inputs_rel, | |
| "--limit", | |
| str(self._settings.task_limit), | |
| ] | |
| export_result = subprocess.run(export_cmd, cwd=root, check=False) | |
| if export_result.returncode == 0: | |
| export_ok = True | |
| else: | |
| logger.warning( | |
| "export_model_inputs.py failed (exit %d); using last exported manifest", | |
| export_result.returncode, | |
| ) | |
| # Verify that we ended up with a valid non-empty task manifest | |
| has_tasks = False | |
| manifest_path = model_inputs_dir / "manifest.json" | |
| if not manifest_path.exists(): | |
| manifest_path = model_inputs_dir / "manifest.csv" | |
| if manifest_path.exists(): | |
| try: | |
| if manifest_path.suffix == ".json": | |
| entries = json.loads(manifest_path.read_text()) | |
| else: | |
| with manifest_path.open(newline="", encoding="utf-8") as handle: | |
| entries = list(csv.DictReader(handle)) | |
| if len(entries) > 0: | |
| has_tasks = True | |
| except Exception as e: | |
| logger.warning("Failed to parse manifest to verify tasks: %s", e) | |
| if not collect_ok or not export_ok or not has_tasks: | |
| logger.warning( | |
| "Realtime collection or export failed/returned 0 tasks (collect_ok=%s, export_ok=%s, has_tasks=%s). " | |
| "Restoring previous cached data from backup to prevent empty leaderboard.", | |
| collect_ok, export_ok, has_tasks | |
| ) | |
| if backup_normalized.exists(): | |
| shutil.rmtree(normalized_dir, ignore_errors=True) | |
| shutil.copytree(backup_normalized, normalized_dir) | |
| if backup_model_inputs.exists(): | |
| shutil.rmtree(model_inputs_dir, ignore_errors=True) | |
| shutil.copytree(backup_model_inputs, model_inputs_dir) | |
| # Clean up backups | |
| shutil.rmtree(backup_normalized, ignore_errors=True) | |
| shutil.rmtree(backup_model_inputs, ignore_errors=True) | |
| self._load_tasks() | |
| def _load_tasks(self) -> None: | |
| summary_path = self.data_root / "summary.json" | |
| if summary_path.exists(): | |
| summary = json.loads(summary_path.read_text()) | |
| self._fetched_at = summary.get("generated_at", "") | |
| else: | |
| self._fetched_at = datetime.now(timezone.utc).isoformat() | |
| source_by_dataset = self._load_dataset_sources() | |
| tasks = self._load_task_specs(source_by_dataset) | |
| self._tasks = {task.task_id: task for task in tasks} | |
| logger.info("Loaded %d TS-Bench tasks from %s", len(tasks), self.data_root) | |
| def _load_dataset_sources(self) -> dict[str, str]: | |
| entities_path = self.data_root / "normalized" / "entities.csv" | |
| mapping: dict[str, str] = {} | |
| if not entities_path.exists(): | |
| return mapping | |
| with entities_path.open(newline="", encoding="utf-8") as handle: | |
| for row in csv.DictReader(handle): | |
| mapping.setdefault(row["dataset_id"], row["source_id"]) | |
| return mapping | |
| def _load_task_specs(self, source_by_dataset: dict[str, str]) -> tuple[_TaskSpec, ...]: | |
| manifest_path = self.data_root / "model_inputs_latest" / "manifest.json" | |
| if not manifest_path.exists(): | |
| manifest_path = self.data_root / "model_inputs_latest" / "manifest.csv" | |
| if not manifest_path.exists(): | |
| raise FileNotFoundError( | |
| f"No TS-Bench manifest under {self.data_root / 'model_inputs_latest'}. " | |
| "Run collect_weekly_data.py and export_model_inputs.py first." | |
| ) | |
| entries: list[dict[str, Any]] | |
| if manifest_path.suffix == ".json": | |
| entries = json.loads(manifest_path.read_text()) | |
| else: | |
| with manifest_path.open(newline="", encoding="utf-8") as handle: | |
| entries = list(csv.DictReader(handle)) | |
| specs: list[_TaskSpec] = [] | |
| for entry in entries: | |
| future_rows = int(entry.get("future_target_rows", 0)) | |
| if future_rows < self._settings.min_future_steps: | |
| continue | |
| task_id = entry["task_id"] | |
| task_dir = self.data_root / "model_inputs_latest" / task_id | |
| task_json = task_dir / "task.json" | |
| if not task_json.exists(): | |
| continue | |
| task = json.loads(task_json.read_text()) | |
| raw_freq = str(task.get("frequency", "1h")) | |
| if raw_freq in self._settings.skip_frequencies: | |
| continue | |
| dataset_id = task["dataset_id"] | |
| source_id = source_by_dataset.get(dataset_id, dataset_id.split("_")[0]) | |
| domain = resolve_domain(dataset_id, source_id, self._source_domains) | |
| gluonts_freq = FREQ_ALIASES.get(raw_freq, raw_freq) | |
| policy = self._task_window_policy(dataset_id, source_id) | |
| pred_len = policy.prediction_points(raw_freq, future_rows) | |
| raw_context_rows = int(entry.get("context_rows", task.get("input_length", 0)) or 0) | |
| context_len = ( | |
| policy.context_points(raw_freq, raw_context_rows) | |
| if raw_context_rows > 0 | |
| else None | |
| ) | |
| variable = task["target_variable_id"] | |
| leaderboard_name = f"{dataset_id}_{variable}" | |
| specs.append( | |
| _TaskSpec( | |
| task_id=task_id, | |
| dataset_id=dataset_id, | |
| source_id=source_id, | |
| domain=domain, | |
| target_variable_id=variable, | |
| frequency=raw_freq, | |
| gluonts_freq=gluonts_freq, | |
| prediction_length=pred_len, | |
| context_length=context_len, | |
| task_dir=task_dir, | |
| leaderboard_name=leaderboard_name, | |
| ) | |
| ) | |
| return tuple(specs) | |
| def list_datasets(self) -> list[str]: | |
| return sorted(self._tasks.keys()) | |
| def get_metadata(self, name: str) -> DatasetMetadata: | |
| task = self._require_task(name) | |
| return DatasetMetadata( | |
| name=name, | |
| domain=task.domain, | |
| frequency=task.gluonts_freq, | |
| num_variates=1, | |
| terms=("short",), | |
| source_type="ts_bench", | |
| ) | |
| def get_prediction_length(self, name: str) -> int: | |
| return self._require_task(name).prediction_length | |
| def stream(self, name: str) -> Iterator[TimeSeriesRecord]: | |
| task = self._require_task(name) | |
| context_path = task.task_dir / "context.csv" | |
| future_path = task.task_dir / "future_target.csv" | |
| if not context_path.exists() or not future_path.exists(): | |
| raise FileNotFoundError(f"Missing task CSVs for {name}") | |
| context = pd.read_csv(context_path) | |
| future = pd.read_csv(future_path) | |
| target_col = task.target_variable_id | |
| if target_col not in context.columns: | |
| raise RuntimeError(f"Target column {target_col} missing in {context_path}") | |
| context_times = list(context["timestamp"]) | |
| future_times = list(future["timestamp"]) | |
| context_vals = pd.to_numeric(context[target_col], errors="coerce").to_numpy(dtype=np.float32) | |
| future_vals = pd.to_numeric(future[target_col], errors="coerce").to_numpy(dtype=np.float32) | |
| if task.context_length is not None and len(context_vals) > task.context_length: | |
| context_vals = context_vals[-task.context_length:] | |
| context_times = context_times[-task.context_length:] | |
| if len(future_vals) > task.prediction_length: | |
| future_vals = future_vals[:task.prediction_length] | |
| future_times = future_times[:task.prediction_length] | |
| values = np.concatenate([context_vals, future_vals]) | |
| mask = ~np.isnan(values) | |
| values = values[mask] | |
| if len(values) < task.prediction_length + 2: | |
| raise RuntimeError( | |
| f"Insufficient data for {name}: {len(values)} points, " | |
| f"need > {task.prediction_length + 2}" | |
| ) | |
| timestamps = context_times + future_times | |
| timestamps = [timestamps[i] for i, keep in enumerate(mask) if keep] | |
| start_ts = pd.Timestamp(timestamps[0]) | |
| period_freq = ( | |
| "M" | |
| if task.gluonts_freq == "M" | |
| else pd.tseries.frequencies.to_offset(task.gluonts_freq) | |
| ) | |
| yield TimeSeriesRecord( | |
| item_id=task.leaderboard_name.replace("/", "_"), | |
| target=values, | |
| start=str(pd.Period(start_ts, freq=period_freq)), | |
| freq=task.gluonts_freq, | |
| ) | |
| def leaderboard_dataset_name(self, name: str) -> str: | |
| task = self._require_task(name) | |
| display_frequency = "1mo" if task.gluonts_freq == "M" else task.gluonts_freq | |
| return f"{task.leaderboard_name}/{display_frequency}/short" | |
| def fetched_at(self) -> str: | |
| return self._fetched_at or datetime.now(timezone.utc).isoformat() | |
| def get_eval_interval(self, name: str, default: int = 300) -> int: | |
| """Return the eval interval in seconds for a dataset task. | |
| Looks up source_intervals by source_id with longest-prefix fallback. | |
| Returns *default* when no match is configured. | |
| """ | |
| task = self._require_task(name) | |
| intervals = self._settings.source_intervals | |
| if not intervals: | |
| return default | |
| # Exact match on source_id | |
| if task.source_id in intervals: | |
| return intervals[task.source_id] | |
| # Longest-prefix match on dataset_id against interval keys | |
| best = max( | |
| (k for k in intervals if task.dataset_id.startswith(k)), | |
| key=len, | |
| default=None, | |
| ) | |
| if best: | |
| return intervals[best] | |
| return default | |
| def _require_task(self, name: str) -> _TaskSpec: | |
| if name not in self._tasks: | |
| raise KeyError(f"Unknown TS-Bench task: {name}") | |
| return self._tasks[name] | |