src/main.py

# coding=utf-8
"""Fine-tuning RegFM for gene expression prediction."""

import argparse
import logging
import os
from datetime import timedelta

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["NCCL_BLOCKING_WAIT"] = "1"
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1"
os.environ["NCCL_DEBUG"] = "INFO"
os.environ["NCCL_TIMEOUT"] = "1800"

import torch

torch.cuda.empty_cache()
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

from transformers import BertConfig, DNATokenizer
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors

from dataset import load_and_cache_examples
from utils import (
    build_dna_tokenizer,
    build_regfm,
    evaluate,
    load_finetuned_checkpoint,
    predict,
    set_seed,
    train,
    visual_cross,
)


logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser()

    # data / model paths
    parser.add_argument("--tfcr_dir", default=None, type=str, required=True, help="TF chromatin region input data dir.")
    parser.add_argument("--dna_dir", default=None, type=str, required=True, help="DNA input data dir.")
    parser.add_argument("--exp_dir", default=None, type=str, required=True, help="Expression input data dir.")
    parser.add_argument("--cis_model_name_or_path", default=None, type=str, required=True, help="Path to cis-DNA pretrained model.")
    parser.add_argument("--trans_model_name_or_path", default=None, type=str, required=True, help="Path to TF/trans pretrained model.")
    parser.add_argument("--exp_config_name", default="", type=str, required=True, help="Expression config name or path.")
    parser.add_argument("--dna_config_name", default="", type=str, required=True, help="DNA config name or path.")
    parser.add_argument("--trans_tokenizer_name", default="", type=str, required=True, help="TF/trans tokenizer name or path.")
    parser.add_argument("--exp_tokenizer_name", default="", type=str, required=True, help="Expression tokenizer name or path.")
    parser.add_argument("--dna_tokenizer_name", default="", type=str, required=True, help="DNA tokenizer name or path.")
    parser.add_argument("--output_dir", default=None, type=str, required=True, help="Output directory.")
    parser.add_argument(
        "--task_name",
        default=None,
        type=str,
        required=True,
        help="Task name selected in the list: " + ", ".join(processors.keys()),
    )

    # modes
    parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
    parser.add_argument("--do_eval", action="store_true", help="Whether to run evaluation.")
    parser.add_argument("--do_predict", action="store_true", help="Whether to run prediction.")
    parser.add_argument("--do_visualcross", action="store_true", help="Whether to extract cross-attention.")
    parser.add_argument("--evaluate_during_training", action="store_true", help="Evaluate during training.")
    parser.add_argument("--overwrite_output_dir", action="store_true", help="Overwrite the output directory.")

    # sequence / data processing
    parser.add_argument("--max_seq_length", default=128, type=int, help="Max TF/expression sequence length.")
    parser.add_argument("--max_dna_seq_length", default=128, type=int, help="Max DNA sequence length.")
    parser.add_argument("--n_process", default=2, type=int, help="Number of processes used for data processing.")

    # training hyperparameters
    parser.add_argument("--per_gpu_train_batch_size", default=8, type=int, help="Train batch size per GPU.")
    parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int, help="Eval batch size per GPU.")
    parser.add_argument("--per_gpu_pred_batch_size", default=8, type=int, help="Predict batch size per GPU.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
    parser.add_argument("--learning_rate", default=5e-5, type=float, help="Initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay.")
    parser.add_argument("--hidden_dropout_prob", default=0.1, type=float, help="Hidden dropout.")
    parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total training epochs.")
    parser.add_argument("--max_steps", default=-1, type=int, help="Override num_train_epochs if > 0.")
    parser.add_argument("--warmup_percent", default=0, type=float, help="Linear warmup percent of total steps.")
    parser.add_argument("--logging_steps", type=int, default=500, help="Log every X update steps.")
    parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X update steps.")
    parser.add_argument("--save_total_limit", type=int, default=None, help="Max number of checkpoints to keep.")
    parser.add_argument("--early_stop", default=0, type=int, help="Early stop patience (0 disables).")

    # predict
    parser.add_argument("--predict_dir", default=None, type=str, help="Output directory for prediction.")
    parser.add_argument("--save_name", type=str, default="", help="Name for saving prediction result.")

    # distributed
    parser.add_argument("--local-rank", type=int, default=-1, help="Distributed training local rank.")

    args = parser.parse_args()

    if (
        os.path.exists(args.output_dir)
        and os.listdir(args.output_dir)
        and args.do_train
        and not args.overwrite_output_dir
    ):
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                args.output_dir
            )
        )

    if args.local_rank == -1:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl", timeout=timedelta(hours=480))
        args.n_gpu = 1
    args.device = device

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s",
        args.local_rank,
        device,
        args.n_gpu,
        bool(args.local_rank != -1),
    )

    set_seed(args)

    args.task_name = args.task_name.lower()
    if args.task_name not in processors:
        raise ValueError("Task not found: %s" % (args.task_name))
    processor = processors[args.task_name]()
    args.output_mode = output_modes[args.task_name]
    label_list = processor.get_labels()
    num_labels = len(label_list)

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()

    config = BertConfig.from_pretrained(
        args.exp_config_name,
        num_labels=num_labels,
        finetuning_task=args.task_name,
    )
    config.hidden_dropout_prob = args.hidden_dropout_prob
    config.attention_probs_dropout_prob = 0.1

    dna_config = BertConfig.from_pretrained(
        args.dna_config_name,
        num_labels=num_labels,
        finetuning_task=args.task_name,
    )
    dna_config.vocab_size = 261
    dna_config.max_position_embeddings = 512
    dna_config.hidden_dropout_prob = args.hidden_dropout_prob
    dna_config.attention_probs_dropout_prob = 0.1

    tokenizer = DNATokenizer.from_pretrained(args.trans_tokenizer_name)
    epi_tokenizer = DNATokenizer.from_pretrained(args.exp_tokenizer_name)
    dna_tokenizer = build_dna_tokenizer(args)

    model = None
    if args.do_train:
        model = build_regfm(args, config, dna_config)
        print(model)
        print(sum(p.numel() for p in model.parameters() if p.requires_grad))
        logger.info("finish loading model")

    if args.local_rank == 0:
        torch.distributed.barrier()

    if model is not None:
        model.to(args.device)
    logger.info("Training/evaluation parameters %s", args)

    if args.do_train:
        train_dataset = load_and_cache_examples(
            args, args.task_name, tokenizer, epi_tokenizer, dna_tokenizer, 3000, evaluate=False
        )
        global_step, tr_loss = train(args, train_dataset, model, tokenizer, epi_tokenizer, dna_tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        tokenizer.save_pretrained(args.output_dir)
        torch.save(model.state_dict(), args.output_dir + "/model.pth")
        torch.save(model, args.output_dir + "/modelwhole.pth")
        torch.save(args, os.path.join(args.output_dir, "training_args.bin"))
        model.to(args.device)

    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
        logger.info("Evaluate the following checkpoint: %s", args.output_dir)
        model = load_finetuned_checkpoint(args.output_dir, args.device, config=config)
        model.to(args.device)
        results = evaluate(args, model, tokenizer, epi_tokenizer, dna_tokenizer)

    if args.do_predict and args.local_rank in [-1, 0]:
        logger.info("Predict using the following checkpoint: %s", args.output_dir)
        model = load_finetuned_checkpoint(args.output_dir, args.device, config=config)
        model.to(args.device)
        predict(args, model, tokenizer, epi_tokenizer, dna_tokenizer)

    if args.do_visualcross and args.local_rank in [-1, 0]:
        logger.info("Visualcross using the following checkpoint: %s", args.output_dir)
        model = load_finetuned_checkpoint(args.output_dir, args.device, config=config)
        model.to(args.device)
        visual_cross(args, model, tokenizer, epi_tokenizer, dna_tokenizer)

    return results


if __name__ == "__main__":
    main()