Elgato_dark/chronos-forecasting
1
0
Fork 0
mirror of https://github.com/amazon-science/chronos-forecasting synced 2026-05-23 09:39:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Rename task -> input in dataset.py (#467)

Browse source 
*Issue #, if available:*

*Description of changes:*

Rename all variables and methods that refer to "task" in `dataset.py` to
use `input` instead:
- `PreparedTask` → `PreparedInput`
- `self.tasks` → `self.inputs`
- `prepare_tasks` → `prepare_inputs`
- `validate_and_prepare_single_dict_task` →
`validate_and_prepare_single_dict_input`
- All `task_` prefixed local variables renamed (e.g., `task_target` →
`target`, `task_context` → `context`, `task_past_tensor` →
`past_tensor`, etc.)


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.
This commit is contained in:
Oleksandr Shchur 2026-02-19 13:30:08 +01:00 committed by GitHub
parent 0e9db70afc
commit f951d9aefa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 164 additions and 164 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

322
src/chronos/chronos2/dataset.py
View file

@ -20,8 +20,8 @@ if TYPE_CHECKING:
TensorOrArray: TypeAlias = torch.Tensor | np.ndarray
class PreparedTask(TypedDict):
"""A preprocessed time series task ready for model training/inference."""
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
@ -47,14 +47,14 @@ def left_pad_and_cat_2D(tensors: list[torch.Tensor]) -> torch.Tensor:
return torch.cat(padded, dim=0)
def validate_and_prepare_single_dict_task(
task: Mapping[str, TensorOrArray | Mapping[str, TensorOrArray]], idx: int, prediction_length: int
) -> PreparedTask:
"""Validates and prepares a single dictionary task for Chronos2Model.
def validate_and_prepare_single_dict_input(
raw_input: Mapping[str, TensorOrArray | Mapping[str, TensorOrArray]], idx: int, prediction_length: int
) -> PreparedInput:
"""Validates and prepares a single dictionary input for Chronos2Model.
Parameters
----------
task
raw_input
A dictionary representing a time series that contains:
- `target` (required): a 1-d or 2-d `torch.Tensor` or `np.ndarray` of shape (history_length,) or (n_variates, history_length).
Forecasts will be generated for items in `target`.
@ -65,27 +65,27 @@ def validate_and_prepare_single_dict_task(
covariates and values must be 1-d `torch.Tensor` or `np.ndarray` with length equal to the `prediction_length`. All keys in
`future_covariates` must be a subset of the keys in `past_covariates`.
idx
Index of this task in the list of tasks, used for error messages
Index of this input in the list of inputs, used for error messages
prediction_length
Number of future time steps to predict, used to validate future covariates
Returns
------
A tuple containing:
- task_context_tensor: Concatenated tensor of target and past covariates of shape (group_size, history_length),
the first `task_n_targets` items along the first axis contain the target variables and the remaining items contain past-only covariates
A PreparedInput containing:
- context: Concatenated tensor of target and past covariates of shape (group_size, history_length),
the first `n_targets` items along the first axis contain the target variables and the remaining items contain past-only covariates
and past values of known future covariates.
- task_future_covariates_tensor: Tensor of future covariates of shape (group_size, prediction_length). The last `task_n_future_covariates`
- future_covariates: Tensor of future covariates of shape (group_size, prediction_length). The last `n_future_covariates`
items along the first axis contain future covariates. All the remaining elements corresponding to target and past-only covariates are NaNs.
- task_n_targets: Number of target variables
- task_n_covariates: Total number of covariates (sum of past-only and known future covariates)
- task_n_future_covariates: Number of known future covariates
- n_targets: Number of target variables
- n_covariates: Total number of covariates (sum of past-only and known future covariates)
- n_future_covariates: Number of known future covariates
"""
allowed_keys = {"target", "past_covariates", "future_covariates"}
# validate keys
keys = set(task.keys())
keys = set(raw_input.keys())
if not keys.issubset(allowed_keys):
raise ValueError(
f"Found invalid keys in element at index {idx}. Allowed keys are {allowed_keys}, but found {keys}"
@ -94,58 +94,58 @@ def validate_and_prepare_single_dict_task(
raise ValueError(f"Element at index {idx} does not contain the required key 'target'")
# validate target
task_target = task["target"]
if isinstance(task_target, np.ndarray):
task_target = torch.from_numpy(task_target)
assert isinstance(task_target, torch.Tensor)
if task_target.ndim > 2:
target = raw_input["target"]
if isinstance(target, np.ndarray):
target = torch.from_numpy(target)
assert isinstance(target, torch.Tensor)
if target.ndim > 2:
raise ValueError(
"When the input is a list of dicts, the `target` should either be 1-d with shape (history_length,) "
f" or 2-d with shape (n_variates, history_length). Found element at index {idx} with shape {tuple(task_target.shape)}."
f" or 2-d with shape (n_variates, history_length). Found element at index {idx} with shape {tuple(target.shape)}."
)
history_length = task_target.shape[-1]
task_target = task_target.view(-1, history_length)
history_length = target.shape[-1]
target = target.view(-1, history_length)
# validate past_covariates
cat_encoders: dict = {}
task_past_covariates = task.get("past_covariates", {})
if not isinstance(task_past_covariates, dict):
past_covariates = raw_input.get("past_covariates", {})
if not isinstance(past_covariates, dict):
raise ValueError(
f"Found invalid type for `past_covariates` in element at index {idx}. "
f'Expected dict with {{"feat_1": tensor_1, "feat_2": tensor_2, ...}}, but found {type(task_past_covariates)}'
f'Expected dict with {{"feat_1": tensor_1, "feat_2": tensor_2, ...}}, but found {type(past_covariates)}'
)
# gather keys and ensure known-future keys come last to match downstream assumptions
task_covariates_keys = sorted(task_past_covariates.keys())
covariates_keys = sorted(past_covariates.keys())
task_future_covariates = task.get("future_covariates", {})
if not isinstance(task_future_covariates, dict):
future_covariates = raw_input.get("future_covariates", {})
if not isinstance(future_covariates, dict):
raise ValueError(
f"Found invalid type for `future_covariates` in element at index {idx}. "
f'Expected dict with {{"feat_1": tensor_1, "feat_2": tensor_2, ...}}, but found {type(task_future_covariates)}'
f'Expected dict with {{"feat_1": tensor_1, "feat_2": tensor_2, ...}}, but found {type(future_covariates)}'
)
task_future_covariates_keys = sorted(task_future_covariates.keys())
if not set(task_future_covariates_keys).issubset(task_covariates_keys):
future_covariates_keys = sorted(future_covariates.keys())
if not set(future_covariates_keys).issubset(covariates_keys):
raise ValueError(
f"Expected keys in `future_covariates` to be a subset of `past_covariates` {task_covariates_keys}, "
f"but found {task_future_covariates_keys} in element at index {idx}"
f"Expected keys in `future_covariates` to be a subset of `past_covariates` {covariates_keys}, "
f"but found {future_covariates_keys} in element at index {idx}"
)
# create ordered keys: past-only first, then known-future (so known-future are the last rows)
task_past_only_keys = [k for k in task_covariates_keys if k not in task_future_covariates_keys] # past_only_keys
task_ordered_covariate_keys = task_past_only_keys + task_future_covariates_keys
past_only_keys = [k for k in covariates_keys if k not in future_covariates_keys]
ordered_covariate_keys = past_only_keys + future_covariates_keys
task_past_covariates_list: list[torch.Tensor] = []
for key in task_ordered_covariate_keys:
tensor = task_past_covariates[key]
past_covariates_list: list[torch.Tensor] = []
for key in ordered_covariate_keys:
tensor = past_covariates[key]
if isinstance(tensor, np.ndarray):
# apply encoding to categorical variates
if not np.issubdtype(tensor.dtype, np.number):
# target encoding, if the target is 1-d
if task_target.shape[0] == 1:
if target.shape[0] == 1:
cat_encoder = TargetEncoder(target_type="continuous", smooth=1.0)
X = tensor.astype(str).reshape(-1, 1)
y = task_target.view(-1).numpy()
y = target.view(-1).numpy()
mask = np.isfinite(y)
X = X[mask]
y = y[mask]
@ -163,18 +163,18 @@ def validate_and_prepare_single_dict_task(
f"Individual `past_covariates` must be 1-d with length equal to the length of `target` (= {history_length}), "
f"found: {key} with shape {tuple(tensor.shape)} in element at index {idx}"
)
task_past_covariates_list.append(tensor)
task_past_covariates_tensor = (
torch.stack(task_past_covariates_list, dim=0)
if task_past_covariates_list
else torch.zeros((0, history_length), device=task_target.device)
past_covariates_list.append(tensor)
past_covariates_tensor = (
torch.stack(past_covariates_list, dim=0)
if past_covariates_list
else torch.zeros((0, history_length), device=target.device)
)
# validate future_covariates (build rows in the same task_ordered_covariate_keys order)
task_future_covariates_list: list[torch.Tensor] = []
for key in task_ordered_covariate_keys:
# validate future_covariates (build rows in the same ordered_covariate_keys order)
future_covariates_list: list[torch.Tensor] = []
for key in ordered_covariate_keys:
# future values of past-only covariates are filled with NaNs
tensor = task_future_covariates.get(key, torch.full((prediction_length,), fill_value=torch.nan))
tensor = future_covariates.get(key, torch.full((prediction_length,), fill_value=torch.nan))
if isinstance(tensor, np.ndarray):
# apply encoding to categorical variates
if not np.issubdtype(tensor.dtype, np.number):
@ -187,52 +187,52 @@ def validate_and_prepare_single_dict_task(
f"Individual `future_covariates` must be 1-d with length equal to the {prediction_length=}, "
f"found: {key} with shape {tuple(tensor.shape)} in element at index {idx}"
)
task_future_covariates_list.append(tensor)
task_future_covariates_tensor = (
torch.stack(task_future_covariates_list, dim=0)
if task_future_covariates_list
else torch.zeros((0, prediction_length), device=task_target.device)
future_covariates_list.append(tensor)
future_covariates_tensor = (
torch.stack(future_covariates_list, dim=0)
if future_covariates_list
else torch.zeros((0, prediction_length), device=target.device)
)
# future values of target series are filled with NaNs
task_future_covariates_target_padding = torch.full(
(task_target.shape[0], prediction_length), fill_value=torch.nan, device=task_target.device
future_covariates_target_padding = torch.full(
(target.shape[0], prediction_length), fill_value=torch.nan, device=target.device
)
task_context_tensor = torch.cat([task_target, task_past_covariates_tensor], dim=0).to(dtype=torch.float32)
task_future_covariates_tensor = torch.cat(
[task_future_covariates_target_padding, task_future_covariates_tensor], dim=0
context_tensor = torch.cat([target, past_covariates_tensor], dim=0).to(dtype=torch.float32)
future_covariates_tensor = torch.cat(
[future_covariates_target_padding, future_covariates_tensor], dim=0
).to(dtype=torch.float32)
task_n_targets = task_target.shape[0]
task_n_covariates = task_past_covariates_tensor.shape[0]
n_targets = target.shape[0]
n_covariates = past_covariates_tensor.shape[0]
# number of known-future covariates
task_n_future_covariates = len(task_future_covariates_keys)
n_future_covariates = len(future_covariates_keys)
return PreparedTask(
context=task_context_tensor,
future_covariates=task_future_covariates_tensor,
n_targets=task_n_targets,
n_covariates=task_n_covariates,
n_future_covariates=task_n_future_covariates,
return PreparedInput(
context=context_tensor,
future_covariates=future_covariates_tensor,
n_targets=n_targets,
n_covariates=n_covariates,
n_future_covariates=n_future_covariates,
)
def prepare_tasks(
raw_tasks: Iterable[Mapping[str, Any]],
def prepare_inputs(
raw_inputs: Iterable[Mapping[str, Any]],
prediction_length: int,
min_past: int = 1,
mode: "DatasetMode | str" = "train",
) -> list[PreparedTask]:
"""Prepare multiple time series tasks for training/inference.
) -> list[PreparedInput]:
"""Prepare multiple time series inputs for training/inference.
This function handles mode-specific preprocessing (e.g., filtering short series)
and calls validate_and_prepare_single_dict_task for each task.
and calls validate_and_prepare_single_dict_input for each input.
"""
tasks: list[PreparedTask] = []
inputs: list[PreparedInput] = []
for idx, raw_task in enumerate(raw_tasks):
for idx, raw_input in enumerate(raw_inputs):
# For non-TEST modes, fix future_covariates (replace None/empty with NaN arrays)
if mode != DatasetMode.TEST:
raw_future_covariates = raw_task.get("future_covariates", {})
raw_future_covariates = raw_input.get("future_covariates", {})
if raw_future_covariates:
raw_future_covariates = cast(dict[str, TensorOrArray | None], raw_future_covariates)
fixed_future_covariates = {}
@ -240,42 +240,42 @@ def prepare_tasks(
fixed_future_covariates[key] = (
np.full(prediction_length, np.nan) if value is None or len(value) == 0 else value
)
raw_task = {**raw_task, "future_covariates": fixed_future_covariates}
raw_input = {**raw_input, "future_covariates": fixed_future_covariates}
raw_task = cast(dict[str, TensorOrArray | Mapping[str, TensorOrArray]], raw_task)
task = validate_and_prepare_single_dict_task(raw_task, idx, prediction_length)
raw_input = cast(dict[str, TensorOrArray | Mapping[str, TensorOrArray]], raw_input)
prepared = validate_and_prepare_single_dict_input(raw_input, idx, prediction_length)
# Filter by minimum length (except in TEST mode)
if mode != DatasetMode.TEST and task["context"].shape[-1] < min_past + prediction_length:
if mode != DatasetMode.TEST and prepared["context"].shape[-1] < min_past + prediction_length:
continue
tasks.append(task)
inputs.append(prepared)
if len(tasks) == 0:
if len(inputs) == 0:
raise ValueError(
"The dataset is empty after filtering based on the length of the time series (length >= min_past + prediction_length). "
"Please provide longer time series or reduce `min_past` or `prediction_length`. "
)
return tasks
return inputs
def validate_prepared_schema(task: Any) -> None:
"""Validate that a task matches the PreparedTask schema."""
if not isinstance(task, Mapping):
def validate_prepared_schema(prepared_input: Any) -> None:
"""Validate that an input matches the PreparedInput schema."""
if not isinstance(prepared_input, Mapping):
raise TypeError(
f"Expected task to be a dict-like, got {type(task).__name__}. "
f"Expected input to be a dict-like, got {type(prepared_input).__name__}. "
"Set convert_inputs=True when calling fit() to preprocess raw inputs."
)
required_keys = {"context", "future_covariates", "n_targets", "n_covariates", "n_future_covariates"}
missing = required_keys - set(task.keys())
missing = required_keys - set(prepared_input.keys())
if missing:
raise TypeError(
f"Task is missing required keys: {missing}. Set convert_inputs=True when calling fit() to preprocess raw inputs."
f"Input is missing required keys: {missing}. Set convert_inputs=True when calling fit() to preprocess raw inputs."
)
context = task["context"]
context = prepared_input["context"]
if not isinstance(context, torch.Tensor) or context.ndim != 2:
raise TypeError(
f"Expected 'context' to be 2-d torch.Tensor, got {type(context).__name__} "
@ -283,7 +283,7 @@ def validate_prepared_schema(task: Any) -> None:
"Set convert_inputs=True when calling fit() to preprocess raw inputs."
)
future_covariates = task["future_covariates"]
future_covariates = prepared_input["future_covariates"]
if not isinstance(future_covariates, torch.Tensor) or future_covariates.ndim != 2:
raise TypeError(
f"Expected 'future_covariates' to be 2-d torch.Tensor, got {type(future_covariates).__name__} "
@ -486,9 +486,9 @@ class Chronos2Dataset(IterableDataset):
covariates.
- `future_covariates` (optional): a dict of future values of known future covariates.
2. Pre-processed inputs (when `convert_inputs=False`): A sequence of `PreparedTask` dicts with keys:
2. Pre-processed inputs (when `convert_inputs=False`): A sequence of `PreparedInput` dicts with keys:
`context`, `future_covariates`, `n_targets`, `n_covariates`, `n_future_covariates`.
Use `prepare_tasks()` to create pre-processed inputs.
Use `prepare_inputs()` to create pre-processed inputs.
context_length
The maximum context length used for training or inference
prediction_length
@ -510,7 +510,7 @@ class Chronos2Dataset(IterableDataset):
def __init__(
self,
inputs: TensorOrArray | Sequence[TensorOrArray] | Sequence[Mapping[str, Any]] | Sequence[PreparedTask],
inputs: TensorOrArray | Sequence[TensorOrArray] | Sequence[Mapping[str, Any]] | Sequence[PreparedInput],
context_length: int,
prediction_length: int,
batch_size: int,
@ -522,7 +522,7 @@ class Chronos2Dataset(IterableDataset):
super().__init__()
assert mode in {DatasetMode.TRAIN, DatasetMode.VALIDATION, DatasetMode.TEST}, f"Invalid mode: {mode}"
self.tasks: Sequence[PreparedTask]
self.inputs: Sequence[PreparedInput]
if convert_inputs:
if isinstance(inputs, (torch.Tensor, np.ndarray)):
inputs = convert_tensor_input_to_list_of_dicts_input(inputs)
@ -530,10 +530,10 @@ class Chronos2Dataset(IterableDataset):
isinstance(inputs, Sequence) and len(inputs) > 0 and isinstance(inputs[0], (torch.Tensor, np.ndarray))
):
inputs = convert_list_of_tensors_input_to_list_of_dicts_input(cast(Sequence[TensorOrArray], inputs))
self.tasks = prepare_tasks(cast(Iterable[Mapping[str, Any]], inputs), prediction_length, min_past, mode)
self.inputs = prepare_inputs(cast(Iterable[Mapping[str, Any]], inputs), prediction_length, min_past, mode)
else:
validate_prepared_schema(inputs[0])
self.tasks = cast(Sequence[PreparedTask], inputs)
self.inputs = cast(Sequence[PreparedInput], inputs)
self.context_length = context_length
self.prediction_length = prediction_length
@ -542,16 +542,16 @@ class Chronos2Dataset(IterableDataset):
self.min_past = min_past
self.mode = mode
def _construct_slice(self, task_idx: int) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, int]:
task = self.tasks[task_idx]
task_past_tensor = task["context"].clone() # shape: (task_n_targets + task_n_covariates, history_length)
task_future_tensor = task["future_covariates"].clone()
task_n_targets = int(task["n_targets"])
task_n_covariates = int(task["n_covariates"])
task_n_future_covariates = int(task["n_future_covariates"])
task_n_past_only_covariates = task_n_covariates - task_n_future_covariates
def _construct_slice(self, input_idx: int) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, int]:
prepared = self.inputs[input_idx]
past_tensor = prepared["context"].clone() # shape: (n_targets + n_covariates, history_length)
future_tensor = prepared["future_covariates"].clone()
n_targets = int(prepared["n_targets"])
n_covariates = int(prepared["n_covariates"])
n_future_covariates = int(prepared["n_future_covariates"])
n_past_only_covariates = n_covariates - n_future_covariates
full_length = task_past_tensor.shape[-1]
full_length = past_tensor.shape[-1]
if self.mode == DatasetMode.TRAIN:
# slice a random subsequence from the full series
@ -565,74 +565,74 @@ class Chronos2Dataset(IterableDataset):
if slice_idx >= self.context_length:
# slice series, if it is longer than context_length
task_context = task_past_tensor[:, slice_idx - self.context_length : slice_idx]
context = past_tensor[:, slice_idx - self.context_length : slice_idx]
else:
task_context = task_past_tensor[:, :slice_idx]
context = past_tensor[:, :slice_idx]
# In the TEST mode, we have no target available and the task_future_covariates can be directly used
# In the TRAIN and VALIDATION modes, the target and task_future_covariates need to be constructed from
# the task_context_tensor by slicing the appropriate indices which we do below
# In the TEST mode, we have no target available and the future_covariates can be directly used
# In the TRAIN and VALIDATION modes, the target and future_covariates need to be constructed from
# the context_tensor by slicing the appropriate indices which we do below
if self.mode in [DatasetMode.TRAIN, DatasetMode.VALIDATION]:
# the first task_n_targets elements in task_context_tensor are the targets
task_future_target = task_past_tensor[:, slice_idx : slice_idx + self.prediction_length].clone()
# the first n_targets elements in context_tensor are the targets
future_target = past_tensor[:, slice_idx : slice_idx + self.prediction_length].clone()
# mask out all rows corresponding to covariates
task_future_target[task_n_targets:] = torch.nan
future_target[n_targets:] = torch.nan
if task_n_future_covariates > 0:
# the last task_n_future_covariates elements in task_context_tensor are the known covariates
task_future_covariates = task_past_tensor[
-task_n_future_covariates:, slice_idx : slice_idx + self.prediction_length
if n_future_covariates > 0:
# the last n_future_covariates elements in context_tensor are the known covariates
future_covariates = past_tensor[
-n_future_covariates:, slice_idx : slice_idx + self.prediction_length
]
else:
# zero-length tensor for easy concatenation later
task_future_covariates = torch.zeros((0, self.prediction_length))
future_covariates = torch.zeros((0, self.prediction_length))
# the leading task_n_targets + task_n_past_only_covariates elements are masked because the target(s)
# the leading n_targets + n_past_only_covariates elements are masked because the target(s)
# and past-only covariates are not known into the future
task_future_covariates_padding = torch.full(
(task_n_targets + task_n_past_only_covariates, self.prediction_length),
future_covariates_padding = torch.full(
(n_targets + n_past_only_covariates, self.prediction_length),
fill_value=torch.nan,
)
task_future_covariates = torch.cat([task_future_covariates_padding, task_future_covariates], dim=0)
future_covariates = torch.cat([future_covariates_padding, future_covariates], dim=0)
else:
task_future_target = None
task_future_covariates = task_future_tensor
future_target = None
future_covariates = future_tensor
# task_context: (task_n_targets + task_n_covariates, min(context_length, history_length))
# task_future_target: (task_n_targets + task_n_covariates, prediction_length), the future values of known future covariates
# context: (n_targets + n_covariates, min(context_length, history_length))
# future_target: (n_targets + n_covariates, prediction_length), the future values of known future covariates
# are ignored during loss computation
# task_future_covariates: (task_n_targets + task_n_past_only_covariates + task_n_future_covariates, prediction_length),
# future_covariates: (n_targets + n_past_only_covariates + n_future_covariates, prediction_length),
# the entries corresponding to targets and past-only covariates are NaNs
return task_context, task_future_target, task_future_covariates, task_n_targets
return context, future_target, future_covariates, n_targets
def _build_batch(self, task_indices: list[int]) -> dict[str, torch.Tensor | int | list[tuple[int, int]] | None]:
"""Build a batch from given task indices."""
batch_context_tensor_list = []
batch_future_target_tensor_list = []
batch_future_covariates_tensor_list = []
def _build_batch(self, input_indices: list[int]) -> dict[str, torch.Tensor | int | list[tuple[int, int]] | None]:
"""Build a batch from given input indices."""
batch_context_list = []
batch_future_target_list = []
batch_future_covariates_list = []
batch_group_ids_list = []
target_idx_ranges: list[tuple[int, int]] = []
target_start_idx = 0
for group_id, task_idx in enumerate(task_indices):
task_context, task_future_target, task_future_covariates, task_n_targets = self._construct_slice(task_idx)
for group_id, input_idx in enumerate(input_indices):
context, future_target, future_covariates, n_targets = self._construct_slice(input_idx)
group_size = task_context.shape[0]
task_group_ids = torch.full((group_size,), fill_value=group_id)
batch_context_tensor_list.append(task_context)
batch_future_target_tensor_list.append(task_future_target)
batch_future_covariates_tensor_list.append(task_future_covariates)
batch_group_ids_list.append(task_group_ids)
target_idx_ranges.append((target_start_idx, target_start_idx + task_n_targets))
group_size = context.shape[0]
group_ids = torch.full((group_size,), fill_value=group_id)
batch_context_list.append(context)
batch_future_target_list.append(future_target)
batch_future_covariates_list.append(future_covariates)
batch_group_ids_list.append(group_ids)
target_idx_ranges.append((target_start_idx, target_start_idx + n_targets))
target_start_idx += group_size
return {
"context": left_pad_and_cat_2D(batch_context_tensor_list),
"context": left_pad_and_cat_2D(batch_context_list),
"future_target": None
if self.mode == DatasetMode.TEST
else torch.cat(cast(list[torch.Tensor], batch_future_target_tensor_list), dim=0),
"future_covariates": torch.cat(batch_future_covariates_tensor_list, dim=0),
else torch.cat(cast(list[torch.Tensor], batch_future_target_list), dim=0),
"future_covariates": torch.cat(batch_future_covariates_list, dim=0),
"group_ids": torch.cat(batch_group_ids_list, dim=0),
"num_output_patches": self.num_output_patches,
"target_idx_ranges": target_idx_ranges,
@ -641,27 +641,27 @@ class Chronos2Dataset(IterableDataset):
def _generate_train_batches(self):
while True:
current_batch_size = 0
task_indices = []
input_indices = []
while current_batch_size < self.batch_size:
task_idx = np.random.randint(len(self.tasks))
task_indices.append(task_idx)
current_batch_size += self.tasks[task_idx]["context"].shape[0]
input_idx = np.random.randint(len(self.inputs))
input_indices.append(input_idx)
current_batch_size += self.inputs[input_idx]["context"].shape[0]
yield self._build_batch(task_indices)
yield self._build_batch(input_indices)
def _generate_sequential_batches(self):
task_idx = 0
while task_idx < len(self.tasks):
input_idx = 0
while input_idx < len(self.inputs):
current_batch_size = 0
task_indices = []
input_indices = []
while task_idx < len(self.tasks) and current_batch_size < self.batch_size:
task_indices.append(task_idx)
current_batch_size += self.tasks[task_idx]["context"].shape[0]
task_idx += 1
while input_idx < len(self.inputs) and current_batch_size < self.batch_size:
input_indices.append(input_idx)
current_batch_size += self.inputs[input_idx]["context"].shape[0]
input_idx += 1
yield self._build_batch(task_indices)
yield self._build_batch(input_indices)
def __iter__(self) -> Iterator:
"""

2
src/chronos/chronos2/pipeline.py
View file

@ -164,7 +164,7 @@ class Chronos2Pipeline(BaseChronosPipeline):
If True, ensures that DataParallel is disabled and training happens on a single GPU
convert_inputs
If True (default), preprocess raw inputs (convert tensors, encode categoricals, validate).
If False, inputs are expected to be already preprocessed using `chronos.chronos2.dataset.prepare_tasks`.
If False, inputs are expected to be already preprocessed using `chronos.chronos2.dataset.prepare_inputs`.
This allows for efficient training on large datasets that don't fit in memory.
**extra_trainer_kwargs
Extra kwargs are directly forwarded to `TrainingArguments`

4
test/test_chronos2.py
View file

@ -1147,13 +1147,13 @@ def test_eager_and_sdpa_produce_identical_outputs(pipeline):
def test_pipeline_can_be_finetuned_with_preprocessed_hf_dataset(pipeline):
"""Test that fine-tuning works with preprocessed inputs from a HuggingFace Dataset."""
from chronos.chronos2.dataset import prepare_tasks
from chronos.chronos2.dataset import prepare_inputs
prediction_length = 8
raw_inputs = [{"target": torch.rand(20)}, {"target": torch.rand(25)}, {"target": torch.rand(30)}]
# Preprocess and convert to HF Dataset (simulating Arrow-based lazy loading)
prepared_tasks = prepare_tasks(raw_inputs, prediction_length=prediction_length, min_past=1, mode="train")
prepared_tasks = prepare_inputs(raw_inputs, prediction_length=prediction_length, min_past=1, mode="train")
hf_dataset = datasets.Dataset.from_list(prepared_tasks).with_format("torch")
# Fine-tune with preprocessed inputs