relik/retriever/trainer/train.py

from copy import deepcopy
import os
from pathlib import Path
from typing import List, Literal, Optional, Union

import hydra
import lightning as pl
import omegaconf
import torch
from lightning import Trainer
from lightning.pytorch.callbacks import (
    EarlyStopping,
    LearningRateMonitor,
    ModelCheckpoint,
    ModelSummary,
)
from lightning.pytorch.loggers import WandbLogger
from omegaconf import OmegaConf
from pprintpp import pformat

from relik.common.log import get_logger
from relik.retriever.callbacks.base import NLPTemplateCallback
from relik.retriever.callbacks.evaluation_callbacks import (
    AvgRankingEvaluationCallback,
    RecallAtKEvaluationCallback,
)
from relik.retriever.callbacks.prediction_callbacks import (
    GoldenRetrieverPredictionCallback,
)
from relik.retriever.callbacks.training_callbacks import NegativeAugmentationCallback
from relik.retriever.callbacks.utils_callbacks import (
    FreeUpIndexerVRAMCallback,
    SavePredictionsCallback,
    SaveRetrieverCallback,
)
from relik.retriever.data.datasets import GoldenRetrieverDataset
from relik.retriever.indexers.base import BaseDocumentIndex
from relik.retriever.lightning_modules.pl_data_modules import (
    GoldenRetrieverPLDataModule,
)
from relik.retriever.lightning_modules.pl_modules import GoldenRetrieverPLModule
from relik.retriever.pytorch_modules.loss import MultiLabelNCELoss
from relik.retriever.pytorch_modules.model import GoldenRetriever
from relik.retriever.pytorch_modules.optim import RAdamW
from relik.retriever.pytorch_modules.scheduler import LinearScheduler

logger = get_logger(__name__)


class RetrieverTrainer:
    def __init__(
        self,
        retriever: GoldenRetriever,
        train_dataset: GoldenRetrieverDataset | None = None,
        val_dataset: GoldenRetrieverDataset
        | list[GoldenRetrieverDataset]
        | None = None,
        test_dataset: GoldenRetrieverDataset
        | list[GoldenRetrieverDataset]
        | None = None,
        num_workers: int = 4,
        optimizer: torch.optim.Optimizer = RAdamW,
        lr: float = 1e-5,
        weight_decay: float = 0.01,
        lr_scheduler: torch.optim.lr_scheduler.LRScheduler = LinearScheduler,
        num_warmup_steps: int = 0,
        loss: torch.nn.Module = MultiLabelNCELoss,
        callbacks: list | None = None,
        accelerator: str = "auto",
        devices: int = 1,
        num_nodes: int = 1,
        strategy: str = "auto",
        accumulate_grad_batches: int = 1,
        gradient_clip_val: float = 1.0,
        val_check_interval: float = 1.0,
        check_val_every_n_epoch: int = 1,
        max_steps: int | None = None,
        max_epochs: int | None = None,
        deterministic: bool = True,
        fast_dev_run: bool = False,
        precision: int | str = 16,
        reload_dataloaders_every_n_epochs: int = 1,
        resume_from_checkpoint_path: str | os.PathLike | None = None,
        trainer_kwargs: dict | None = None,
        # eval parameters
        metric_to_monitor: str = "validate_recall@{top_k}",
        monitor_mode: str = "max",
        top_k: int | List[int] = 100,
        # early stopping parameters
        early_stopping: bool = True,
        early_stopping_patience: int = 10,
        early_stopping_kwargs: dict | None = None,
        # wandb logger parameters
        log_to_wandb: bool = True,
        wandb_entity: str | None = None,
        wandb_experiment_name: str | None = None,
        wandb_project_name: str = "golden-retriever",
        wandb_save_dir: str | os.PathLike = "./",  # TODO: i don't like this default
        wandb_log_model: bool = True,
        wandb_online_mode: bool = False,
        wandb_watch: str = "all",
        wandb_kwargs: dict | None = None,
        # checkpoint parameters
        model_checkpointing: bool = True,
        checkpoint_dir: str | os.PathLike | None = None,
        checkpoint_filename: str | os.PathLike | None = None,
        save_top_k: int = 1,
        save_last: bool = False,
        checkpoint_kwargs: dict | None = None,
        # prediction callback parameters
        prediction_batch_size: int = 128,
        # hard negatives callback parameters
        max_hard_negatives_to_mine: int = 15,
        hard_negatives_threshold: float = 0.0,
        metrics_to_monitor_for_hard_negatives: str | None = None,
        mine_hard_negatives_with_probability: float = 1.0,
        # other parameters
        seed: int = 42,
        float32_matmul_precision: str = "medium",
        **kwargs,
    ):
        # put all the parameters in the class
        self.retriever = retriever
        # datasets
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.test_dataset = test_dataset
        self.num_workers = num_workers
        # trainer parameters
        self.optimizer = optimizer
        self.lr = lr
        self.weight_decay = weight_decay
        self.lr_scheduler = lr_scheduler
        self.num_warmup_steps = num_warmup_steps
        self.loss = loss
        self.callbacks = callbacks
        self.accelerator = accelerator
        self.devices = devices
        self.num_nodes = num_nodes
        self.strategy = strategy
        self.accumulate_grad_batches = accumulate_grad_batches
        self.gradient_clip_val = gradient_clip_val
        self.val_check_interval = val_check_interval
        self.check_val_every_n_epoch = check_val_every_n_epoch
        self.max_steps = max_steps
        self.max_epochs = max_epochs
        self.deterministic = deterministic
        self.fast_dev_run = fast_dev_run
        self.precision = precision
        self.reload_dataloaders_every_n_epochs = reload_dataloaders_every_n_epochs
        self.resume_from_checkpoint_path = resume_from_checkpoint_path
        self.trainer_kwargs = trainer_kwargs or {}
        # eval parameters
        self.metric_to_monitor = metric_to_monitor
        self.monitor_mode = monitor_mode
        self.top_k = top_k
        # early stopping parameters
        self.early_stopping = early_stopping
        self.early_stopping_patience = early_stopping_patience
        self.early_stopping_kwargs = early_stopping_kwargs
        # wandb logger parameters
        self.log_to_wandb = log_to_wandb
        self.wandb_entity = wandb_entity
        self.wandb_experiment_name = wandb_experiment_name
        self.wandb_project_name = wandb_project_name
        self.wandb_save_dir = wandb_save_dir
        self.wandb_log_model = wandb_log_model
        self.wandb_online_mode = wandb_online_mode
        self.wandb_watch = wandb_watch
        self.wandb_kwargs = wandb_kwargs
        # checkpoint parameters
        self.model_checkpointing = model_checkpointing
        self.checkpoint_dir = checkpoint_dir
        self.checkpoint_filename = checkpoint_filename
        self.save_top_k = save_top_k
        self.save_last = save_last
        self.checkpoint_kwargs = checkpoint_kwargs
        # prediction callback parameters
        self.prediction_batch_size = prediction_batch_size
        # hard negatives callback parameters
        self.max_hard_negatives_to_mine = max_hard_negatives_to_mine
        self.hard_negatives_threshold = hard_negatives_threshold
        self.metrics_to_monitor_for_hard_negatives = (
            metrics_to_monitor_for_hard_negatives
        )
        self.mine_hard_negatives_with_probability = mine_hard_negatives_with_probability
        # other parameters
        self.seed = seed
        self.float32_matmul_precision = float32_matmul_precision

        if self.max_epochs is None and self.max_steps is None:
            raise ValueError(
                "Either `max_epochs` or `max_steps` should be specified in the trainer configuration"
            )

        if self.max_epochs is not None and self.max_steps is not None:
            logger.info(
                "Both `max_epochs` and `max_steps` are specified in the trainer configuration. "
                "Will use `max_epochs` for the number of training steps"
            )
            self.max_steps = None

        # reproducibility
        pl.seed_everything(self.seed)
        # set the precision of matmul operations
        torch.set_float32_matmul_precision(self.float32_matmul_precision)

        # lightning data module declaration
        self.lightning_datamodule = self.configure_lightning_datamodule()

        if self.max_epochs is not None:
            logger.info(f"Number of training epochs: {self.max_epochs}")
            self.max_steps = (
                len(self.lightning_datamodule.train_dataloader()) * self.max_epochs
            )

        # optimizer declaration
        self.optimizer, self.lr_scheduler = self.configure_optimizers()

        # lightning module declaration
        self.lightning_module = self.configure_lightning_module()

        # logger and experiment declaration
        # update self.wandb_kwargs
        wandb_args = dict(
            entity=self.wandb_entity,
            project=self.wandb_project_name,
            name=self.wandb_experiment_name,
            save_dir=self.wandb_save_dir,
            log_model=self.wandb_log_model,
            offline=not self.wandb_online_mode,
            watch=self.wandb_watch,
            lightning_module=self.lightning_module,
        )
        if self.wandb_kwargs is not None:
            wandb_args.update(self.wandb_kwargs)
        self.wandb_kwargs = wandb_args
        self.wandb_logger: Optional[WandbLogger] = None
        self.experiment_path: Optional[Path] = None

        # setup metrics to monitor for a bunch of callbacks
        if isinstance(self.top_k, int):
            self.top_k = [self.top_k]
        # save the target top_k
        self.target_top_k = self.top_k[0]
        self.metric_to_monitor = self.metric_to_monitor.format(top_k=self.target_top_k)

        # explicitly configure some callbacks that will be needed not only by the
        # pl.Trainer but also in this class
        # model checkpoint callback
        if self.save_last:
            logger.warning(
                "We will override the `save_last` of `ModelCheckpoint` to `False`. "
                "Instead, we will use a separate `ModelCheckpoint` callback to save the last checkpoint"
            )
        checkpoint_kwargs = dict(
            monitor=self.metric_to_monitor,
            mode=self.monitor_mode,
            verbose=True,
            save_top_k=self.save_top_k,
            filename=self.checkpoint_filename,
            dirpath=self.checkpoint_dir,
            auto_insert_metric_name=False,
        )
        if self.checkpoint_kwargs is not None:
            checkpoint_kwargs.update(self.checkpoint_kwargs)
        self.checkpoint_kwargs = checkpoint_kwargs
        self.model_checkpoint_callback: ModelCheckpoint | None = None
        self.checkpoint_path: str | os.PathLike | None = None
        # last checkpoint callback
        self.latest_model_checkpoint_callback: ModelCheckpoint | None = None
        self.last_checkpoint_kwargs: dict | None = None
        if self.save_last:
            last_checkpoint_kwargs = deepcopy(self.checkpoint_kwargs)
            last_checkpoint_kwargs["save_top_k"] = 1
            last_checkpoint_kwargs["filename"] = "last-{epoch}-{step}"
            last_checkpoint_kwargs["monitor"] = "step"
            last_checkpoint_kwargs["mode"] = "max"
            self.last_checkpoint_kwargs = last_checkpoint_kwargs

        # early stopping callback
        early_stopping_kwargs = dict(
            monitor=self.metric_to_monitor,
            mode=self.monitor_mode,
            patience=self.early_stopping_patience,
        )
        if self.early_stopping_kwargs is not None:
            early_stopping_kwargs.update(self.early_stopping_kwargs)
        self.early_stopping_kwargs = early_stopping_kwargs
        self.early_stopping_callback: EarlyStopping | None = None

        # other callbacks declaration
        self.callbacks_store: List[pl.Callback] = []  # self.configure_callbacks()
        # add default callbacks
        self.callbacks_store += [
            ModelSummary(max_depth=2),
            LearningRateMonitor(logging_interval="step"),
        ]

        # lazy trainer declaration
        self.trainer: pl.Trainer | None = None

    def configure_lightning_datamodule(self, *args, **kwargs):
        # lightning data module declaration
        if self.val_dataset is not None and isinstance(
            self.val_dataset, GoldenRetrieverDataset
        ):
            self.val_dataset = [self.val_dataset]
        if self.test_dataset is not None and isinstance(
            self.test_dataset, GoldenRetrieverDataset
        ):
            self.test_dataset = [self.test_dataset]

        self.lightning_datamodule = GoldenRetrieverPLDataModule(
            train_dataset=self.train_dataset,
            val_datasets=self.val_dataset,
            test_datasets=self.test_dataset,
            num_workers=self.num_workers,
            *args,
            **kwargs,
        )
        return self.lightning_datamodule

    def configure_lightning_module(self, *args, **kwargs):
        # add loss object to the retriever
        if self.retriever.loss_type is None:
            self.retriever.loss_type = self.loss()

        # lightning module declaration
        self.lightning_module = GoldenRetrieverPLModule(
            model=self.retriever,
            optimizer=self.optimizer,
            lr_scheduler=self.lr_scheduler,
            *args,
            **kwargs,
        )

        return self.lightning_module

    def configure_optimizers(self, *args, **kwargs):
        # check if it is the class or the instance
        if isinstance(self.optimizer, type):
            param_optimizer = list(self.retriever.named_parameters())
            no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in param_optimizer if "layer_norm_layer" in n
                    ],
                    "weight_decay": self.weight_decay,
                    "lr": 1e-4,
                },
                {
                    "params": [
                        p
                        for n, p in param_optimizer
                        if all(nd not in n for nd in no_decay)
                        and "layer_norm_layer" not in n
                    ],
                    "weight_decay": self.weight_decay,
                },
                {
                    "params": [
                        p
                        for n, p in param_optimizer
                        if "layer_norm_layer" not in n
                        and any(nd in n for nd in no_decay)
                    ],
                    "weight_decay": 0.0,
                },
            ]
            self.optimizer = self.optimizer(
                # params=self.retriever.parameters(),
                params=optimizer_grouped_parameters,
                lr=self.lr,
                # weight_decay=self.weight_decay,
            )
        else:
            self.optimizer = self.optimizer

        # LR Scheduler declaration
        # check if it is the class, the instance or a function
        if self.lr_scheduler is not None:
            if isinstance(self.lr_scheduler, type):
                self.lr_scheduler = self.lr_scheduler(
                    optimizer=self.optimizer,
                    num_warmup_steps=self.num_warmup_steps,
                    num_training_steps=self.max_steps,
                )

        return self.optimizer, self.lr_scheduler

    @staticmethod
    def configure_logger(
        name: str,
        save_dir: str | os.PathLike,
        offline: bool,
        entity: str,
        project: str,
        log_model: Literal["all"] | bool,
        watch: str | None = None,
        lightning_module: torch.nn.Module | None = None,
        *args,
        **kwargs,
    ) -> WandbLogger:
        """
        Configure the wandb logger

        Args:
            name (`str`):
                The name of the experiment
            save_dir (`str`, `os.PathLike`):
                The directory where to save the experiment
            offline (`bool`):
                Whether to run wandb offline
            entity (`str`):
                The wandb entity
            project (`str`):
                The wandb project name
            log_model (`Literal["all"]`, `bool`):
                Whether to log the model to wandb
            watch (`str`, optional, defaults to `None`):
                The mode to watch the model
            lightning_module (`torch.nn.Module`, optional, defaults to `None`):
                The lightning module to watch
            *args:
                Additional args
            **kwargs:
                Additional kwargs

        Returns:
            `lightning.loggers.WandbLogger`:
                The wandb logger
        """
        wandb_logger = WandbLogger(
            name=name,
            save_dir=save_dir,
            offline=offline,
            project=project,
            log_model=log_model and not offline,
            entity=entity,
            *args,
            **kwargs,
        )
        if watch is not None and lightning_module is not None:
            watch_kwargs = dict(model=lightning_module)
            if watch is not None:
                watch_kwargs["log"] = watch
            wandb_logger.watch(**watch_kwargs)
        return wandb_logger

    @staticmethod
    def configure_early_stopping(
        monitor: str,
        mode: str,
        patience: int = 3,
        *args,
        **kwargs,
    ) -> EarlyStopping:
        logger.info(f"Enabling EarlyStopping callback with patience: {patience}")
        early_stopping_callback = EarlyStopping(
            monitor=monitor,
            mode=mode,
            patience=patience,
            *args,
            **kwargs,
        )
        return early_stopping_callback

    def configure_model_checkpoint(
        self,
        monitor: str,
        mode: str,
        verbose: bool = True,
        save_top_k: int = 1,
        save_last: bool = False,
        filename: str | os.PathLike | None = None,
        dirpath: str | os.PathLike | None = None,
        auto_insert_metric_name: bool = False,
        *args,
        **kwargs,
    ) -> ModelCheckpoint:
        logger.info("Enabling Model Checkpointing")
        if dirpath is None:
            dirpath = (
                self.experiment_path / "checkpoints" if self.experiment_path else None
            )
        if filename is None:
            filename = (
                "checkpoint-" + monitor + "_{" + monitor + ":.4f}-epoch_{epoch:02d}"
            )
        self.checkpoint_path = dirpath / filename if dirpath is not None else None
        logger.info(f"Checkpoint directory: {dirpath}")
        logger.info(f"Checkpoint filename: {filename}")

        kwargs = dict(
            monitor=monitor,
            mode=mode,
            verbose=verbose,
            save_top_k=save_top_k,
            save_last=save_last,
            filename=filename,
            dirpath=dirpath,
            auto_insert_metric_name=auto_insert_metric_name,
            *args,
            **kwargs,
        )

        # update the kwargs
        # TODO: this is bad
        # kwargs.update(
        #     dirpath=self.checkpoint_dir,
        #     filename=self.checkpoint_filename,
        # )
        # modelcheckpoint_kwargs = dict(
        #     dirpath=self.checkpoint_dir,
        #     filename=self.checkpoint_filename,
        # )
        # modelcheckpoint_kwargs.update(kwargs)
        self.model_checkpoint_callback = ModelCheckpoint(**kwargs)
        return self.model_checkpoint_callback

    def configure_hard_negatives_callback(self):
        metrics_to_monitor = (
            self.metrics_to_monitor_for_hard_negatives or self.metric_to_monitor
        )
        hard_negatives_callback = NegativeAugmentationCallback(
            k=self.target_top_k,
            batch_size=self.prediction_batch_size,
            precision=self.precision,
            stages=["validate"],
            metrics_to_monitor=metrics_to_monitor,
            threshold=self.hard_negatives_threshold,
            max_negatives=self.max_hard_negatives_to_mine,
            add_with_probability=self.mine_hard_negatives_with_probability,
            refresh_every_n_epochs=1,
        )
        return hard_negatives_callback

    def training_callbacks(self):
        if self.model_checkpointing:
            self.model_checkpoint_callback = self.configure_model_checkpoint(
                **self.checkpoint_kwargs
            )
            self.callbacks_store.append(self.model_checkpoint_callback)
            if self.save_last:
                self.latest_model_checkpoint_callback = self.configure_model_checkpoint(
                    **self.last_checkpoint_kwargs
                )
                self.callbacks_store.append(self.latest_model_checkpoint_callback)

            self.callbacks_store.append(SaveRetrieverCallback())
        if self.early_stopping:
            self.early_stopping_callback = self.configure_early_stopping(
                **self.early_stopping_kwargs
            )
        return self.callbacks_store

    def configure_metrics_callbacks(
        self, save_predictions: bool = False
    ) -> List[NLPTemplateCallback]:
        """
        Configure the metrics callbacks for the trainer. This method is called
        by the `eval_callbacks` method, and it is used to configure the callbacks
        that will be used to evaluate the model during training.

        Args:
            save_predictions (`bool`, optional, defaults to `False`):
                Whether to save the predictions to disk or not

        Returns:
            `List[NLPTemplateCallback]`:
                The list of callbacks to use for evaluation
        """
        # prediction callback
        metrics_callbacks: List[NLPTemplateCallback] = [
            RecallAtKEvaluationCallback(k, verbose=True) for k in self.top_k
        ]
        metrics_callbacks += [
            AvgRankingEvaluationCallback(k, verbose=True) for k in self.top_k
        ]
        if save_predictions:
            metrics_callbacks.append(SavePredictionsCallback())
        return metrics_callbacks

    def configure_prediction_callbacks(
        self,
        batch_size: int = 64,
        precision: int | str = 32,
        k: int | None = None,
        force_reindex: bool = True,
        metrics_callbacks: list[NLPTemplateCallback] | None = None,
        *args,
        **kwargs,
    ):
        if k is None:
            # we need the largest k for the prediction callback
            # get the max top_k for the prediction callback
            k = sorted(self.top_k, reverse=True)[0]
        if metrics_callbacks is None:
            metrics_callbacks = self.configure_metrics_callbacks()

        prediction_callback = GoldenRetrieverPredictionCallback(
            batch_size=batch_size,
            precision=precision,
            k=k,
            force_reindex=force_reindex,
            other_callbacks=metrics_callbacks,
            *args,
            **kwargs,
        )
        return prediction_callback

    def train(self, *args, **kwargs):
        """
        Train the model

        Args:
            *args:
                Additional args
            **kwargs:
                Additional kwargs

        Returns:
            `None`
        """
        if self.log_to_wandb:
            logger.info("Instantiating Wandb Logger")
            # log the args to wandb
            # logger.info(pformat(self.wandb_kwargs))
            self.wandb_logger = self.configure_logger(**self.wandb_kwargs)
            self.experiment_path = Path(self.wandb_logger.experiment.dir)

        # set-up training specific callbacks
        self.callbacks_store = self.training_callbacks()
        # add the evaluation callbacks
        self.callbacks_store.append(
            self.configure_prediction_callbacks(
                batch_size=self.prediction_batch_size,
                precision=self.precision,
            )
        )
        # add the hard negatives callback after the evaluation callback
        if self.max_hard_negatives_to_mine > 0:
            self.callbacks_store.append(self.configure_hard_negatives_callback())

        self.callbacks_store.append(FreeUpIndexerVRAMCallback())

        if self.trainer is None:
            logger.info("Instantiating the Trainer")
            self.trainer = pl.Trainer(
                accelerator=self.accelerator,
                devices=self.devices,
                num_nodes=self.num_nodes,
                strategy=self.strategy,
                accumulate_grad_batches=self.accumulate_grad_batches,
                max_epochs=self.max_epochs,
                max_steps=self.max_steps,
                gradient_clip_val=self.gradient_clip_val,
                val_check_interval=self.val_check_interval,
                check_val_every_n_epoch=self.check_val_every_n_epoch,
                deterministic=self.deterministic,
                fast_dev_run=self.fast_dev_run,
                precision=self.precision,
                reload_dataloaders_every_n_epochs=self.reload_dataloaders_every_n_epochs,
                callbacks=self.callbacks_store,
                logger=self.wandb_logger,
                **self.trainer_kwargs,
            )

        # # save this class as config to file
        # if self.experiment_path is not None:
        #     logger.info("Saving the configuration to file")
        #     self.experiment_path.mkdir(parents=True, exist_ok=True)
        #     OmegaConf.save(
        #         OmegaConf.create(to_config(self)),
        #         self.experiment_path / "trainer_config.yaml",
        #     )
        self.trainer.fit(
            self.lightning_module,
            datamodule=self.lightning_datamodule,
            ckpt_path=self.resume_from_checkpoint_path,
        )

    def test(
        self,
        lightning_module: GoldenRetrieverPLModule | None = None,
        checkpoint_path: str | os.PathLike | None = None,
        lightning_datamodule: GoldenRetrieverPLDataModule | None = None,
        force_reindex: bool = False,
        *args,
        **kwargs,
    ):
        """
        Test the model

        Args:
            lightning_module (`GoldenRetrieverPLModule`, optional, defaults to `None`):
                The lightning module to test
            checkpoint_path (`str`, `os.PathLike`, optional, defaults to `None`):
                The path to the checkpoint to load
            lightning_datamodule (`GoldenRetrieverPLDataModule`, optional, defaults to `None`):
                The lightning data module to use for testing
            *args:
                Additional args
            **kwargs:
                Additional kwargs

        Returns:
            `None`
        """
        if self.test_dataset is None:
            logger.warning("No test dataset provided. Skipping testing.")
            return

        if self.trainer is None:
            self.trainer = pl.Trainer(
                accelerator=self.accelerator,
                devices=self.devices,
                num_nodes=self.num_nodes,
                strategy=self.strategy,
                deterministic=self.deterministic,
                fast_dev_run=self.fast_dev_run,
                precision=self.precision,
                callbacks=[
                    self.configure_prediction_callbacks(
                        batch_size=self.prediction_batch_size,
                        precision=self.precision,
                        force_reindex=force_reindex,
                    )
                ],
                **self.trainer_kwargs,
            )
        if lightning_module is not None:
            best_lightning_module = lightning_module
        else:
            try:
                if self.fast_dev_run:
                    best_lightning_module = self.lightning_module
                else:
                    # load best model for testing
                    if checkpoint_path is not None:
                        best_model_path = checkpoint_path
                    elif self.checkpoint_path is not None:
                        best_model_path = self.checkpoint_path
                    elif self.model_checkpoint_callback:
                        best_model_path = self.model_checkpoint_callback.best_model_path
                    else:
                        raise ValueError(
                            "Either `checkpoint_path` or `model_checkpoint_callback` should "
                            "be provided to the trainer"
                        )
                    logger.info(f"Loading best model from {best_model_path}")

                    best_lightning_module = (
                        GoldenRetrieverPLModule.load_from_checkpoint(best_model_path)
                    )
            except Exception as e:
                logger.info(f"Failed to load the model from checkpoint: {e}")
                logger.info("Using last model instead")
                best_lightning_module = self.lightning_module

        lightning_datamodule = lightning_datamodule or self.lightning_datamodule
        # module test
        self.trainer.test(best_lightning_module, datamodule=lightning_datamodule)


def train(conf: omegaconf.DictConfig) -> None:
    logger.info("Starting training with config:")
    logger.info(pformat(OmegaConf.to_container(conf)))

    logger.info("Instantiating the Retriever")
    retriever: GoldenRetriever = hydra.utils.instantiate(
        conf.retriever, _recursive_=False
    )

    logger.info("Instantiating datasets")
    train_dataset: GoldenRetrieverDataset = hydra.utils.instantiate(
        conf.data.train_dataset, _recursive_=False
    )
    val_dataset: GoldenRetrieverDataset = hydra.utils.instantiate(
        conf.data.val_dataset, _recursive_=False
    )
    test_dataset: GoldenRetrieverDataset = hydra.utils.instantiate(
        conf.data.test_dataset, _recursive_=False
    )

    logger.info("Loading the document index")
    document_index: BaseDocumentIndex = hydra.utils.instantiate(
        conf.data.document_index, _recursive_=False
    )
    retriever.document_index = document_index

    logger.info("Instantiating the Trainer")
    trainer: Trainer = hydra.utils.instantiate(
        conf.train,
        retriever=retriever,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        test_dataset=test_dataset,
        _recursive_=False,
    )

    logger.info("Starting training")
    trainer.train()

    logger.info("Starting testing")
    trainer.test()

    logger.info("Training and testing completed")


@hydra.main(config_path="../../conf", config_name="default", version_base="1.3")
def main(conf: omegaconf.DictConfig):
    train(conf)


def train_hydra(conf: omegaconf.DictConfig) -> None:
    # reproducibility
    pl.seed_everything(conf.train.seed)
    torch.set_float32_matmul_precision(conf.train.float32_matmul_precision)

    logger.info(f"Starting training for [bold cyan]{conf.model_name}[/bold cyan] model")
    if conf.train.pl_trainer.fast_dev_run:
        logger.info(
            f"Debug mode {conf.train.pl_trainer.fast_dev_run}. Forcing debugger configuration"
        )
        # Debuggers don't like GPUs nor multiprocessing
        # conf.train.pl_trainer.accelerator = "cpu"
        conf.train.pl_trainer.devices = 1
        conf.train.pl_trainer.strategy = "auto"
        conf.train.pl_trainer.precision = 32
        if "num_workers" in conf.data.datamodule:
            conf.data.datamodule.num_workers = {
                k: 0 for k in conf.data.datamodule.num_workers
            }
        # Switch wandb to offline mode to prevent online logging
        conf.logging.log = None
        # remove model checkpoint callback
        conf.train.model_checkpoint_callback = None

    if "print_config" in conf and conf.print_config:
        # pprint(OmegaConf.to_container(conf), console=logger, expand_all=True)
        logger.info(pformat(OmegaConf.to_container(conf)))

    # data module declaration
    logger.info("Instantiating the Data Module")
    pl_data_module: GoldenRetrieverPLDataModule = hydra.utils.instantiate(
        conf.data.datamodule, _recursive_=False
    )
    # force setup to get labels initialized for the model
    pl_data_module.prepare_data()
    # main module declaration
    pl_module: Optional[GoldenRetrieverPLModule] = None

    if not conf.train.only_test:
        pl_data_module.setup("fit")

        # count the number of training steps
        if (
            "max_epochs" in conf.train.pl_trainer
            and conf.train.pl_trainer.max_epochs > 0
        ):
            num_training_steps = (
                len(pl_data_module.train_dataloader())
                * conf.train.pl_trainer.max_epochs
            )
            if "max_steps" in conf.train.pl_trainer:
                logger.info(
                    "Both `max_epochs` and `max_steps` are specified in the trainer configuration. "
                    "Will use `max_epochs` for the number of training steps"
                )
                conf.train.pl_trainer.max_steps = None
        elif (
            "max_steps" in conf.train.pl_trainer and conf.train.pl_trainer.max_steps > 0
        ):
            num_training_steps = conf.train.pl_trainer.max_steps
            conf.train.pl_trainer.max_epochs = None
        else:
            raise ValueError(
                "Either `max_epochs` or `max_steps` should be specified in the trainer configuration"
            )
        logger.info(f"Expected number of training steps: {num_training_steps}")

        if "lr_scheduler" in conf.model.pl_module and conf.model.pl_module.lr_scheduler:
            # set the number of warmup steps as x% of the total number of training steps
            if conf.model.pl_module.lr_scheduler.num_warmup_steps is None:
                if (
                    "warmup_steps_ratio" in conf.model.pl_module
                    and conf.model.pl_module.warmup_steps_ratio is not None
                ):
                    conf.model.pl_module.lr_scheduler.num_warmup_steps = int(
                        conf.model.pl_module.lr_scheduler.num_training_steps
                        * conf.model.pl_module.warmup_steps_ratio
                    )
                else:
                    conf.model.pl_module.lr_scheduler.num_warmup_steps = 0
            logger.info(
                f"Number of warmup steps: {conf.model.pl_module.lr_scheduler.num_warmup_steps}"
            )

        logger.info("Instantiating the Model")
        pl_module: GoldenRetrieverPLModule = hydra.utils.instantiate(
            conf.model.pl_module, _recursive_=False
        )
        if (
            "pretrain_ckpt_path" in conf.train
            and conf.train.pretrain_ckpt_path is not None
        ):
            logger.info(
                f"Loading pretrained checkpoint from {conf.train.pretrain_ckpt_path}"
            )
            pl_module.load_state_dict(
                torch.load(conf.train.pretrain_ckpt_path)["state_dict"], strict=False
            )

        if "compile" in conf.model.pl_module and conf.model.pl_module.compile:
            try:
                pl_module = torch.compile(pl_module, backend="inductor")
            except Exception:
                logger.info(
                    "Failed to compile the model, you may need to install PyTorch 2.0"
                )

    # callbacks declaration
    callbacks_store = [ModelSummary(max_depth=2)]

    experiment_logger: Optional[WandbLogger] = None
    experiment_path: Optional[Path] = None
    if conf.logging.log:
        logger.info("Instantiating Wandb Logger")
        experiment_logger = hydra.utils.instantiate(conf.logging.wandb_arg)
        if pl_module is not None:
            # it may happen that the model is not instantiated if we are only testing
            # in that case, we don't need to watch the model
            experiment_logger.watch(pl_module, **conf.logging.watch)
        experiment_path = Path(experiment_logger.experiment.dir)
        # Store the YaML config separately into the wandb dir
        yaml_conf: str = OmegaConf.to_yaml(cfg=conf)
        (experiment_path / "hparams.yaml").write_text(yaml_conf)
        # Add a Learning Rate Monitor callback to log the learning rate
        callbacks_store.append(LearningRateMonitor(logging_interval="step"))

    early_stopping_callback: Optional[EarlyStopping] = None
    if conf.train.early_stopping_callback is not None:
        early_stopping_callback = hydra.utils.instantiate(
            conf.train.early_stopping_callback
        )
        callbacks_store.append(early_stopping_callback)

    model_checkpoint_callback: Optional[ModelCheckpoint] = None
    if conf.train.model_checkpoint_callback is not None:
        model_checkpoint_callback = hydra.utils.instantiate(
            conf.train.model_checkpoint_callback,
            dirpath=experiment_path / "checkpoints" if experiment_path else None,
        )
        callbacks_store.append(model_checkpoint_callback)

    if "callbacks" in conf.train and conf.train.callbacks is not None:
        for _, callback in conf.train.callbacks.items():
            # callback can be a list of callbacks or a single callback
            if isinstance(callback, omegaconf.listconfig.ListConfig):
                for cb in callback:
                    if cb is not None:
                        callbacks_store.append(
                            hydra.utils.instantiate(cb, _recursive_=False)
                        )
            else:
                if callback is not None:
                    callbacks_store.append(hydra.utils.instantiate(callback))

    # trainer
    logger.info("Instantiating the Trainer")
    trainer: Trainer = hydra.utils.instantiate(
        conf.train.pl_trainer, callbacks=callbacks_store, logger=experiment_logger
    )

    if not conf.train.only_test:
        # module fit
        trainer.fit(pl_module, datamodule=pl_data_module)

    if conf.train.pl_trainer.fast_dev_run:
        best_pl_module = pl_module
    else:
        # load best model for testing
        if conf.train.checkpoint_path:
            best_model_path = conf.evaluation.checkpoint_path
        elif model_checkpoint_callback:
            best_model_path = model_checkpoint_callback.best_model_path
        else:
            raise ValueError(
                "Either `checkpoint_path` or `model_checkpoint_callback` should "
                "be specified in the evaluation configuration"
            )
        logger.info(f"Loading best model from {best_model_path}")

        try:
            best_pl_module = GoldenRetrieverPLModule.load_from_checkpoint(
                best_model_path
            )
        except Exception as e:
            logger.info(f"Failed to load the model from checkpoint: {e}")
            logger.info("Using last model instead")
            best_pl_module = pl_module
        if "compile" in conf.model.pl_module and conf.model.pl_module.compile:
            try:
                best_pl_module = torch.compile(best_pl_module, backend="inductor")
            except Exception:
                logger.info(
                    "Failed to compile the model, you may need to install PyTorch 2.0"
                )

    # module test
    trainer.test(best_pl_module, datamodule=pl_data_module)


if __name__ == "__main__":
    main()