main.py

import hydra
import torch
import transformers
import os
from omegaconf import OmegaConf

if "TOKENIZERS_PARALLELISM" not in os.environ:
    #NOTE: resolves potential deadlock between tokenizer and data loader (+ warning) at cost of parallelism in tokenizer.
    # Should look at this again.
    # TODO: Clean
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

try:
    import lightning as pl
    from lightning.pytorch.loggers import WandbLogger, TensorBoardLogger
    from lightning.pytorch import seed_everything
except ModuleNotFoundError as e:
    print(e)
    import pytorch_lightning as pl
    from pytorch_lightning.loggers import WandbLogger, TensorBoardLogger
    from pytorch_lightning import seed_everything
from tools.callbacks import ModelCheckpoint, make_progress_bar

# Logger
from tools.logger import logger

import kg_dataset
import datamodules
from tokenizer_tools import load_tokenizer
from kg_transformer_model import KGTransformerModel

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    # Loads a dictionary of type hints describing the config variable
    from tools.type_hints import *

from datetime import timedelta
CHECKPOINT_TIME = timedelta(minutes=58)

@hydra.main(version_base=None, config_path="conf", config_name="config")
def run(config: "JobConfig") -> None:
    logger.debug("Our Logger")
    print(OmegaConf.to_yaml(config))
    hydra_cfg: "HydraConfig" = hydra.core.hydra_config.HydraConfig.get()
    # Process config
    config = process_deprecated(config)

    if config.model.is_encoder_decoder in ("auto", "detect"): # Detect if it's an encoder-decoder model
        archetype = config["model"]["archetype"]
        try: 
            archetype_class = getattr(transformers, archetype)
            archetype_config = archetype_class.config_class 
        except AttributeError: 
            logger.warning("The model you are trying to load is not on huggingface yet! Loading the T5Config manually!")
            from transformers import T5Config
            archetype_config = T5Config   
        config.model.is_encoder_decoder = archetype_config().is_encoder_decoder
        logger.debug("Detected config.model.is_encoder_decoder is %s", config.model.is_encoder_decoder)
        
    if config.disable_logging:
        logger.warning("Logging disabled: No log files will be generated")
    else:
        logger.debug("output written to %s \n", config.output_dir)

    # Load the datamodule class associated with the dataset
    dataset_class: type[kg_dataset.KGCContextDataset] = getattr(kg_dataset, config.dataset["class"])
    datamodule_class = dataset_class.DATAMODULE_NAME
    logger.info(f"Using datamodule {datamodule_class}")
    if isinstance(datamodule_class, str):
        datamodule_class: type[pl.LightningDataModule] = getattr(datamodules, datamodule_class)
    tokenizer = load_tokenizer(config, "conf/special_tokens_t5.csv") 
    dm: "pl.LightningDataModule | datamodules.NbhoodDataModule | datamodules.FewShotNbhModule" = datamodule_class(config=config, tokenizer=tokenizer)

    #reproducibility
    seed_everything(1052024, workers=True) # sets seeds for numpy, torch and python.random. NOTE: before weight initialization
    if config.resume_from:
        model = KGTransformerModel.load_from_checkpoint(
            config.resume_from, config=config, data_module=dm
        )
    else:
        model = KGTransformerModel(config, data_module=dm)  
    
    if config.disable_logging:
        checkpoint_callbacks = []
        train_loggers = False #False disables logging, None sets it to default!.
    else:
        checkpoint_callbacks = []
        if config.checkpoint.train_time_interval is not None and isinstance(config.checkpoint.train_time_interval, int) and config.checkpoint.train_time_interval > 0:
            logger.info("Checkpointing every %d minutes", config.checkpoint.train_time_interval)
            checkpoint_monitor_time = ModelCheckpoint(
                filename="{epoch}-{step}",
                monitor="epoch",
                mode="max",
                save_last=True,
                save_top_k=1,
                train_time_interval=timedelta(minutes=config.checkpoint.train_time_interval),
            )
            checkpoint_monitor_time.CHECKPOINT_NAME_LAST = "last-step_{step}"
            checkpoint_callbacks.append(checkpoint_monitor_time)
        else:
            logger.info("Time interval checkpointing disabled.", config.checkpoint.train_time_interval)
        
        checkpoint_monitor_epoch = ModelCheckpoint(
            filename="{epoch}-{step}",
            monitor="epoch",
            mode="max",
            save_top_k=config.checkpoint.keep_top_k,
            every_n_epochs=config.checkpoint.every_n_epochs,
            save_on_train_epoch_end=True # Save before validation; avoid loss if validation fails.
        )
        checkpoint_monitor_epoch.CHECKPOINT_NAME_LAST = "last-epoch_{epoch}-step_{step}"
        checkpoint_callbacks.append(checkpoint_monitor_epoch)
        checkpoint_callbacks.append(make_progress_bar(refresh_rate=10))
        
        # Loggers :

        train_loggers = []
        if config.tensorboard.use:
            try:
                tensorboard_logger = TensorBoardLogger(save_dir=config.output_dir, version=config.tensorboard.version, name=config.tensorboard.name)
            except ModuleNotFoundError:
                logger.error("config.tensorboard.use is True but tensorboard or tensorboardX not found. Tensorboard logging disabled.", exc_info=True)
            else:
                train_loggers.append(tensorboard_logger)
        
        if config.wandb.use and not config.disable_logging:
            wandb_logger = WandbLogger(
                name=config.wandb.run_name,
                project=config.wandb.project_name,
                config=config,
                save_dir=config.output_dir
            )
            wandb_logger.experiment.config.update(config)
            train_loggers.append(wandb_logger)
        if len(train_loggers) == 0:
            logger.warning(" --- This job runs without any loggers (e.g. no tensorboard or wandb) ---")

    train_options = dict(config.train)
    # remove non pl.trainer options:
    del train_options["num_workers"]
    del train_options["batch_size"]
    del train_options["drop_subject"]
    
    if train_options["max_steps"] == "one-per-triple":
        import os
        dataset_file = os.path.join('data', config.dataset.name, "train.del")
        with open(dataset_file, "r") as f:
            num_triples = len(f.readlines())
        train_options["max_steps"] = num_triples
        print("[AUTO OPTION] Setting max_steps to", num_triples)
    
    val_check_interval = config.valid.get("val_check_interval", None)
    if val_check_interval is not None: 
        if isinstance(val_check_interval, str) and val_check_interval.lower() == "none": # NOTE: should write null
            config.valid.val_check_interval = val_check_interval = None
            logger.info("Setting 'val_check_interval' to None; write 'null' in config file instead of 'None' to avoid this message.")

    train_options = {
        'limit_train_batches' : config.train.get("limit_train_batches", 1.0),
        **train_options, # <-- Unpack all options
        
        #'accelerator': config.train.accelerator,
        #'devices': config.train.devices,
        #'strategy': config.train.strategy,
        #'precision': config.train.precision,
        
        # Adding non config.train options:

        'callbacks': checkpoint_callbacks,
        'default_root_dir': config.output_dir,
        'check_val_every_n_epoch': config.valid.every_n_epochs if config.valid.every_n_epochs != 0 else None,
        'val_check_interval' : val_check_interval,
        'limit_val_batches' : config.valid["limit_val_batches"],
        'num_sanity_val_steps' : config.valid.get('num_sanity_val_steps', 2),
        'logger' : train_loggers,
        'deterministic': True,
    }
    # default logger used by trainer (if tensorboard is installed)
   

    if train_options["max_steps"] != -1: # only profile on small runs
        try:
            from lightning.pytorch.profilers import AdvancedProfiler, SimpleProfiler, PyTorchProfiler
        except ModuleNotFoundError:
            from pytorch_lightning.profilers import AdvancedProfiler, SimpleProfiler, PyTorchProfiler
        #profiler = SimpleProfiler(dirpath="outputs", filename="perf_logs_simple.log")
        profiler = PyTorchProfiler(dirpath="outputs", filename="perf_logs_advanced.log")
    else:
        profiler = None
    
    seed_everything(1052024, workers=True) # Once again for trainer, in case of different models.
    trainer = pl.Trainer(**train_options, profiler=profiler,)
    if trainer.fast_dev_run:
        # NOTE: This fixes a incompatibility with Adafactor, transformers and lightning in the recent versions.
        # Related Issue: https://github.com/Daraan/KnowledgeGraphTransformer/issues/4 
        trainer.num_sanity_val_steps = 2
    if len(config.resume_from) != 0:
        trainer.fit(model, dm, ckpt_path=config.resume_from)  # , ckpt_path=ckpt_path)
    else:
        trainer.fit(model, dm)  # , ckpt_path=ckpt_path)

    if config.eval.execute_after_training and not (trainer.fast_dev_run or trainer.interrupted):
        import eval
        if trainer.num_devices == 1:
            # NOTE: Should only be done if only 1 device is used
            if not config.disable_logging:
                eval_loggers = eval.make_eval_loggers(config, split="test")
                trainer.loggers = eval_loggers
            trainer.test(model, datamodule=dm)
        elif len(checkpoint_callbacks) > 0 and any(isinstance(cb, ModelCheckpoint) for cb in checkpoint_callbacks):
            for i in range(1, len(checkpoint_callbacks)+1):
                if isinstance(checkpoint_callbacks[-i], ModelCheckpoint):
                    break
            checkpoint_path = checkpoint_callbacks[-i].best_model_path
            eval.run(checkpoint_path, None, split="test")
        else:
            logger.warning("Cannot run evaluation cleanly on a single device as no checkpoint was saved. This might lead to errors in the evaluation.")
            eval.eval_model(model, None, split="test")


def process_deprecated(config):
    # Legacy code
    if hasattr(config, "use_neighborhood"):
        config.context.use = config.use_neighborhood
        del config.use_neighborhood
    if hasattr(config, "use_wandb"):
        config.wandb.use = config.use_wandb
        del config.use_wandb
    if not hasattr(config, "descriptions"):
        config.descriptions = {"use": False}
    
    # Encoder decoder check
    if config.model.is_encoder_decoder in ("auto", "detect"): # Detect if it's an encoder-decoder model
        archetype : str = config["model"]["archetype"]
        try: 
            archetype_class: type[transformers.PreTrainedModel] = getattr(transformers, archetype)
            archetype_config: type[transformers.PretrainedConfig] = archetype_class.config_class 
        except AttributeError: 
            from transformers import T5Config
            archetype_config = T5Config
        config.model.is_encoder_decoder = archetype_config().is_encoder_decoder
        logger.info("Detected config.model.is_encoder_decoder is %s", config.model.is_encoder_decoder)
    elif not config.model.is_encoder_decoder in (True, False):
        logger.warning("config.model.is_encoder_decoder has a unsupported value, expected boolean or 'auto'.")
    return config


if __name__ == '__main__':
    torch.set_float32_matmul_precision('medium')
    import os
    #os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
    #os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    os.environ['TRANSFORMERS_OFFLINE']="1" #for when hugginface is down
    os.environ['HF_DATASETS_OFFLINE']="1"
    run()