Refactor preprocessing logic for Chronos-2

src/chronos/chronos2/dataset.py

@@ -5,13 +5,15 @@
 
 import math
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, Sequence, TypeAlias, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, Sequence, TypeAlias, cast
 
 import numpy as np
 import torch
 from sklearn.preprocessing import OrdinalEncoder, TargetEncoder
 from torch.utils.data import IterableDataset
 
+from chronos.chronos2.preprocess import PreparedInput
+
 if TYPE_CHECKING:
     import datasets
     import fev
@@ -20,16 +22,6 @@
 TensorOrArray: TypeAlias = torch.Tensor | np.ndarray
 
 
-class PreparedInput(TypedDict):
-    """A preprocessed time series input ready for model training/inference."""
-
-    context: torch.Tensor  # (n_variates, history_length), float32
-    future_covariates: torch.Tensor  # (n_variates, prediction_length), float32
-    n_targets: int
-    n_covariates: int
-    n_future_covariates: int
-
-
 def left_pad_and_cat_2D(tensors: list[torch.Tensor]) -> torch.Tensor:
     """
     Left pads tensors in the list to the length of the longest tensor along the second axis, then concats