import os


def init_wandb_training(training_args):
    """
    Helper function for setting up Weights & Biases logging tools.
    """
    if training_args.output_dir is not None:
        os.makedirs(os.path.join(training_args.output_dir, "wandb"), exist_ok=True)
        project_name = os.path.basename(training_args.output_dir)
        os.environ["WANDB_DIR"] = os.path.join(training_args.output_dir, "wandb")
        os.environ["WANDB_PROJECT"] = project_name
        print(f"Creating wandb directory at {os.path.join(training_args.output_dir, 'wandb')}, project name: {project_name}")
    if training_args.wandb_entity is not None:
        os.environ["WANDB_ENTITY"] = training_args.wandb_entity
    if training_args.wandb_project is not None:
        os.environ["WANDB_PROJECT"] = training_args.wandb_project
