from gluonts.time_feature import get_lags_for_frequency
lags_sequence = get_lags_for_frequency(freq)print(lags_sequence)
>>> [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 23, 24, 25, 35, 36, 37]

from gluonts.time_feature import (    time_features_from_frequency_str,    TimeFeature,    get_lags_for_frequency,)from gluonts.dataset.field_names import FieldNamefrom gluonts.transform import (    AddAgeFeature,    AddObservedValuesIndicator,    AddTimeFeatures,    AsNumpyArray,    Chain,    ExpectedNumInstanceSampler,    InstanceSplitter,    RemoveFields,    SelectFields,    SetField,    TestSplitSampler,    Transformation,    ValidationSplitSampler,    VstackFeatures,    RenameFields,)

from transformers import PretrainedConfig
def create_transformation(freq: str, config: PretrainedConfig) -> Transformation:    remove_field_names = []    if config.num_static_real_features == 0:        remove_field_names.append(FieldName.FEAT_STATIC_REAL)    if config.num_dynamic_real_features == 0:        remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)    if config.num_static_categorical_features == 0:        remove_field_names.append(FieldName.FEAT_STATIC_CAT)
    # a bit like torchvision.transforms.Compose    return Chain(        # step 1: remove static/dynamic fields if not specified        [RemoveFields(field_names=remove_field_names)]        # step 2: convert the data to NumPy (potentially not needed)        + (            [                AsNumpyArray(                    field=FieldName.FEAT_STATIC_CAT,                    expected_ndim=1,                    dtype=int,                )            ]            if config.num_static_categorical_features > 0            else []        )        + (            [                AsNumpyArray(                    field=FieldName.FEAT_STATIC_REAL,                    expected_ndim=1,                )            ]            if config.num_static_real_features > 0            else []        )        + [            AsNumpyArray(                field=FieldName.TARGET,                # we expect an extra dim for the multivariate case:                expected_ndim=1 if config.input_size == 1 else 2,            ),            # step 3: handle the NaN's by filling in the target with zero            # and return the mask (which is in the observed values)            # true for observed values, false for nan's            # the decoder uses this mask (no loss is incurred for unobserved values)            # see loss_weights inside the xxxForPrediction model            AddObservedValuesIndicator(                target_field=FieldName.TARGET,                output_field=FieldName.OBSERVED_VALUES,            ),            # step 4: add temporal features based on freq of the dataset            # month of year in the case when freq="M"            # these serve as positional encodings            AddTimeFeatures(                start_field=FieldName.START,                target_field=FieldName.TARGET,                output_field=FieldName.FEAT_TIME,                time_features=time_features_from_frequency_str(freq),                pred_length=config.prediction_length,            ),            # step 5: add another temporal feature (just a single number)            # tells the model where in its life the value of the time series is,            # sort of a running counter            AddAgeFeature(                target_field=FieldName.TARGET,                output_field=FieldName.FEAT_AGE,                pred_length=config.prediction_length,                log_scale=True,            ),            # step 6: vertically stack all the temporal features into the key FEAT_TIME            VstackFeatures(                output_field=FieldName.FEAT_TIME,                input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]                + (                    [FieldName.FEAT_DYNAMIC_REAL]                    if config.num_dynamic_real_features > 0                    else []                ),            ),            # step 7: rename to match HuggingFace names            RenameFields(                mapping={                    FieldName.FEAT_STATIC_CAT: "static_categorical_features",                    FieldName.FEAT_STATIC_REAL: "static_real_features",                    FieldName.FEAT_TIME: "time_features",                    FieldName.TARGET: "values",                    FieldName.OBSERVED_VALUES: "observed_mask",                }            ),        ]    )

from gluonts.transform.sampler import InstanceSamplerfrom typing import Optional
def create_instance_splitter(    config: PretrainedConfig,    mode: str,    train_sampler: Optional[InstanceSampler] = None,    validation_sampler: Optional[InstanceSampler] = None,) -> Transformation:    assert mode in ["train", "validation", "test"]
    instance_sampler = {        "train": train_sampler        or ExpectedNumInstanceSampler(            num_instances=1.0, min_future=config.prediction_length        ),        "validation": validation_sampler        or ValidationSplitSampler(min_future=config.prediction_length),        "test": TestSplitSampler(),    }[mode]
    return InstanceSplitter(        target_field="values",        is_pad_field=FieldName.IS_PAD,        start_field=FieldName.START,        forecast_start_field=FieldName.FORECAST_START,        instance_sampler=instance_sampler,        past_length=config.context_length + max(config.lags_sequence),        future_length=config.prediction_length,        time_series_fields=["time_features", "observed_mask"],    )

def create_test_dataloader(    config: PretrainedConfig,    freq,    data,    batch_size: int,    **kwargs,):    PREDICTION_INPUT_NAMES = [        "past_time_features",        "past_values",        "past_observed_mask",        "future_time_features",    ]    if config.num_static_categorical_features > 0:        PREDICTION_INPUT_NAMES.append("static_categorical_features")
    if config.num_static_real_features > 0:        PREDICTION_INPUT_NAMES.append("static_real_features")
    transformation = create_transformation(freq, config)    transformed_data = transformation.apply(data, is_train=False)
    # we create a Test Instance splitter which will sample the very last    # context window seen during training only for the encoder.    instance_sampler = create_instance_splitter(config, "test")
    # we apply the transformations in test mode    testing_instances = instance_sampler.apply(transformed_data, is_train=False)
    return as_stacked_batches(        testing_instances,        batch_size=batch_size,        output_type=torch.tensor,        field_names=PREDICTION_INPUT_NAMES,    )

train_dataloader = create_train_dataloader(    config=config,    freq=freq,    data=train_dataset,    batch_size=256,    num_batches_per_epoch=100,)
test_dataloader = create_test_dataloader(    config=config,    freq=freq,    data=test_dataset,    batch_size=64,)

batch = next(iter(train_dataloader))for k, v in batch.items():    print(k, v.shape, v.type())
>>> past_time_features torch.Size([256, 85, 2]) torch.FloatTensor    past_values torch.Size([256, 85]) torch.FloatTensor    past_observed_mask torch.Size([256, 85]) torch.FloatTensor    future_time_features torch.Size([256, 24, 2]) torch.FloatTensor    static_categorical_features torch.Size([256, 1]) torch.LongTensor    future_values torch.Size([256, 24]) torch.FloatTensor    future_observed_mask torch.Size([256, 24]) torch.FloatTensor

# perform forward passoutputs = model(    past_values=batch["past_values"],    past_time_features=batch["past_time_features"],    past_observed_mask=batch["past_observed_mask"],    static_categorical_features=batch["static_categorical_features"]    if config.num_static_categorical_features > 0    else None,    static_real_features=batch["static_real_features"]    if config.num_static_real_features > 0    else None,    future_values=batch["future_values"],    future_time_features=batch["future_time_features"],    future_observed_mask=batch["future_observed_mask"],    output_hidden_states=True,)

print("Loss:", outputs.loss.item())
>>> Loss: 9.069628715515137

from accelerate import Acceleratorfrom torch.optim import AdamW
accelerator = Accelerator()device = accelerator.device
model.to(device)optimizer = AdamW(model.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=1e-1)
model, optimizer, train_dataloader = accelerator.prepare(    model,    optimizer,    train_dataloader,)
model.train()for epoch in range(40):    for idx, batch in enumerate(train_dataloader):        optimizer.zero_grad()        outputs = model(            static_categorical_features=batch["static_categorical_features"].to(device)            if config.num_static_categorical_features > 0            else None,            static_real_features=batch["static_real_features"].to(device)            if config.num_static_real_features > 0            else None,            past_time_features=batch["past_time_features"].to(device),            past_values=batch["past_values"].to(device),            future_time_features=batch["future_time_features"].to(device),            future_values=batch["future_values"].to(device),            past_observed_mask=batch["past_observed_mask"].to(device),            future_observed_mask=batch["future_observed_mask"].to(device),        )        loss = outputs.loss
        # Backpropagation        accelerator.backward(loss)        optimizer.step()
        if idx % 100 == 0:            print(loss.item())

model.eval()
forecasts = []
for batch in test_dataloader:    outputs = model.generate(        static_categorical_features=batch["static_categorical_features"].to(device)        if config.num_static_categorical_features > 0        else None,        static_real_features=batch["static_real_features"].to(device)        if config.num_static_real_features > 0        else None,        past_time_features=batch["past_time_features"].to(device),        past_values=batch["past_values"].to(device),        future_time_features=batch["future_time_features"].to(device),        past_observed_mask=batch["past_observed_mask"].to(device),    )    forecasts.append(outputs.sequences.cpu().numpy())

forecasts[0].shape
>>> (64, 100, 24)

forecasts = np.vstack(forecasts)print(forecasts.shape)
>>> (366, 100, 24)