LiveHouse-TS / scripts /patch_snapshots.py
ziyuzhou02's picture
Deploy GitHub main 3feb6cda1511
e317359 verified
Raw
History Blame Contribute Delete
3.39 kB
from pathlib import Path
import json
import pandas as pd
from tsfm_bench.data.registry import load_data_source
source_config = Path("configs/datasets/ts_bench.yaml")
source = load_data_source(source_config)
# Map leaderboard name to task name
lb_to_task = {}
for task_name in source.list_datasets():
lb_name = source.leaderboard_dataset_name(task_name)
lb_to_task[lb_name] = task_name
# Now walk space/results/ and find all forecast JSONs
results_root = Path("space/results")
patched_count = 0
for p in results_root.glob("**/forecasts/*.json"):
try:
data = json.loads(p.read_text())
lb_dataset = data.get("dataset")
if lb_dataset in lb_to_task:
task_name = lb_to_task[lb_dataset]
record = next(iter(source.stream(task_name)))
n_context = data.get("context_length")
n_future = data.get("prediction_length")
start_period = pd.Period(record.start, freq=record.freq)
dr = pd.period_range(start=start_period, periods=n_context + n_future, freq=record.freq)
timestamps = [str(p.to_timestamp()) for p in dr]
data["timestamps"] = timestamps
p.write_text(json.dumps(data, indent=2) + "\n")
patched_count += 1
else:
# Fallback mapping if task names changed slightly
# We can find task that has the dataset name as substring
matched_task = None
dataset_name = lb_dataset.split("/")[0]
for t_name in source.list_datasets():
if dataset_name in t_name:
matched_task = t_name
break
if matched_task:
record = next(iter(source.stream(matched_task)))
n_context = data.get("context_length")
n_future = data.get("prediction_length")
start_period = pd.Period(record.start, freq=record.freq)
dr = pd.period_range(start=start_period, periods=n_context + n_future, freq=record.freq)
timestamps = [str(p.to_timestamp()) for p in dr]
data["timestamps"] = timestamps
p.write_text(json.dumps(data, indent=2) + "\n")
patched_count += 1
else:
# Fallback to evaluated_at date_range
eval_at = data.get("evaluated_at")
freq = data.get("freq", "1h")
n_context = data.get("context_length")
n_future = data.get("prediction_length")
try:
# Map '1h' to 'H' or 'h' if pandas complains, or handle basic units
p_freq = freq
if p_freq == "1h":
p_freq = "h"
end_dt = pd.to_datetime(eval_at).tz_convert(None)
dr = pd.date_range(end=end_dt, periods=n_context + n_future, freq=p_freq)
timestamps = [str(p) for p in dr]
data["timestamps"] = timestamps
p.write_text(json.dumps(data, indent=2) + "\n")
patched_count += 1
except Exception as e:
print(f"Dataset {lb_dataset} not found and fallback failed: {e}")
except Exception as e:
print(f"Error patching {p.name}: {e}")
print(f"Successfully patched {patched_count} forecast snapshots.")