code: def main(
      data_dir: str,
      output_dir: str,
      seed: int,
      # Experiment setup
      model_name: str,
      quantiles: str,
      dm_name: str,
      support_set_size: int,
      standardize: bool,
      # Common hyperparameters
      pretrained: str,
      learning_rate: float,
      lr_on_plateau_patience: float,
      lr_on_plateau_factor: float,
      context_length_multiple: int,
      support_length_multiple: int,
      model_prediction_length: Optional[int],
      max_epochs: int,
      max_training_time: str,
      num_batches: int,
      batch_size_train: int,
      batch_size_val_test: int,
      accumulate_grad_batches: Optional[int],
      early_stopping_min_delta: float,
      early_stopping_patience: int,
      monitor_metric: str,
      n_logging_samples: int,
      log_plot_every_n_epochs: int,
      num_workers: int,
      # Model hyperparameters
      **kwargs: int,
  ):
      """
      Trains a model in the meta learning framework
      """
      args_to_save = locals()
      random.seed(seed)
      torch.manual_seed(seed)
      pl.seed_everything(seed, workers=True)
      kwargs["dm_super_seed"] = seed
  
      quantiles = quantiles.split(",")
  
      # -------------------- Configure the data loader ---------------------------------------------------
      dm = get_data_module(
          dm_name,
          support_set_size=support_set_size,
          support_length_multiple=support_length_multiple,
          standardize=standardize,
          data_dir=Path(data_dir),
          batch_size_train=batch_size_train,
          batch_size_val_test=batch_size_val_test,
          context_length_multiple=context_length_multiple,
          num_workers=num_workers,
          **{
              key[len(dm_name) + 1 :]: value
              for key, value in kwargs.items()
              if key.startswith(dm_name)
          },
      )
      dm.setup()
      # -------------------- Initialize the model ---------------------------------------------------
      model = get_model(
          model_name,
          prediction_length=model_prediction_length
          or dm.prediction_length,
          quantiles=quantiles,
          **{
              key[len(model_name) + 1 :]: value
              for key, value in kwargs.items()
              if key.startswith(model_name)
          },
      )
      if pretrained:
          model = load_weights(model=model, path_to_weights=pretrained)
  
      module = MetaLightningModule(
          model,
          loss=QuantileLoss(quantiles),
          crps=CRPS(quantiles, rescale=dm.standardize),
          lr_scheduler_monitor=monitor_metric,
          crps_scaled=CRPS(quantiles) if dm.standardize else None,
          quantile_width=QuantileWidth(quantiles),
          lr=learning_rate,
          lr_on_plateau_patience=lr_on_plateau_patience,
          lr_on_plateau_factor=lr_on_plateau_factor,
          val_dataset_names=dm.dataset_names_val,
          test_dataset_names=dm.dataset_names_test,
      )
  
      # -------------------- configure callbacks ---------------------------------------------------
      # for the plotting callback we need to get some validation data
      log_batch_train, log_batch_val = dm.get_log_batches(
          n_logging_samples
      )
  
      # early stopping callback
      early_stop_callback = EarlyStopping(
          monitor=monitor_metric,
          min_delta=early_stopping_min_delta,
          patience=early_stopping_patience,
          verbose=True,
          mode="min",
      )
  
      checkpoint_callback = ModelCheckpoint(
          monitor=monitor_metric,
          filename=monitor_metric + "--{epoch:02d}--{train_loss:.2f}",
          save_top_k=1,
          mode="min",
          save_last=True,
      )
  
      callbacks = [
          ParameterCountCallback(),
          InitialSaveCallback(args_to_save=args_to_save),
          early_stop_callback,
          checkpoint_callback,
      ]
      if dm_name == "dm_super":
          callbacks.append(
              LossPlotLoggerCallback(
                  every_n_epochs=log_plot_every_n_epochs
              )
          )
          if len(kwargs["dm_super_dataset_names_val"].split(",")) > 1:
              callbacks.append(
                  MacroCRPSPlotCallback(
                      every_n_epochs=log_plot_every_n_epochs
                  )
              )
      elif dm_name == "dm_cheat":
          callbacks.append(
              CheatLossPlotLoggerCallback(
                  dataset_names_val=dm.dataset_names_val,
                  every_n_epochs=log_plot_every_n_epochs,
              )
          )
  
      # add callback that only works with attention
      attention_models = ["iwata", "cnn_iwata", "tcn"]
      if model_name in attention_models:
          callbacks.extend(
              [
                  ForecastSupportSetAttentionPlotLoggerCallback(
                      log_batch_train,
                      quantiles=quantiles,
                      split="train",
                      every_n_epochs=log_plot_every_n_epochs,
                  ),
                  ForecastSupportSetAttentionPlotLoggerCallback(
                      log_batch_val,
                      quantiles=quantiles,
                      split="val",
                      every_n_epochs=log_plot_every_n_epochs,
                  ),
              ]
          )
      else:
          callbacks.extend(
              [
                  ForecastPlotLoggerCallback(
                      log_batch_val,
                      quantiles=quantiles,
                      split="val",
                      every_n_epochs=log_plot_every_n_epochs,
                  ),
                  ForecastPlotLoggerCallback(
                      log_batch_train,
                      quantiles=quantiles,
                      split="train",
                      every_n_epochs=log_plot_every_n_epochs,
                  ),
              ]
          )
  
      # -------------------- train model ---------------------------------------------------
      trainer = pl.Trainer(
          gpus=1,
          log_every_n_steps=1,
          limit_train_batches=num_batches,
          val_check_interval=1.0,
          max_epochs=max_epochs,
          max_time=max_training_time,
          logger=pll.CSVLogger(output_dir),
          accumulate_grad_batches=accumulate_grad_batches or None,
          callbacks=callbacks,
          # profiler="pytorch",
      )
  
      trainer.fit(module, dm)
      trainer.test(module, dm, ckpt_path="best")
  
      # -------------------- evaluate model and store results ---------------------------------------------------
      # load weights with best validation loss during training period
      model_best = load_weights(
          model=model,
          path_to_weights=checkpoint_callback.best_model_path,
      )
      if dm_name == "dm_super" and model_name in attention_models:
          save_dir = get_save_dir_from_csvlogger(trainer.logger)
          large_training_evaluation(
              save_dir=save_dir,
              model_name=model_name,
              seed=seed,
              support_length_multiple=support_length_multiple,
              dm_super_cheat=kwargs["dm_super_cheat"],
              datasets=kwargs["dm_super_dataset_names_test"].split(","),
              dm_super_catch22_train=kwargs["dm_super_catch22_train"],
              quantiles=quantiles,
              model=model_best,
              data_dir=Path(data_dir),
          )
  
file_path: gluonts/src/gluonts/nursery/few_shot_prediction/src/scripts/train.py
imports: import warnings
  from typing import cast, Optional
  from pathlib import Path
  import os
  import random
  import click
  import pytorch_lightning
  import pytorch_lightning.loggers
  import torch
  from pytorch_lightning.callbacks.early_stopping import EarlyStopping
  from pytorch_lightning.callbacks import ModelCheckpoint
  from meta.models.module import MetaLightningModule
  from meta.models.common import load_weights
  from meta.models import get_model, MODEL_REGISTRY
  from meta.datasets import get_data_module, DATA_MODULE_REGISTRY
  from meta.metrics import CRPS, QuantileLoss, QuantileWidth
  from meta.callbacks.common import get_save_dir_from_csvlogger
  from meta.callbacks import ForecastPlotLoggerCallback, ParameterCountCallback, InitialSaveCallback, ForecastSupportSetAttentionPlotLoggerCallback, LossPlotLoggerCallback, CheatLossPlotLoggerCallback, MacroCRPSPlotCallback
  from experiments.evaluation.eval_real_cheat import large_training_evaluation
label_desc: Use image node_id 1 for calling a global function with example usage: main(data_dir, output_dir, seed, model_name, quantiles, dm_name, support_set_size, standardize, pretrained, learning_rate, lr_on_plateau_patience, lr_on_plateau_factor, context_length_multiple, support_length_multiple, model_prediction_length, max_epochs, max_training_time, num_batches, batch_size_train, batch_size_val_test, accumulate_grad_batches, early_stopping_min_delta, early_stopping_patience, monitor_metric, n_logging_samples, log_plot_every_n_epochs, num_workers) without return types