"""The REaLTabFormer implements the model training and data processing
for tabular and relational data.
"""
import json
import logging
import math
import os
import random
import shutil
import time
import warnings
from collections import OrderedDict
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from sklearn.metrics.pairwise import manhattan_distances

# from sklearn.metrics import accuracy_score
from transformers import (
    EarlyStoppingCallback,
    EncoderDecoderConfig,
    EncoderDecoderModel,
    PreTrainedModel,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    Trainer,
    TrainingArguments,
)
from transformers.models.gpt2 import GPT2Config, GPT2LMHeadModel

import realtabformer

from .data_utils import (
    ModelFileName,
    ModelType,
    SpecialTokens,
    TabularArtefact,
    build_vocab,
    make_dataset,
    make_relational_dataset,
    process_data,
)
from .rtf_analyze import SyntheticDataBench
from .rtf_datacollator import RelationalDataCollator
from .rtf_exceptions import SampleEmptyLimitError
from .rtf_sampler import RelationalSampler, TabularSampler
from .rtf_trainer import ResumableTrainer
from .rtf_validators import ObservationValidator


def _normalize_gpt2_state_dict(state_dict):
    state = []
    for key, value in state_dict.items():
        if key.startswith("transformer."):
            # The saved state prefixes the weight names
            # with `transformer.` whereas the
            # encoder expects the weight names to not
            # have the prefix.
            key = key.replace("transformer.", "")

        state.append((key, value))

    return OrderedDict(state)


def _validate_get_device(device: str) -> str:
    if (device == "cuda") and (torch.cuda.device_count() == 0):
        if torch.backends.mps.is_available():
            _device = "mps"
        else:
            _device = "cpu"

        warnings.warn(
            f"The device={device} is not available, using device={_device} instead."
        )
        device = _device

    return device


class REaLTabFormer:
    def __init__(
        self,
        model_type: str,
        tabular_config: Optional[GPT2Config] = None,
        relational_config: Optional[EncoderDecoderConfig] = None,
        parent_realtabformer_path: Optional[Path] = None,
        freeze_parent_model: Optional[bool] = True,
        checkpoints_dir: str = "rtf_checkpoints",
        samples_save_dir: str = "rtf_samples",
        epochs: int = 100,
        batch_size: int = 8,
        random_state: int = 1029,
        train_size: float = 1,
        output_max_length: int = 512,
        early_stopping_patience: int = 5,
        early_stopping_threshold: float = 0,
        mask_rate: float = 0,
        numeric_nparts: int = 1,
        numeric_precision: int = 4,
        numeric_max_len: int = 10,
        **training_args_kwargs,
    ) -> None:
        """Set up a REaLTabFormer instance.

        Args:
            model_type: Explicit declaration of which type of model will be used.
                Can take `tabular` and `relational` as valid values.
            tabular_config: GPT2Config instance to customize the GPT2 model for tabular data
                generation.
            relational_config: EncoderDecoderConfig instance that defines the encoder and decoder
                configs for the encoder-decoder model used for the relational data generation. See
                link for example: https://huggingface.co/docs/transformers/model_doc/encoder-decoder
            parent_realtabformer_path: Path to a saved tabular REaLTabFormer model trained on the
                parent table of a relational tabular data.
            freeze_parent_model: Boolean flag indicating whether the parent-based encoder will be
                frozen or not.
            checkpoints_dir:  Directory where the training checkpoints will be saved
            samples_save_dir: Save the samples generated by this model in this directory.
            epochs: Number of epochs for training the GPT2LM model
            batch_size: Batch size used for training. Must be adjusted based on the available
                compute resource. TrainingArguments is set to use `gradient_accumulation_steps=4`
                which will have an effective batch_size of 32 for the default value.
            train_size: Fraction of the data that will be passed to the `.fit` method that will
                  be used for training. The remaining will be used as validation data.
            output_max_length: Truncation length for the number of output token ids in the
                relational model. This limit applies to the processed data and not the raw number
                of variables. This is not used in the tabular data model.
            early_stopping_patience: Number of evaluation rounds without improvement before
                stopping the training.
            early_stopping_threshold: See link
                https://huggingface.co/docs/transformers/main_classes/callback#transformers.EarlyStoppingCallback.early_stopping_threshold(float,
            mask_rate: The rate of tokens in the transformed observation that will be replaced
                with the [RMASK] token for regularization during training.
            training_args_kwargs: Keyword arguments for the `TrainingArguments` used in training
                the model. Arguments such as `output_dir`, `num_train_epochs`,
                `per_device_train_batch_size`, `per_device_eval_batch_size` if passed will be
                replaced by `checkpoints_dir`, `epochs`, `batch_size arguments`. The comprehensive
                set of options can be found in
              https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments
        """
        self.model: PreTrainedModel = None

        # This will be set during and will also be deleted after training.
        self.dataset = None

        if model_type not in ModelType.types():
            self._invalid_model_type(model_type)

        self.model_type = model_type

        if self.model_type == ModelType.tabular:
            self._init_tabular(tabular_config)
        elif self.model_type == ModelType.relational:
            self.parent_vocab = None
            self.parent_gpt2_config = None
            self.parent_gpt2_state_dict = None
            self.parent_col_transform_data = None
            self.freeze_parent_model = freeze_parent_model

            if parent_realtabformer_path is not None:
                parent_realtabformer_path = Path(parent_realtabformer_path)

                parent_config = json.loads(
                    (
                        parent_realtabformer_path / ModelFileName.rtf_config_json
                    ).read_text()
                )
                self.parent_col_transform_data = parent_config["col_transform_data"]
                self.parent_vocab = parent_config["vocab"]
                self.parent_gpt2_config = parent_config["tabular_config"]
                self.parent_gpt2_state_dict = _normalize_gpt2_state_dict(
                    torch.load(parent_realtabformer_path / ModelFileName.rtf_model_pt)
                )

            if output_max_length is None:
                warnings.warn(
                    "The `output_max_length` is None. This could result to extended model training if the output length has large variations."
                )

            self.output_max_length = output_max_length
            self._init_relational(relational_config)
        else:
            self._invalid_model_type(self.model_type)

        self.checkpoints_dir = Path(checkpoints_dir)
        self.samples_save_dir = Path(samples_save_dir)
        self.epochs = epochs
        self.batch_size = batch_size

        self.early_stopping_patience = early_stopping_patience
        self.early_stopping_threshold = early_stopping_threshold

        self.training_args_kwargs = dict(
            evaluation_strategy="steps",
            output_dir=self.checkpoints_dir.as_posix(),
            metric_for_best_model="loss",  # This will be replaced with "eval_loss" if `train_size` < 1
            overwrite_output_dir=True,
            num_train_epochs=self.epochs,
            per_device_train_batch_size=self.batch_size,
            per_device_eval_batch_size=self.batch_size,
            gradient_accumulation_steps=4,
            remove_unused_columns=True,
            logging_steps=100,
            save_steps=100,
            eval_steps=100,
            load_best_model_at_end=True,
            save_total_limit=early_stopping_patience + 1,
            optim="adamw_torch",
        )

        # Remove experiment params from `training_args_kwargs`
        for p in [
            "output_dir",
            "num_train_epochs",
            "per_device_train_batch_size",
            "per_device_eval_batch_size",
        ]:
            if p in training_args_kwargs:
                warnings.warn(
                    f"Argument {p} was passed in training_args_kwargs but will be ignored..."
                )
                training_args_kwargs.pop(p)

        self.training_args_kwargs.update(training_args_kwargs)

        self.train_size = train_size
        self.mask_rate = mask_rate

        self.columns: List[str] = []
        self.column_dtypes: Dict[str, type] = {}
        self.column_has_missing: Dict[str, bool] = {}
        self.drop_na_cols: List[str] = []
        self.processed_columns: List[str] = []
        self.numeric_columns: List[str] = []
        self.datetime_columns: List[str] = []
        self.vocab: Dict[str, dict] = {}
        # Output length for generator model
        # including special tokens.
        self.tabular_max_length = None
        self.relational_max_length = None
        # Number of derived columns for the relational
        # and tabular data after performing the data transformation.
        # This will be used as record size validator in the
        # sampling stage.
        self.tabular_col_size = None
        self.relational_col_size = None

        # This stores the transformation
        # parameters for numeric columns.
        self.col_transform_data: Optional[Dict] = None

        # This is the col_transform_data
        # for the relational models's in_df.
        self.in_col_transform_data: Optional[Dict] = None

        self.col_idx_ids: Dict[int, list] = {}

        self.random_state = random_state

        self.numeric_nparts = numeric_nparts
        self.numeric_precision = numeric_precision
        self.numeric_max_len = numeric_max_len

        # A unique identifier for the experiment set after the
        # model is trained.
        self.experiment_id = None
        self.trainer_state = None

        # Target column, when set, a copy of the column values will be
        # implicitly placed at the beginning of the dataframe.
        self.target_col = None

        self.realtabformer_version = realtabformer.__version__

    def _invalid_model_type(self, model_type):
        raise ValueError(
            f"Model type: {model_type} is not valid. REaLTabFormer only supports \
                `tabular` and `relational` values."
        )

    def _init_tabular(self, tabular_config):
        if tabular_config is not None:
            warnings.warn(
                "The `bos_token_id`, `eos_token_id`, and `vocab_size` attributes will \
                    be replaced when the `.fit` method is run."
            )
        else:
            # Default is 12, use 6 for distill-gpt2 as default
            tabular_config = GPT2Config(n_layer=6)

        self.tabular_config = tabular_config
        self.model = None

    def _init_relational(self, relational_config):
        if relational_config is not None:
            warnings.warn(
                "The `bos_token_id`, `eos_token_id`, and `vocab_size` attributes for the \
                    encoder and decoder will be replaced when the `.fit` method is run."
            )
        else:
            # Default is 12, use 6 for distill-gpt2 as default
            relational_config = EncoderDecoderConfig(
                encoder=GPT2Config(n_layer=6).to_dict(),
                decoder=GPT2Config(n_layer=6).to_dict(),
            )

        if self.parent_gpt2_config is not None:
            warnings.warn(
                "A trained model for the parent table is available. The encoder will use the \
                    pretrained config and weights."
            )
            relational_config.encoder = GPT2Config(**self.parent_gpt2_config)

        self.relational_config = relational_config
        self.model = None

    def _extract_column_info(self, df: pd.DataFrame) -> None:
        # Track the column order of the original data
        self.columns = df.columns.to_list()

        # Store the dtypes of the columns
        self.column_dtypes = df.dtypes.astype(str).to_dict()

        # Track which variables have missing values
        self.column_has_missing = (df.isnull().sum() > 0).to_dict()

        # Get the columns where there should be no missing values
        self.drop_na_cols = [
            col for col, has_na in self.column_has_missing.items() if not has_na
        ]

        # Identify the numeric columns. These will undergo
        # special preprocessing.
        self.numeric_columns = df.select_dtypes(include=np.number).columns.to_list()

        # Identify the datetime columns. These will undergo
        # special preprocessing.
        self.datetime_columns = df.select_dtypes(include="datetime").columns.to_list()

    def _generate_vocab(self, df: pd.DataFrame) -> dict:
        return build_vocab(df, special_tokens=SpecialTokens.tokens(), add_columns=False)

    def _check_model(self):
        assert self.model is not None, "Model is None. Train the model first!"

    def _split_train_eval_dataset(self, dataset: Dataset):
        test_size = 1 - self.train_size
        if test_size > 0:
            dataset = dataset.train_test_split(
                test_size=test_size, seed=self.random_state
            )
            dataset["train_dataset"] = dataset.pop("train")
            dataset["eval_dataset"] = dataset.pop("test")

            # Override `metric_for_best_model` from "loss" to "eval_loss"
            self.training_args_kwargs["metric_for_best_model"] = "eval_loss"
            # Make this explicit so that no assumption is made on the
            # direction of the metric improvement.
            self.training_args_kwargs["greater_is_better"] = False
        else:
            dataset = dict(train_dataset=dataset)
            self.training_args_kwargs["evaluation_strategy"] = "no"
            self.training_args_kwargs["load_best_model_at_end"] = False

        return dataset

    def fit(
        self,
        df: pd.DataFrame,
        in_df: Optional[pd.DataFrame] = None,
        join_on: Optional[str] = None,
        resume_from_checkpoint: Union[bool, str] = False,
        device="cuda",
        num_bootstrap: int = 500,
        frac: float = 0.165,
        frac_max_data: int = 10000,
        qt_max: Union[str, float] = 0.05,
        qt_max_default: float = 0.05,
        qt_interval: int = 100,
        qt_interval_unique: int = 100,
        distance: manhattan_distances = manhattan_distances,
        quantile: float = 0.95,
        n_critic: int = 5,
        n_critic_stop: int = 2,
        gen_rounds: int = 3,
        sensitivity_max_col_nums: int = 20,
        use_ks: bool = False,
        full_sensitivity: bool = False,
        sensitivity_orig_frac_multiple: int = 4,
        orig_samples_rounds: int = 5,
        load_from_best_mean_sensitivity: bool = False,
        target_col: str = None,
    ):
        """Train the REaLTabFormer model on the tabular data.

        Args:
            df: Pandas DataFrame containing the tabular data that will be generated during sampling.
              This data goes to the decoder for the relational model.
            in_df: Pandas DataFrame containing observations related to `df`, and from which the
              model will generate data. This data goes to the encoder for the relational model.
            join_on: Column name that links the `df` and the `in_df` tables.
            resume_from_checkpoint: If True, resumes training from the latest checkpoint in the
              checkpoints_dir. If path, resumes the training from the given checkpoint.
            device: Device where the model and the training will be run.
              Use torch devices, e.g., `cpu`, `cuda`, `mps` (experimental)
            num_bootstrap: Number of Bootstrap samples
            frac: The fraction of the data used for training.
            frac_max_data: The maximum number of rows that the training data will have.
            qt_max: The maximum quantile for the discriminator.
            qt_max_default: The default maximum quantile for the discriminator.
            qt_interval: Interval for the quantile check during the training process.
            qt_interval_unique: Interval for the quantile check during the training process.
            distance: Distance metric used for discriminator.
            quantile: The quantile value that the discriminator will be trained to.
            n_critic: Interval between epochs to perform a discriminator assessment.
            n_critic_stop: The number of critic rounds without improvement after which the training
              will be stopped.
            gen_rounds: The number of generator rounds.
            sensitivity_max_col_nums: The maximum number of columns used to compute sensitivity.
            use_ks: Whether to use KS test or not.
            full_sensitivity: Whether to use full sensitivity or not.
            sensitivity_orig_frac_multiple: The size of the training data relative to the chosen
              `frac` that will be used in computing the sensitivity. The larger this value is, the
              more robust the sensitivity threshold will be. However,
              `(sensitivity_orig_frac_multiple + 2)` multiplied by `frac` must be less than 1.
            orig_samples_rounds: This is the number of train/hold-out samples that will be used to
              compute the epoch sensitivity value.
            load_from_best_mean_sensitivity: Whether to load from best mean sensitivity or not.
            target_col: The target column name.

        Returns:
            Trainer
        """
        device = _validate_get_device(device)

        # Set target col for teacher forcing
        self.target_col = target_col

        # Set the seed for, *hopefully*, replicability.
        # This may cause an unexpected behavior when using
        # the resume_from_checkpoint option.
        if self.random_state:
            random.seed(self.random_state)
            np.random.seed(self.random_state)
            torch.manual_seed(self.random_state)
            torch.cuda.manual_seed_all(self.random_state)

        if self.model_type == ModelType.tabular:
            if n_critic <= 0:
                trainer = self._fit_tabular(df, device=device)
                trainer.train(resume_from_checkpoint=resume_from_checkpoint)
            else:
                trainer = self._train_with_sensitivity(
                    df,
                    device,
                    num_bootstrap=num_bootstrap,
                    frac=frac,
                    frac_max_data=frac_max_data,
                    qt_max=qt_max,
                    qt_max_default=qt_max_default,
                    qt_interval=qt_interval,
                    qt_interval_unique=qt_interval_unique,
                    distance=distance,
                    quantile=quantile,
                    n_critic=n_critic,
                    n_critic_stop=n_critic_stop,
                    gen_rounds=gen_rounds,
                    resume_from_checkpoint=resume_from_checkpoint,
                    sensitivity_max_col_nums=sensitivity_max_col_nums,
                    use_ks=use_ks,
                    full_sensitivity=full_sensitivity,
                    sensitivity_orig_frac_multiple=sensitivity_orig_frac_multiple,
                    orig_samples_rounds=orig_samples_rounds,
                    load_from_best_mean_sensitivity=load_from_best_mean_sensitivity,
                )

            del self.dataset

        elif self.model_type == ModelType.relational:
            assert (
                in_df is not None
            ), "The REaLTabFormer for relational data requires two tables for training."
            assert join_on is not None, "The column to join the data must not be None."

            trainer = self._fit_relational(df, in_df, join_on=join_on, device=device)
            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
        else:
            self._invalid_model_type(self.model_type)

        try:
            self.experiment_id = f"id{int((time.time() * 10 ** 10)):024}"
            torch.cuda.empty_cache()

            return trainer
        except Exception as exception:
            if device == torch.device("cuda"):
                del self.model
                torch.cuda.empty_cache()
                self.model = None

            raise exception

    def _train_with_sensitivity(
        self,
        df: pd.DataFrame,
        device: str = "cuda",
        num_bootstrap: int = 500,
        frac: float = 0.165,
        frac_max_data: int = 10000,
        qt_max: Union[str, float] = 0.05,
        qt_max_default: float = 0.05,
        qt_interval: int = 100,
        qt_interval_unique: int = 100,
        distance: manhattan_distances = manhattan_distances,
        quantile: float = 0.95,
        n_critic: int = 5,
        n_critic_stop: int = 2,
        gen_rounds: int = 3,
        sensitivity_max_col_nums: int = 20,
        use_ks: bool = False,
        resume_from_checkpoint: Union[bool, str] = False,
        full_sensitivity: bool = False,
        sensitivity_orig_frac_multiple: int = 4,
        orig_samples_rounds: int = 5,
        load_from_best_mean_sensitivity: bool = False,
    ):
        assert gen_rounds >= 1

        _frac = min(frac, frac_max_data / len(df))
        if frac != _frac:
            warnings.warn(
                f"The frac ({frac}) set results to a sample larger than \
                    frac_max_data={frac_max_data}. Setting frac to {_frac}."
            )
            frac = _frac

        trainer: Trainer = None
        dup_rate = df.duplicated().mean()

        if isinstance(qt_max, str):
            if qt_max == "compute":
                # The idea behind this is if the empirical has
                # natural duplicates, we can use that as
                # basis for what a typical rate for duplicates a
                # random sample should have. Any signidican excess
                # from this indicates overfitting.
                # The choice of dividing the duplicate rate by 2
                # is arbitrary but reasonable to prevent delayed
                # stopping when overfitting.
                dup_rate = dup_rate / 2
                qt_max = dup_rate if dup_rate > 0 else qt_max_default
            else:
                raise ValueError(f"Unexpected qt_max value: {qt_max}")
        elif not isinstance(qt_max, str) and dup_rate >= qt_max:
            warnings.warn(
                f'The qt_max ({qt_max}) set is lower than the duplicate \rate ({dup_rate}) in \
                    the data. This will not give a reliable early stopping condition. Consider \
                        using qt_max="compute" argument.'
            )

        if dup_rate == 0:
            # We do this because for data without unique values, we
            # expect that a generated sample should have equal likelihood
            # in the minimum distance with the hold out.
            warnings.warn(
                f"Duplicate rate ({dup_rate}) in the data is zero. The `qt_interval` will be set \
                    to qt_interval_unique={qt_interval_unique}."
            )
            qt_interval = qt_interval_unique

        # Estimate the sensitivity threshold
        print("Computing the sensitivity threshold...")

        if not full_sensitivity:
            # Dynamically compute the qt_interval to fit the data
            # if the resulting sample has lower resolution.
            # For example, we can't use qt_interval=1000 if the number
            # of samples left at qt_max of the distance matrix is less than
            # 1000.
            # The formula means:
            # - 2       -> accounts for the fact that we concatenate the rows and columns
            #             of the distance matrix.
            # - frac    -> the proportion of the training data that is used to compute the
            #             the distance matrix.
            # - qt_max  -> the maximum quantile of assessment.
            # We divide by 2 to increase the resolution a bit
            _qt_interval = min(qt_interval, (2 * frac * len(df) * qt_max) // 2)
            _qt_interval = max(_qt_interval, 2)
            _qt_interval = int(_qt_interval)

            if _qt_interval < qt_interval:
                warnings.warn(
                    f"qt_interval adjusted from {qt_interval} to {_qt_interval}..."
                )
                qt_interval = _qt_interval

        # Computing this here before splitting may have some data
        # leakage issue, but it should be almost negligible. Doing
        # the computation of the threshold on the full data with the
        # train size aligned will give a more reliable estimate of
        # the sensitivity threshold.
        sensitivity_values = SyntheticDataBench.compute_sensitivity_threshold(
            train_data=df,
            num_bootstrap=num_bootstrap,
            # Divide by two so that the train_data in this computation matches the size
            # of the final df used to train the model. This is essential so that the
            # sensitivity_threshold value is consistent with the val_sensitivity.
            # Concretely, the computation of the distribution of min distances is
            # relative to the number of training observations.
            # The `frac` in  this method corresponds to the size of both the test and the
            # synthetic samples.
            frac=frac / 2,
            qt_max=qt_max,
            qt_interval=qt_interval,
            distance=distance,
            return_values=True,
            quantile=quantile,
            max_col_nums=sensitivity_max_col_nums,
            use_ks=use_ks,
            full_sensitivity=full_sensitivity,
            sensitivity_orig_frac_multiple=sensitivity_orig_frac_multiple,
        )
        sensitivity_threshold = np.quantile(sensitivity_values, quantile)
        mean_sensitivity_value = np.mean(sensitivity_values)
        best_mean_sensitivity_value = np.inf

        assert isinstance(sensitivity_threshold, float)
        print("Sensitivity threshold:", sensitivity_threshold, "qt_max:", qt_max)

        # # Create a hold out sample for the discriminator model
        # hold_df = df.sample(frac=frac, random_state=self.random_state)
        # df = df.loc[df.index.difference(hold_df.index)]

        # Start training
        logging.info("Start training...")

        # Remove existing checkpoints
        for chkp in self.checkpoints_dir.glob("checkpoint-*"):
            shutil.rmtree(chkp, ignore_errors=True)

        sensitivity_scores = []
        bdm_path = self.checkpoints_dir / TabularArtefact.best_disc_model
        mean_closest_bdm_path = (
            self.checkpoints_dir / TabularArtefact.mean_best_disc_model
        )
        not_bdm_path = self.checkpoints_dir / TabularArtefact.not_best_disc_model
        last_epoch_path = self.checkpoints_dir / TabularArtefact.last_epoch_model

        # Remove existing artefacts in the best model dir
        shutil.rmtree(bdm_path, ignore_errors=True)
        bdm_path.mkdir(parents=True, exist_ok=True)

        shutil.rmtree(mean_closest_bdm_path, ignore_errors=True)
        mean_closest_bdm_path.mkdir(parents=True, exist_ok=True)

        shutil.rmtree(not_bdm_path, ignore_errors=True)
        not_bdm_path.mkdir(parents=True, exist_ok=True)

        shutil.rmtree(last_epoch_path, ignore_errors=True)
        last_epoch_path.mkdir(parents=True, exist_ok=True)

        last_epoch = 0
        not_best_val_sensitivity = np.inf

        if resume_from_checkpoint:
            chkp_list = sorted(
                self.checkpoints_dir.glob("checkpoint-*"), key=os.path.getmtime
            )
            if chkp_list:
                # Get the most recent checkpoint based on
                # creation time.
                chkp = chkp_list[-1]
                trainer_state = json.loads((chkp / "trainer_state.json").read_text())
                last_epoch = math.ceil(trainer_state["epoch"])

                trainer = self._fit_tabular(
                    df,
                    device=device,
                    num_train_epochs=last_epoch,
                    target_epochs=self.epochs,
                )

        np.random.seed(self.random_state)
        random.seed(self.random_state)

        for p_epoch in range(last_epoch, self.epochs, n_critic):
            gen_total = int(len(df) * frac)

            num_train_epochs = min(p_epoch + n_critic, self.epochs)
            # Perform the discriminator sampling every `n_critic` epochs
            if trainer is None:
                trainer = self._fit_tabular(
                    df,
                    device=device,
                    num_train_epochs=num_train_epochs,
                    target_epochs=self.epochs,
                )
                trainer.train(resume_from_checkpoint=False)
            else:
                trainer = self._build_tabular_trainer(
                    device=device,
                    num_train_epochs=num_train_epochs,
                    target_epochs=self.epochs,
                )
                trainer.train(resume_from_checkpoint=True)

            try:
                # Generate samples
                gen_df = self.sample(n_samples=gen_rounds * gen_total, device=device)
            except SampleEmptyLimitError:
                # Continue training if the model is still not
                # able to generate stable observations.
                continue

            val_sensitivities = []

            if full_sensitivity:
                for _ in range(gen_rounds):
                    hold_df = df.sample(n=gen_total)

                    for g_idx in range(gen_rounds):
                        val_sensitivities.append(
                            SyntheticDataBench.compute_sensitivity_metric(
                                original=df.loc[df.index.difference(hold_df.index)],
                                synthetic=gen_df.iloc[
                                    g_idx * gen_total : (g_idx + 1) * gen_total
                                ],
                                test=hold_df,
                                qt_max=qt_max,
                                qt_interval=qt_interval,
                                distance=distance,
                                max_col_nums=sensitivity_max_col_nums,
                                use_ks=use_ks,
                            )
                        )
            else:
                for g_idx in range(gen_rounds):
                    for _ in range(orig_samples_rounds):
                        original_df = df.sample(
                            n=sensitivity_orig_frac_multiple * gen_total, replace=False
                        )
                        hold_df = df.loc[df.index.difference(original_df.index)].sample(
                            n=gen_total, replace=False
                        )

                        val_sensitivities.append(
                            SyntheticDataBench.compute_sensitivity_metric(
                                original=original_df,
                                synthetic=gen_df.iloc[
                                    g_idx * gen_total : (g_idx + 1) * gen_total
                                ],
                                test=hold_df,
                                qt_max=qt_max,
                                qt_interval=qt_interval,
                                distance=distance,
                                max_col_nums=sensitivity_max_col_nums,
                                use_ks=use_ks,
                            )
                        )

            val_sensitivity = np.mean(val_sensitivities)

            sensitivity_scores.append(val_sensitivity)

            if val_sensitivity < sensitivity_threshold:
                # Just save the model while the
                # validation sensitivity is still within
                # the accepted range.
                # This way we can load the acceptable
                # model back when the threshold is breached.
                trainer.save_model(bdm_path.as_posix())
                trainer.state.save_to_json((bdm_path / "trainer_state.json").as_posix())

            elif not_best_val_sensitivity > (val_sensitivity - sensitivity_threshold):
                print("Saving not-best model...")
                trainer.save_model(not_bdm_path.as_posix())
                trainer.state.save_to_json(
                    (not_bdm_path / "trainer_state.json").as_posix()
                )
                not_best_val_sensitivity = val_sensitivity - sensitivity_threshold

            _delta_mean_sensitivity_value = abs(
                mean_sensitivity_value - val_sensitivity
            )

            if _delta_mean_sensitivity_value < best_mean_sensitivity_value:
                best_mean_sensitivity_value = _delta_mean_sensitivity_value
                trainer.save_model(mean_closest_bdm_path.as_posix())
                trainer.state.save_to_json(
                    (mean_closest_bdm_path / "trainer_state.json").as_posix()
                )

            print(
                f"Critic round: {p_epoch + n_critic}, \
                    sensitivity_threshold: {sensitivity_threshold}, \
                        val_sensitivity: {val_sensitivity}, \
                            val_sensitivities: {val_sensitivities}"
            )

            if len(sensitivity_scores) > n_critic_stop:
                n_no_improve = 0
                for sensitivity_score in sensitivity_scores[-n_critic_stop:]:
                    # We count no improvement if the score is not
                    # better than the best, and that the score is not
                    # better than the previous score.
                    if sensitivity_score > sensitivity_threshold:
                        n_no_improve += 1

                if n_no_improve == n_critic_stop:
                    print("Stopping training, no improvement in critic...")
                    break

        # Save last epoch artefacts before loading the best model.
        trainer.save_model(last_epoch_path.as_posix())
        trainer.state.save_to_json((last_epoch_path / "trainer_state.json").as_posix())

        loaded_model_path = None

        if not load_from_best_mean_sensitivity:
            if (bdm_path / "pytorch_model.bin").exists():
                loaded_model_path = bdm_path
        else:
            if (mean_closest_bdm_path / "pytorch_model.bin").exists():
                loaded_model_path = mean_closest_bdm_path

        if loaded_model_path is None:
            # There should always be at least one `mean_closest_bdm_path` but
            # in case it doesn't exist, try loading from `not_bdm_path`.
            warnings.warn(
                "No best model was saved. Loading the closest model to the sensitivity_threshold."
            )
            loaded_model_path = not_bdm_path

        self.model = self.model.from_pretrained(loaded_model_path.as_posix())
        self.trainer_state = json.loads(
            (loaded_model_path / "trainer_state.json").read_text()
        )

        return trainer

    def _set_up_relational_coder_configs(self) -> None:
        def _get_coder(coder_name) -> GPT2Config:
            return getattr(self.relational_config, coder_name)

        for coder_name in ["encoder", "decoder"]:
            coder = _get_coder(coder_name)

            coder.bos_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.BOS]
            coder.eos_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.EOS]
            coder.pad_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.PAD]
            coder.vocab_size = len(self.vocab[coder_name]["id2token"])

            if coder_name == "decoder":
                self.relational_config.bos_token_id = coder.bos_token_id
                self.relational_config.eos_token_id = coder.eos_token_id
                self.relational_config.pad_token_id = coder.pad_token_id
                self.relational_config.decoder_start_token_id = coder.eos_token_id

            # Make sure that we have at least the number of
            # columns in the transformed data as positions.
            # This will prevent runtime error.
            # `RuntimeError: CUDA error: device-side assert triggered`
            assert self.relational_max_length
            if (
                coder_name == "decoder"
                and coder.n_positions < self.relational_max_length
            ):
                coder.n_positions = 128 + self.relational_max_length
            elif coder_name == "encoder" and getattr(
                coder, "n_positions", getattr(coder, "max_position_embeddings")
            ) < len(self.vocab[coder_name]["column_token_ids"]):
                positions = 128 + len(self.vocab[coder_name]["column_token_ids"])
                try:
                    coder.n_positions = positions
                except:
                    coder.max_position_embeddings = positions

        # This must be set to True for the EncoderDecoderModel to work at least
        # with GPT2 as the decoder.
        self.relational_config.decoder.add_cross_attention = True

    def _fit_relational(
        self, out_df: pd.DataFrame, in_df: pd.DataFrame, join_on: str, device="cuda"
    ):
        # bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
        # bert2bert.config.eos_token_id = tokenizer.sep_token_id
        # bert2bert.config.pad_token_id = tokenizer.pad_token_id
        # bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size

        # All join values in the out_df must be present in the in_df.
        assert len(set(out_df[join_on].unique()).difference(in_df[join_on])) == 0

        # Get the list of index of observations that are related based on
        # the join_on variable.
        common_out_idx = (
            out_df.reset_index(drop=True)
            .groupby(join_on)
            .apply(lambda x: x.index.to_list())
        )

        # Track the mapping of index from input to the list of output indices.
        in_out_idx = pd.Series(
            # Reset the index so that we are sure that the index ids are set properly.
            dict(in_df[join_on].reset_index(drop=True).items())
        ).map(lambda x: common_out_idx.get(x, []))

        # Remove the unique id column from the in_df and the out_df
        in_df = in_df.drop(join_on, axis=1)
        out_df = out_df.drop(join_on, axis=1)

        self._extract_column_info(out_df)
        out_df, self.col_transform_data = process_data(
            out_df,
            numeric_max_len=self.numeric_max_len,
            numeric_precision=self.numeric_precision,
            numeric_nparts=self.numeric_nparts,
        )
        self.processed_columns = out_df.columns.to_list()
        self.vocab["decoder"] = self._generate_vocab(out_df)
        self.relational_col_size = out_df.shape[1]

        # NOTE: the index starts at zero, but should be adjusted
        # to account for the special tokens. For relational data,
        # the index should start at 3 ([[EOS], [BOS], [BMEM]]).
        self.col_idx_ids = {
            ix: self.vocab["decoder"]["column_token_ids"][col]
            for ix, col in enumerate(self.processed_columns)
        }

        # Add these special tokens at specific key values
        # which are used in `REaLSampler._get_relational_col_idx_ids`
        self.col_idx_ids[-1] = [
            self.vocab["decoder"]["token2id"][SpecialTokens.BMEM],
            self.vocab["decoder"]["token2id"][SpecialTokens.EOS],
        ]
        self.col_idx_ids[-2] = [self.vocab["decoder"]["token2id"][SpecialTokens.EMEM]]

        # TODO: handle the col_transform_data from the in_df as well.
        in_df, self.in_col_transform_data = process_data(
            in_df,
            numeric_max_len=self.numeric_max_len,
            numeric_precision=self.numeric_precision,
            numeric_nparts=self.numeric_nparts,
            col_transform_data=self.parent_col_transform_data,
        )
        if self.parent_vocab is None:
            self.vocab["encoder"] = self._generate_vocab(in_df)
        else:
            self.vocab["encoder"] = self.parent_vocab

        # Load the dataframe into a HuggingFace Dataset
        dataset = make_relational_dataset(
            in_df=in_df,
            out_df=out_df,
            vocab=self.vocab,
            in_out_idx=in_out_idx,
            output_max_length=self.output_max_length,
            mask_rate=self.mask_rate,
            return_token_type_ids=False,
        )

        # Compute the longest sequence of labels in the dataset and add a buffer of 1.
        self.relational_max_length = (
            max(
                dataset.map(lambda example: dict(length=len(example["labels"])))[
                    "length"
                ]
            )
            + 1
        )

        # Create train-eval split if specified
        dataset = self._split_train_eval_dataset(dataset)

        # Set up the config and the model
        self._set_up_relational_coder_configs()

        # Build the model.
        self.model = EncoderDecoderModel(self.relational_config)
        if self.parent_gpt2_state_dict is not None:
            pretrain_load = self.model.encoder.load_state_dict(
                self.parent_gpt2_state_dict, strict=False
            )
            assert (
                not pretrain_load.missing_keys
            ), "There should be no missing_keys after loading the pretrained GPT2 state!"

            if self.freeze_parent_model:
                # We freeze the weights if we use the pretrained
                # parent table model.
                for param in self.model.encoder.parameters():
                    param.requires_grad = False

        # Tell pytorch to run this model on the GPU.
        device = torch.device(device)
        if device == torch.device("cuda"):
            self.model.cuda()

        # Set TrainingArguments and the Seq2SeqTrainer
        training_args_kwargs = dict(self.training_args_kwargs)

        default_args_kwargs = dict(
            # predict_with_generate=True,
            # warmup_steps=2000,
            fp16=(
                device == torch.device("cuda")
            ),  # Use fp16 by default if using cuda device
        )

        for k, v in default_args_kwargs.items():
            if k not in training_args_kwargs:
                training_args_kwargs[k] = v

        callbacks = None
        if training_args_kwargs["load_best_model_at_end"]:
            callbacks = [
                EarlyStoppingCallback(
                    self.early_stopping_patience, self.early_stopping_threshold
                )
            ]

        # instantiate trainer
        trainer = Seq2SeqTrainer(
            model=self.model,
            args=Seq2SeqTrainingArguments(**training_args_kwargs),
            callbacks=callbacks,
            data_collator=RelationalDataCollator(),
            **dataset,
        )

        return trainer

    def _fit_tabular(
        self,
        df: pd.DataFrame,
        device="cuda",
        num_train_epochs: int = None,
        target_epochs: int = None,
    ) -> Trainer:
        self._extract_column_info(df)
        df, self.col_transform_data = process_data(
            df,
            numeric_max_len=self.numeric_max_len,
            numeric_precision=self.numeric_precision,
            numeric_nparts=self.numeric_nparts,
            target_col=self.target_col,
        )
        self.processed_columns = df.columns.to_list()
        self.vocab = self._generate_vocab(df)
        self.tabular_col_size = df.shape[0]

        # NOTE: the index starts at zero, but should be adjusted
        # to account for the special tokens. For tabular data,
        # the index should start at 1.
        self.col_idx_ids = {
            ix: self.vocab["column_token_ids"][col]
            for ix, col in enumerate(self.processed_columns)
        }

        # Load the dataframe into a HuggingFace Dataset
        dataset = make_dataset(
            df, self.vocab, mask_rate=self.mask_rate, return_token_type_ids=False
        )

        # Store the sequence length for the processed data
        self.tabular_max_length = len(dataset[0]["input_ids"])

        # Create train-eval split if specified
        dataset = self._split_train_eval_dataset(dataset)
        self.dataset = dataset

        # Set up the config and the model
        self.tabular_config.bos_token_id = self.vocab["token2id"][SpecialTokens.BOS]
        self.tabular_config.eos_token_id = self.vocab["token2id"][SpecialTokens.EOS]
        self.tabular_config.vocab_size = len(self.vocab["id2token"])

        # Make sure that we have at least the number of
        # columns in the transformed data as positions.
        if self.tabular_config.n_positions < len(self.vocab["column_token_ids"]):
            self.tabular_config.n_positions = 128 + len(self.vocab["column_token_ids"])

        self.model = GPT2LMHeadModel(self.tabular_config)

        # Tell pytorch to run this model on the GPU.
        device = torch.device(device)
        if device == torch.device("cuda"):
            self.model.cuda()

        return self._build_tabular_trainer(
            device=device,
            num_train_epochs=num_train_epochs,
            target_epochs=target_epochs,
        )

    def _build_tabular_trainer(
        self,
        device="cuda",
        num_train_epochs: int = None,
        target_epochs: int = None,
    ) -> Trainer:
        device = torch.device(device)

        # Set TrainingArguments and the Trainer
        logging.info("Set up the TrainingArguments and the Trainer...")
        training_args_kwargs: Dict[str, Any] = dict(self.training_args_kwargs)

        default_args_kwargs = dict(
            fp16=(
                device == torch.device("cuda")
            ),  # Use fp16 by default if using cuda device
        )

        for k, v in default_args_kwargs.items():
            if k not in training_args_kwargs:
                training_args_kwargs[k] = v

        if num_train_epochs is not None:
            training_args_kwargs["num_train_epochs"] = num_train_epochs

        # # NOTE: The `ResumableTrainer` will default to its original
        # # behavior (Trainer) if `target_epochs`` is None.
        # # Set the `target_epochs` to `num_train_epochs` if not specified.
        # if target_epochs is None:
        #     target_epochs = training_args_kwargs.get("num_train_epochs")

        callbacks = None
        if training_args_kwargs["load_best_model_at_end"]:
            callbacks = [
                EarlyStoppingCallback(
                    self.early_stopping_patience, self.early_stopping_threshold
                )
            ]

        assert self.dataset
        trainer = ResumableTrainer(
            target_epochs=target_epochs,
            save_epochs=None,
            model=self.model,
            args=TrainingArguments(**training_args_kwargs),
            data_collator=None,  # Use the default_data_collator
            callbacks=callbacks,
            **self.dataset,
        )

        return trainer

    def sample(
        self,
        n_samples: int = None,
        input_unique_ids: Optional[Union[pd.Series, List]] = None,
        input_df: Optional[pd.DataFrame] = None,
        input_ids: Optional[torch.tensor] = None,
        gen_batch: Optional[int] = 128,
        device: str = "cuda",
        seed_input: Optional[Union[pd.DataFrame, Dict[str, Any]]] = None,
        save_samples: Optional[bool] = False,
        constrain_tokens_gen: Optional[bool] = True,
        validator: Optional[ObservationValidator] = None,
        continuous_empty_limit: int = 10,
        suppress_tokens: Optional[List[int]] = None,
        forced_decoder_ids: Optional[List[List[int]]] = None,
        related_num: Optional[Union[int, List[int]]] = None,
        **generate_kwargs,
    ) -> pd.DataFrame:
        """Generate synthetic tabular data samples

        Args:
            n_samples: Number of synthetic samples to generate for the tabular data.
            input_unique_ids: The unique identifier that will be used to link the input
              data to the generated values when sampling for relational data.
            input_df: Pandas DataFrame containing the tabular input data.
            input_ids: (NOTE: the `input_df` argument is the preferred input)
              The input_ids that conditions the generation of the relational data.
            gen_batch: Controls the batch size of the data generation process. This parameter
              should be adjusted based on the compute resources.
            device: The device used by the generator.
              Use torch devices, e.g., `cpu`, `cuda`, `mps` (experimental)
            seed_input: A dictionary of `col_name:values` for the seed data. Only `col_names`
              that are actually in the first sequence of the training input will be used.
            constrain_tokens_gen: Set whether we impose a constraint at each step of the generation
              limited only to valid tokens for the column.
            validator: An instance of `ObservationValidator` for validating the generated samples.
              The validators are applied to observations only, and don't support inter-observation
              validation. See `ObservationValidator` docs on how to set up a validator.
            continuous_invalid_limit: The sampling will raise an exception if
              `continuous_empty_limit` empty sample batches have been produced continuously. This
              will prevent an infinite loop if the quality of the data generated is not good and
              always produces invalid observations.
            suppress_tokens: (from docs) A list of tokens that will be supressed at generation.
              The SupressTokens logit processor will set their log probs to -inf so that they are
              not sampled. This is a useful feature for imputing missing values.
            forced_decoder_ids: (from docs) A list of pairs of integers which indicates a mapping
              from generation indices to token indices that will be forced before sampling. For
              example, [[1, 123]] means the second generated token will always be a token of
              index 123. This is a useful feature for constraining the model to generate only
              specific stratification variables in surveys, e.g., GEO1, URBAN/RURAL variables.
            related_num: A column name in the input_df containing the number of observations that the child
             table is expected to have for the parent observation. It can also be an integer if the input_df
             corresponds to a set of observations having the same number of expected observations.
             This parameter is only valid for the relational model.
            generate_kwargs: Additional keywords arguments that will be supplied to `.generate`
              method. For a comprehensive list of arguments, see:
              https://huggingface.co/docs/transformers/v4.24.0/en/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate

        Returns:
            DataFrame with n_samples rows of generated data
        """
        self._check_model()
        device = _validate_get_device(device)

        # Clear the cache
        torch.cuda.empty_cache()

        if self.model_type == ModelType.tabular:
            assert n_samples
            assert self.tabular_max_length is not None
            assert self.tabular_col_size is not None
            assert self.col_transform_data is not None

            tabular_sampler = TabularSampler.sampler_from_model(
                rtf_model=self, device=device
            )
            # (
            #     model_type=self.model_type,
            #     model=self.model,
            #     vocab=self.vocab,
            #     processed_columns=self.processed_columns,
            #     max_length=self.tabular_max_length,
            #     col_size=self.tabular_col_size,
            #     col_idx_ids=self.col_idx_ids,
            #     columns=self.columns,
            #     datetime_columns=self.datetime_columns,
            #     column_dtypes=self.column_dtypes,
            #     drop_na_cols=self.drop_na_cols,
            #     col_transform_data=self.col_transform_data,
            #     random_state=self.random_state,
            #     device=device,
            # )
            synth_df = tabular_sampler.sample_tabular(
                n_samples=n_samples,
                gen_batch=gen_batch,
                device=device,
                seed_input=seed_input,
                constrain_tokens_gen=constrain_tokens_gen,
                validator=validator,
                continuous_empty_limit=continuous_empty_limit,
                suppress_tokens=suppress_tokens,
                forced_decoder_ids=forced_decoder_ids,
                **generate_kwargs,
            )

        elif self.model_type == ModelType.relational:
            assert (input_ids is not None) or (input_df is not None)
            assert self.relational_max_length is not None
            assert self.relational_col_size is not None
            assert self.col_transform_data is not None
            assert self.in_col_transform_data is not None

            relational_sampler = RelationalSampler.sampler_from_model(
                rtf_model=self, device=device
            )
            # (
            #     model_type=self.model_type,
            #     model=self.model,
            #     vocab=self.vocab,
            #     processed_columns=self.processed_columns,
            #     max_length=self.relational_max_length,
            #     col_size=self.relational_col_size,
            #     col_idx_ids=self.col_idx_ids,
            #     columns=self.columns,
            #     datetime_columns=self.datetime_columns,
            #     column_dtypes=self.column_dtypes,
            #     drop_na_cols=self.drop_na_cols,
            #     col_transform_data=self.col_transform_data,
            #     in_col_transform_data=self.in_col_transform_data,
            #     random_state=self.random_state,
            #     device=device,
            # )

            synth_df = relational_sampler.sample_relational(
                input_unique_ids=input_unique_ids,
                input_df=input_df,
                input_ids=input_ids,
                device=device,
                gen_batch=gen_batch,
                constrain_tokens_gen=constrain_tokens_gen,
                validator=validator,
                continuous_empty_limit=continuous_empty_limit,
                suppress_tokens=suppress_tokens,
                forced_decoder_ids=forced_decoder_ids,
                related_num=related_num,
                **generate_kwargs,
            )

        if save_samples:
            samples_fname = (
                self.samples_save_dir
                / f"rtf_{self.model_type}-exp_{self.experiment_id}-{int(time.time())}-samples_{synth_df.shape[0]}.pkl"
            )
            samples_fname.parent.mkdir(parents=True, exist_ok=True)
            synth_df.to_pickle(samples_fname)

        return synth_df

    def predict(
        self,
        data: pd.DataFrame,
        target_col: str,
        target_pos_val: Any = None,
        batch: int = 32,
        obs_sample: int = 30,
        fillunk: bool = True,
        device: str = "cuda",
        disable_progress_bar: bool = True,
        **generate_kwargs,
    ) -> pd.Series:
        """
        Use the trained model to make predictions on a given dataframe.

        Args:
            data: The data to make predictions on, in the form of a Pandas dataframe.
            target_col: The name of the target column in the data to predict.
            target_pos_val: The positive value in the target column to use for binary
              classification. This is produces a one-to-many prediction relative to
              `target_pos_val` for targets that are multi-categorical.
            batch: The batch size to use when making predictions.
            obs_sample: The number of observations to sample from the data when making predictions.
            fillunk: If True, the function will fill any missing values in the data before making
              predictions. Fill unknown tokens with the mode of the batch in the given step.
            device: The device to use for prediction. Can be either "cpu" or "cuda".
            **generate_kwargs: Additional keyword arguments to pass to the model's `generate`
              method.

        Returns:
            A Pandas series containing the predicted values for the target column.
        """

        assert (
            self.model_type == ModelType.tabular
        ), "The predict method is only implemented for tabular data..."
        self._check_model()
        device = _validate_get_device(device)
        batch = min(batch, data.shape[0])

        # Clear the cache
        torch.cuda.empty_cache()

        # assert self.tabular_max_length is not None
        # assert self.tabular_col_size is not None
        # assert self.col_transform_data is not None

        tabular_sampler = TabularSampler.sampler_from_model(self, device=device)

        # TabularSampler(
        #     model_type=self.model_type,
        #     model=self.model,
        #     vocab=self.vocab,
        #     processed_columns=self.processed_columns,
        #     max_length=self.tabular_max_length,
        #     col_size=self.tabular_col_size,
        #     col_idx_ids=self.col_idx_ids,
        #     columns=self.columns,
        #     datetime_columns=self.datetime_columns,
        #     column_dtypes=self.column_dtypes,
        #     drop_na_cols=self.drop_na_cols,
        #     col_transform_data=self.col_transform_data,
        #     random_state=self.random_state,
        #     device=device,
        # )

        return tabular_sampler.predict(
            data=data,
            target_col=target_col,
            target_pos_val=target_pos_val,
            batch=batch,
            obs_sample=obs_sample,
            fillunk=fillunk,
            device=device,
            disable_progress_bar=disable_progress_bar,
            **generate_kwargs,
        )

    def save(self, path: Union[str, Path], allow_overwrite: Optional[bool] = False):
        """Save REaLTabFormer Model

        Saves the model weights and a configuration file in the given directory.
        Args:
            path: Path where to save the model
        """
        self._check_model()
        assert self.experiment_id is not None

        if isinstance(path, str):
            path = Path(path)

        # Add experiment id to the save path
        path = path / self.experiment_id

        config_file = path / ModelFileName.rtf_config_json
        model_file = path / ModelFileName.rtf_model_pt

        if path.is_dir() and not allow_overwrite:
            if config_file.exists() or model_file.exists():
                raise ValueError(
                    "This directory is not empty, and contains either a config or a model."
                    " Consider setting `allow_overwrite=True` if you want to overwrite these."
                )
            else:
                warnings.warn(
                    f"Directory {path} exists, but `allow_overwrite=False`."
                    " This will raise an error next time when the model artifacts \
                        exist on this directory"
                )

        path.mkdir(parents=True, exist_ok=True)

        # Save attributes
        rtf_attrs = self.__dict__.copy()
        rtf_attrs.pop("model")

        # We don't need to store the `parent_config`
        # since a saved model should have the weights loaded from
        # the trained model already.
        for ignore_key in [
            "parent_vocab",
            "parent_gpt2_config",
            "parent_gpt2_state_dict",
            "parent_col_transform_data",
        ]:
            if ignore_key in rtf_attrs:
                rtf_attrs.pop(ignore_key)

        # GPT2Config is not JSON serializable, let us manually
        # extract the attributes.
        if rtf_attrs.get("tabular_config"):
            rtf_attrs["tabular_config"] = rtf_attrs["tabular_config"].to_dict()

        if rtf_attrs.get("relational_config"):
            rtf_attrs["relational_config"] = rtf_attrs["relational_config"].to_dict()

        rtf_attrs["checkpoints_dir"] = rtf_attrs["checkpoints_dir"].as_posix()
        rtf_attrs["samples_save_dir"] = rtf_attrs["samples_save_dir"].as_posix()

        config_file.write_text(json.dumps(rtf_attrs))

        # Save model weights
        torch.save(self.model.state_dict(), model_file.as_posix())

        if self.model_type == ModelType.tabular:
            # Copy the special model checkpoints for
            # tabular models.
            for artefact in TabularArtefact.artefacts():
                print("Copying artefacts from:", artefact)
                if (self.checkpoints_dir / artefact).exists():
                    shutil.copytree(
                        self.checkpoints_dir / artefact,
                        path / artefact,
                        dirs_exist_ok=True,
                    )

    @classmethod
    def load_from_dir(cls, path: Union[str, Path]):
        """Load a saved REaLTabFormer model

        Load trained REaLTabFormer model from directory.
        Args:
            path: Directory where REaLTabFormer model is saved
        Returns:
            REaLTabFormer instance
        """

        if isinstance(path, str):
            path = Path(path)

        config_file = path / ModelFileName.rtf_config_json
        model_file = path / ModelFileName.rtf_model_pt

        assert path.is_dir(), f"Directory {path} does not exist."
        assert config_file.exists(), f"Config file {config_file} does not exist."
        assert model_file.exists(), f"Model file {model_file} does not exist."

        # Load the saved attributes
        rtf_attrs = json.loads(config_file.read_text())

        # Create new REaLTabFormer model instance
        try:
            realtf = cls(model_type=rtf_attrs["model_type"])
        except KeyError:
            # Back-compatibility for saved models
            # before the support for relational data
            # was implemented.
            realtf = cls(model_type="tabular")

        # Set all attributes and handle the
        # special case for the GPT2Config.
        for k, v in rtf_attrs.items():
            if k == "gpt_config":
                # Back-compatibility for saved models
                # before the support for relational data
                # was implemented.
                v = GPT2Config.from_dict(v)
                k = "tabular_config"

            elif k == "tabular_config":
                v = GPT2Config.from_dict(v)

            elif k == "relational_config":
                v = EncoderDecoderConfig.from_dict(v)

            elif k in ["checkpoints_dir", "samples_save_dir"]:
                v = Path(v)

            elif k == "vocab":
                if realtf.model_type == ModelType.tabular:
                    # Cast id back to int since JSON converts them to string.
                    v["id2token"] = {int(ii): vv for ii, vv in v["id2token"].items()}
                elif realtf.model_type == ModelType.relational:
                    v["encoder"]["id2token"] = {
                        int(ii): vv for ii, vv in v["encoder"]["id2token"].items()
                    }
                    v["decoder"]["id2token"] = {
                        int(ii): vv for ii, vv in v["decoder"]["id2token"].items()
                    }
                else:
                    raise ValueError(f"Invalid model_type: {realtf.model_type}")

            elif k == "col_idx_ids":
                v = {int(ii): vv for ii, vv in v.items()}

            setattr(realtf, k, v)

        # Implement back-compatibility for REaLTabFormer version < 0.0.1.8.2
        # since the attribute `col_idx_ids` is not implemented before.
        if "col_idx_ids" not in rtf_attrs:
            if realtf.model_type == ModelType.tabular:
                realtf.col_idx_ids = {
                    ix: realtf.vocab["column_token_ids"][col]
                    for ix, col in enumerate(realtf.processed_columns)
                }
            elif realtf.model_type == ModelType.relational:
                # NOTE: the index starts at zero, but should be adjusted
                # to account for the special tokens. For relational data,
                # the index should start at 3 ([[EOS], [BOS], [BMEM]]).
                realtf.col_idx_ids = {
                    ix: realtf.vocab["decoder"]["column_token_ids"][col]
                    for ix, col in enumerate(realtf.processed_columns)
                }

                # Add these special tokens at specific key values
                # which are used in `REaLSampler._get_relational_col_idx_ids`
                realtf.col_idx_ids[-1] = [
                    realtf.vocab["decoder"]["token2id"][SpecialTokens.BMEM],
                    realtf.vocab["decoder"]["token2id"][SpecialTokens.EOS],
                ]
                realtf.col_idx_ids[-2] = [
                    realtf.vocab["decoder"]["token2id"][SpecialTokens.EMEM]
                ]

        # Load model weights
        if realtf.model_type == ModelType.tabular:
            realtf.model = GPT2LMHeadModel(realtf.tabular_config)
        elif realtf.model_type == ModelType.relational:
            realtf.model = EncoderDecoderModel(realtf.relational_config)
        else:
            raise ValueError(f"Invalid model_type: {realtf.model_type}")

        realtf.model.load_state_dict(
            torch.load(model_file.as_posix(), map_location="cpu")
        )

        return realtf
