slide_generation/src/train_text_summarizer.py

try:
    import wandb
    WANDB_FLAG = True
except ImportError:
    WANDB_FLAG = False

import os

import jax
import hydra
from omegaconf import DictConfig, OmegaConf

from transformers import BartTokenizer
import tensorflow as tf

from schema import register_configs
from utils.for_training import TextSummarizerTrainer, create_text_summarizer_dataset

import logging
logger = logging.getLogger(__name__)


# Register data classes that manage all the parameters dealt with by hydra config to confirm their type annotation.
register_configs()

@hydra.main(config_path='config', config_name='text_summarization.yaml', version_base="1.2")
def main(cfg: DictConfig):
    ## 0. Experiment Management
    # Setup weight & biases.
    if WANDB_FLAG:
        config = OmegaConf.to_container(
            cfg, resolve=True, throw_on_missing=True
        )
        tag_list = [
            "bs_{}".format(cfg.training.batch_size),
            "tstep_{}".format(cfg.training.num_train_steps),
            "opt_{}".format(cfg.opt.mode),
            "lr_{}".format(cfg.opt.peak_value if cfg.opt.mode == "scheduler" else cfg.opt.lr),
        ]
        if cfg.opt.mode == "scheduler":
            tag_list.append("wstep_{}".format(cfg.opt.warmup_steps))
        wandb.init(
            config=config,
            project="slide_generation",
            # name="",
            group="bart_debug",
            tags=tag_list,
            reinit=True,
            settings=wandb.Settings(start_method="thread"),
            notes="" # leave comments if you want to log more detailed messages
        )
    # Hide any GPUs form TensorFlow. Otherwise TF might reserve memory and make
    # it unavailable to JAX.
    tf.config.experimental.set_visible_devices([], "GPU")
    # Log basic information,
    logger.info("Jax local devices: {}".format(jax.local_devices()))

    ## 1. Training Preparation
    # Get the path to the result directory and the project root.
    result_dir = os.getcwd()
    project_root = hydra.utils.get_original_cwd()

    # Prepare each group of configuration.
    training_cfg = cfg.training
    layout_model_cfg = cfg.model
    dataset_cfg = cfg.dataset
    optax_cfg = cfg.opt

    n_devices = jax.local_device_count()
    assert cfg.training.batch_size % n_devices == 0

    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-xsum")

    dataset_dir_path = os.path.join(project_root, dataset_cfg.dataset_dir)
    train_dataset, val_dataset = create_text_summarizer_dataset(
        dataset_dir_path=dataset_dir_path,
        tokenizer=tokenizer,
        batch_size=cfg.training.batch_size,
        max_length=cfg.training.max_seq_length,
        layout_resolution_w=dataset_cfg.resolution_w,
        layout_resolution_h=dataset_cfg.resolution_h,
    )
    # How many epochs do we need (if count is None, that means the dataset will be repeated indefinitely)
    train_dataset = train_dataset.repeat(count=None)
    val_dataset = val_dataset.repeat(count=None)
    logger.info("Finished dataset creation.")

    rng = jax.random.PRNGKey(cfg.seed)
    train_rng, param_rng = jax.random.split(rng)

    # Make a trainer class instance.
    trainer = TextSummarizerTrainer.create_trainer(
        rng=param_rng,
        training_cfg=training_cfg,
        model_cfg=layout_model_cfg,
        dataset_cfg=dataset_cfg,
        optax_cfg=optax_cfg,
        save_folder=result_dir,
    )

    trainer.train(
        rng=train_rng,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
    )


if __name__ == '__main__':
    main() # pylint: disable=no-value-for-parameter