Function train

tasks/AutoChem/code/experiment.py:177–468

(args)

def train(args):
    os.environ["WANDB_SILENT"] = "true"  
    os.environ["WANDB_ERROR_REPORTING"] = "false"  
    os.makedirs(args.out_dir, exist_ok=True)
    if args.local_rank == 0:
        tb_log_dir =os.getenv("TENSORBOARD_LOG_PATH", "/tensorboard_logs/")
        #tb_log_dir = os.path.join(args.out_dir, "logs")  
        writer = SummaryWriter(log_dir=tb_log_dir)
        #wandb.init( 
            #project="chemical_predict", 
            #name=datetime.datetime.now().strftime('%Y-%m-%d--%H:%M')+'-'+args.data_name
            #)

    # Load the model and tokenizer
    pretrained_model_path = args.pretrained_model_path
    num_epoch = args.num_epoch
    batch_size= args.per_device_train_batch_size
    yield_predictor_path = args.yield_predictor_path
    lr=args.lr
    max_length = args.max_length

    data_path=args.data_path
    data_name = args.data_name
    # Save the base model
    lora_adapter_path = args.lora_adapter_path

    load_ds_dir = args.load_ds_dir
    load_ds_ckpt_id = args.load_ds_ckpt_id

    use_lora = args.use_lora

    log_path = os.path.join(args.out_dir, args.log_file)

    logging.basicConfig(
        filename=log_path,
        level=logging.INFO,
    )
    logger = logging.getLogger()

    # args.global_rank = torch.distributed.get_rank()
    get_accelerator().set_device(args.local_rank)
    device = torch.device(get_accelerator().device_name(), args.local_rank)
    # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
    # torch.distributed.init_process_group(backend='nccl')
    deepspeed.init_distributed()

    print('using device', device)

    print('Load model...')
    logger.info('Load model...')


    if use_lora:

        # Define LoRA configuration
        if not os.path.exists(lora_adapter_path):
            model = AutoModel.from_pretrained(pretrained_model_path)
            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_path)