Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Export task-level model input CSV files from normalized TS-Bench tables. | |
| This script does not fetch data. It reads: | |
| normalized/observations.csv | |
| normalized/forecast_tasks.csv | |
| For each task it writes: | |
| model_inputs/<task_id>/context.csv | |
| model_inputs/<task_id>/future_target.csv | |
| The context file is what a model can consume. The future_target file is only | |
| for evaluator/scoring code. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import json | |
| import sys | |
| from collections import defaultdict | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any, Dict, Iterable, List, Tuple | |
| UTC = timezone.utc | |
| def parse_dt(value: str) -> datetime: | |
| text = value.strip() | |
| if text.endswith("Z"): | |
| text = text[:-1] + "+00:00" | |
| dt = datetime.fromisoformat(text) | |
| if dt.tzinfo is None: | |
| dt = dt.replace(tzinfo=UTC) | |
| return dt.astimezone(UTC) | |
| def read_csv(path: Path) -> List[Dict[str, str]]: | |
| with path.open("r", newline="", encoding="utf-8") as handle: | |
| return list(csv.DictReader(handle)) | |
| def write_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: List[str]) -> int: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| count = 0 | |
| with path.open("w", newline="", encoding="utf-8") as handle: | |
| writer = csv.DictWriter(handle, fieldnames=fieldnames) | |
| writer.writeheader() | |
| for row in rows: | |
| writer.writerow({name: row.get(name, "") for name in fieldnames}) | |
| count += 1 | |
| return count | |
| def pivot_rows(rows: List[Dict[str, str]]) -> Tuple[List[Dict[str, Any]], List[str]]: | |
| grouped: Dict[str, Dict[str, Any]] = {} | |
| variables = sorted({row["variable_id"] for row in rows}) | |
| for row in rows: | |
| timestamp = row["timestamp"] | |
| if timestamp not in grouped: | |
| grouped[timestamp] = {"timestamp": timestamp} | |
| grouped[timestamp][row["variable_id"]] = row["value"] | |
| return [grouped[key] for key in sorted(grouped)], ["timestamp"] + variables | |
| def export_task( | |
| task: Dict[str, str], | |
| observations: List[Dict[str, str]], | |
| output_root: Path, | |
| ) -> Dict[str, Any]: | |
| dataset_id = task["dataset_id"] | |
| entity_id = task["target_entity_id"] | |
| target_variable_id = task["target_variable_id"] | |
| forecast_issue_time = parse_dt(task["forecast_issue_time"]) | |
| context_start = parse_dt(task["context_start"]) | |
| context_end = parse_dt(task["context_end"]) | |
| horizon_start = parse_dt(task["horizon_start"]) | |
| horizon_end = parse_dt(task["horizon_end"]) | |
| context_rows = [] | |
| future_rows = [] | |
| for row in observations: | |
| if row["dataset_id"] != dataset_id or row["entity_id"] != entity_id: | |
| continue | |
| timestamp = parse_dt(row["timestamp"]) | |
| available_time = parse_dt(row["available_time"]) | |
| if context_start <= timestamp <= context_end and available_time <= forecast_issue_time: | |
| context_rows.append(row) | |
| if horizon_start <= timestamp <= horizon_end and row["variable_id"] == target_variable_id: | |
| future_rows.append(row) | |
| context_pivot, context_fields = pivot_rows(context_rows) | |
| future_pivot, future_fields = pivot_rows(future_rows) | |
| task_dir = output_root / task["task_id"] | |
| context_count = write_csv(task_dir / "context.csv", context_pivot, context_fields) | |
| future_count = write_csv(task_dir / "future_target.csv", future_pivot, future_fields) | |
| (task_dir / "task.json").write_text(json.dumps(task, ensure_ascii=False, indent=2), encoding="utf-8") | |
| return { | |
| "task_id": task["task_id"], | |
| "dataset_id": dataset_id, | |
| "entity_id": entity_id, | |
| "target_variable_id": target_variable_id, | |
| "context_rows": context_count, | |
| "future_target_rows": future_count, | |
| "context_csv": str(task_dir / "context.csv"), | |
| "future_target_csv": str(task_dir / "future_target.csv"), | |
| } | |
| def parse_args(argv: List[str]) -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Export model input CSVs from normalized TS-Bench tables.") | |
| parser.add_argument("--input", default="data/sample_week/normalized", help="Directory containing normalized CSV tables.") | |
| parser.add_argument("--output", default="data/sample_week/model_inputs", help="Output directory for task-level model inputs.") | |
| parser.add_argument("--limit", type=int, default=20, help="Maximum number of tasks to export.") | |
| return parser.parse_args(argv) | |
| def main(argv: List[str]) -> int: | |
| args = parse_args(argv) | |
| input_dir = Path(args.input) | |
| output_dir = Path(args.output) | |
| observations_path = input_dir / "observations.csv" | |
| tasks_path = input_dir / "forecast_tasks.csv" | |
| if not observations_path.exists() or not tasks_path.exists(): | |
| raise FileNotFoundError(f"Expected {observations_path} and {tasks_path}") | |
| observations = read_csv(observations_path) | |
| tasks = read_csv(tasks_path)[: args.limit] | |
| manifest = [export_task(task, observations, output_dir) for task in tasks] | |
| write_csv( | |
| output_dir / "manifest.csv", | |
| manifest, | |
| [ | |
| "task_id", | |
| "dataset_id", | |
| "entity_id", | |
| "target_variable_id", | |
| "context_rows", | |
| "future_target_rows", | |
| "context_csv", | |
| "future_target_csv", | |
| ], | |
| ) | |
| (output_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") | |
| print(f"Exported {len(manifest)} tasks to {output_dir}") | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main(sys.argv[1:])) | |