code: def cli_main():
      parser = ArgumentParser()
      parser.add_argument("--batch-size", default=6, type=int)
      parser.add_argument(
          "--dataset",
          default="librimix",
          type=str,
          choices=["wsj0mix", "librimix"],
      )
      parser.add_argument(
          "--root-dir",
          type=Path,
          help="The path to the directory where the directory ``Libri2Mix`` or ``Libri3Mix`` is stored.",
      )
      parser.add_argument(
          "--librimix-tr-split",
          default="train-360",
          choices=["train-360", "train-100"],
          help="The training partition of librimix dataset. (default: ``train-360``)",
      )
      parser.add_argument(
          "--librimix-task",
          default="sep_clean",
          type=str,
          choices=["sep_clean", "sep_noisy", "enh_single", "enh_both"],
          help="The task to perform (separation or enhancement, noisy or clean). (default: ``sep_clean``)",
      )
      parser.add_argument(
          "--num-speakers",
          default=2,
          type=int,
          help="The number of speakers in the mixture. (default: 2)",
      )
      parser.add_argument(
          "--sample-rate",
          default=8000,
          type=int,
          help="Sample rate of audio files in the given dataset. (default: 8000)",
      )
      parser.add_argument(
          "--exp-dir",
          default=Path("./exp"),
          type=Path,
          help="The directory to save checkpoints and logs.",
      )
      parser.add_argument(
          "--epochs",
          metavar="NUM_EPOCHS",
          default=200,
          type=int,
          help="The number of epochs to train. (default: 200)",
      )
      parser.add_argument(
          "--learning-rate",
          default=1e-3,
          type=float,
          help="Initial learning rate. (default: 1e-3)",
      )
      parser.add_argument(
          "--num-gpu",
          default=1,
          type=int,
          help="The number of GPUs for training. (default: 1)",
      )
      parser.add_argument(
          "--num-node",
          default=1,
          type=int,
          help="The number of nodes in the cluster for training. (default: 1)",
      )
      parser.add_argument(
          "--num-workers",
          default=4,
          type=int,
          help="The number of workers for dataloader. (default: 4)",
      )
  
      args = parser.parse_args()
  
      model = _get_model(num_sources=args.num_speakers)
  
      optimizer = torch.optim.Adam(
          model.parameters(), lr=args.learning_rate
      )
      lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
          optimizer, mode="min", factor=0.5, patience=5
      )
      train_loader, valid_loader, eval_loader = _get_dataloader(
          args.dataset,
          args.root_dir,
          args.num_speakers,
          args.sample_rate,
          args.batch_size,
          args.num_workers,
          args.librimix_task,
          args.librimix_tr_split,
      )
      loss = si_sdr_loss
      metric_dict = {
          "sdri": sdri_metric,
          "sisdri": sisdri_metric,
      }
      model = ConvTasNetModule(
          model=model,
          train_loader=train_loader,
          val_loader=valid_loader,
          loss=loss,
          optim=optimizer,
          metrics=metric_dict,
          lr_scheduler=lr_scheduler,
      )
      checkpoint_dir = args.exp_dir / "checkpoints"
      checkpoint = ModelCheckpoint(
          checkpoint_dir,
          monitor="Losses/val_loss",
          mode="min",
          save_top_k=5,
          save_weights_only=True,
          verbose=True,
      )
      callbacks = [
          checkpoint,
          EarlyStopping(
              monitor="Losses/val_loss",
              mode="min",
              patience=30,
              verbose=True,
          ),
      ]
      trainer = Trainer(
          default_root_dir=args.exp_dir,
          max_epochs=args.epochs,
          num_nodes=args.num_node,
          accelerator="gpu",
          strategy="ddp_find_unused_parameters_false",
          devices=args.num_gpu,
          limit_train_batches=1.0,  # Useful for fast experiment
          gradient_clip_val=5.0,
          callbacks=callbacks,
      )
      trainer.fit(model)
      model.load_from_checkpoint(checkpoint.best_model_path)
      state_dict = torch.load(
          checkpoint.best_model_path, map_location="cpu"
      )
      state_dict = {
          k.replace("model.", ""): v
          for k, v in state_dict["state_dict"].items()
      }
      torch.save(state_dict, args.exp_dir / "best_model.pth")
      trainer.test(model, eval_loader)
  
imports: from argparse import ArgumentParser
  from pathlib import Path
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, TypedDict, Union
  import torch
  import torchaudio
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
  from torch import nn
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from utils import metrics
  from utils.dataset import utils
label_desc: Use image node_id 6 for calling a global function with example usage: cli_main() without return types