diff --git a/src/uni2ts/data/builder/simple.py b/src/uni2ts/data/builder/simple.py index bb54215..c283c54 100644 --- a/src/uni2ts/data/builder/simple.py +++ b/src/uni2ts/data/builder/simple.py @@ -14,6 +14,7 @@ # limitations under the License. import argparse +from collections import defaultdict from dataclasses import dataclass from itertools import product from pathlib import Path @@ -26,18 +27,19 @@ from uni2ts.common.env import env from uni2ts.common.typing import GenFunc +from uni2ts.data.builder._base import DatasetBuilder from uni2ts.data.dataset import EvalDataset, SampleTimeSeriesType, TimeSeriesDataset from uni2ts.data.indexer import HuggingFaceDatasetIndexer from uni2ts.transform import Transformation -from uni2ts.data.builder._base import DatasetBuilder - - -# Manually set the freq of the datasets whose freq can be inferred automatically. -freq_dict = { - 'weather': '10T', 'weather_eval': '10T', - -} +# Manually set the freq of the datasets whose freq can be inferred automatically. Default freq is H. +freq_dict = defaultdict( + lambda: "H", + { + "weather": "10T", + "weather_eval": "10T", + }, +) def _from_long_dataframe( @@ -58,7 +60,11 @@ def example_gen_func() -> Generator[dict[str, Any], None, None]: yield { "target": item_df.to_numpy(), "start": item_df.index[0], - "freq": pd.infer_freq(df.index) if pd.infer_freq(df.index) is not None else freq_dict[dataset], + "freq": ( + pd.infer_freq(df.index) + if pd.infer_freq(df.index) is not None + else freq_dict[dataset] + ), "item_id": item_id, } @@ -92,7 +98,11 @@ def example_gen_func() -> Generator[dict[str, Any], None, None]: yield { "target": df.iloc[:, i].to_numpy(), "start": df.index[0], - "freq": pd.infer_freq(df.index) if pd.infer_freq(df.index) is not None else freq_dict[dataset], + "freq": ( + pd.infer_freq(df.index) + if pd.infer_freq(df.index) is not None + else freq_dict[dataset] + ), "item_id": f"item_{i}", } @@ -123,7 +133,11 @@ def example_gen_func() -> Generator[dict[str, Any], None, None]: yield { "target": df.to_numpy().T, "start": df.index[0], - "freq": pd.infer_freq(df.index) if pd.infer_freq(df.index) is not None else freq_dict[dataset], + "freq": ( + pd.infer_freq(df.index) + if pd.infer_freq(df.index) is not None + else freq_dict[dataset] + ), "item_id": "item_0", }