code: def train_with_cleanup(
      trainer: ArgillaTrainer, output_dir: str, train: bool = True
  ) -> None:
      try:
          if train:
              trainer.train(output_dir)
          else:
              trainer.save(output_dir)
      finally:
          if Path(output_dir).exists():
              shutil.rmtree(output_dir)
  
imports: import os
  import shutil
  from pathlib import Path
  from argilla.training import ArgillaTrainer
label_desc: Use image node_id 1 for calling a global function with example usage: train_with_cleanup(trainer, output_dir, train) without return types