first commit

README.md

@@ -0,0 +1,28 @@
+# TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer
+
+## Pre-requisites
+
+- `Python>=3.9` installed.
+- `pip install -r requirements.txt`.
+
+## Usage
+
+To train,
+
+```shell
+python main.py train -d DATA_PATH -t TARGET_COLUMN -p TASK_TYPE -o OUT_DIR
+```
+
+After training, to generate,
+
+```shell
+python main.py sample -c OUT_DIR -n N_ROWS -o OUT_CSV_PATH
+```
+
+For instance, to train and sample iris dataset (one can get the dataset by `sklearn.datasets.load_iris`),
+one can run the following:
+
+```shell
+python main.py train -d iris.csv -t target -p mult -o out
+python main.py sample -c out -n 150 -o synthetic-iris.csv
+```

main.py

@@ -0,0 +1,50 @@
+import argparse
+import os.path
+
+import pandas as pd
+import torch
+
+from tabtreeformer import TabTreeFormer
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="op")
+
+    train_parser = subparsers.add_parser("train")
+    train_parser.add_argument("--data-path", "-d", type=str, required=True,
+                              help="Path to data (.csv file).")
+    train_parser.add_argument("--target", "-t", type=str, required=True,
+                              help="Target column name.")
+    train_parser.add_argument("--ttype", "-p", type=str, required=True, choices=["bin", "mult", "reg"],
+                              help="Task type.")
+    train_parser.add_argument("--out", "-o", type=str, required=True,
+                              help="Path to output directory.")
+
+    sample_parser = subparsers.add_parser("sample")
+    sample_parser.add_argument("--ckpt-path", "-c", type=str, required=True,
+                               help="Path to checkpoint directory (output directory during training).")
+    sample_parser.add_argument("--n-rows", "-n", type=int, required=True,
+                               help="Number of rows to sample.")
+    sample_parser.add_argument("--out", "-o", type=str, required=True,
+                               help="Path to output synthetic data (.csv file).")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.op == "train":
+        data = pd.read_csv(args.data_path)
+        ttf = TabTreeFormer()
+        ttf.train(data, args.target, args.ttype, args.out)
+        torch.save(ttf, os.path.join(args.out, "ttf.pkl"))
+    elif args.op == "sample":
+        ttf: TabTreeFormer = torch.load(os.path.join(args.ckpt_path, "ttf.pkl"))
+        sampled = ttf.sample(args.n_rows)
+        sampled.to_csv(args.out, index=False)
+    else:
+        raise ValueError("Invalid op.")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+numpy
+lightgbm
+optuna
+pandas
+scikit-learn
+torch
+transformers

tabtreeformer/__init__.py

@@ -0,0 +1 @@
+from .model import TabTreeFormer

tabtreeformer/data.py

@@ -0,0 +1,423 @@
+"""Data Handling."""
+
+from typing import List, Literal, Optional, Tuple, Sequence, Union
+
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.preprocessing import KBinsDiscretizer, OrdinalEncoder, StandardScaler
+from transformers import BatchEncoding, LogitsProcessor
+
+
+class Dataset:
+    def __init__(self, data: pd.DataFrame,
+                 target: str,
+                 ttype: Literal["bin", "mult", "reg"]):
+        """
+        Parameters
+        ----------
+        data : pd.DataFrame
+            The dataset data.
+        target : str
+            The target column's name.
+        ttype : "bin" | "mult" | "reg"
+            Task type. Values can be "bin" for binary classification, "mult" for multiclass classification,
+            "reg" for regression.
+        """
+        n_unique = data[target].nunique()
+
+        self.data = data.copy()
+        self.target = target
+        self.ttype = ttype
+        self.n_classes = n_unique if ttype != "reg" else -1
+
+        self.num_columns = []
+        self.cat_columns = []
+        for c in self.data.columns:
+            if data[c].nunique() <= max(3, 10) or not pd.api.types.is_numeric_dtype(data[c]):
+                self.cat_columns.append(c)
+            else:
+                self.num_columns.append(c)
+        self.columns = data.columns.tolist()
+
+        cat_data = self.data[self.cat_columns].values
+        self._oe = OrdinalEncoder()
+        self.max_n_categories = 0
+        if cat_data.shape[-1] > 0:
+            self._cat_data = self._oe.fit_transform(cat_data)
+            self.max_n_categories = max(len(x) for x in self._oe.categories_)
+        else:
+            self._cat_data = cat_data
+        num_data = self.data[self.num_columns].values
+        self.max_n_bins = 0
+        self._kbins = KBinsDiscretizer(strategy="kmeans", n_bins=10, encode="ordinal", )
+        if len(self.num_columns) <= 0:
+            self._bin_data = num_data
+        else:
+            self._bin_data = self._kbins.fit_transform(num_data)
+            self.max_n_bins = max(x for x in self._kbins.n_bins_)
+        self.max_n_quantiles = 0
+        self._qbins = KBinsDiscretizer(strategy="quantile", n_bins=1000, encode="ordinal", )
+        self._sc = StandardScaler()
+        if len(self.num_columns) <= 0:
+            self._quantiles = num_data
+            self._num_data = num_data
+        else:
+            self._quantiles = self._qbins.fit_transform(num_data)
+            self.max_n_quantiles = max(x for x in self._qbins.n_bins_)
+            self._num_data = self._sc.fit_transform(num_data)
+
+        self._orders = []
+        self._target_descr = None
+        for c in self.columns:
+            if c in self.num_columns:
+                descr = "num", self.num_columns.index(c)
+            else:
+                descr = "cat", self.cat_columns.index(c)
+            self._orders.append(descr)
+            if c == self.target:
+                self._target_descr = descr
+
+        self.index_matrix = []
+        self.index_description = []
+        self._cat_indices = []
+        self._quantile_indices = []
+        idx = 0
+        for dtype, index in self._orders:
+            if dtype == "num":
+                if self._bin_data is not None:
+                    self.index_matrix.append(self._bin_data[:, index])
+                    if self._kbins is not None:
+                        self.index_description.append(("bin", self._kbins.n_bins_[index]))
+                    else:
+                        self.index_description.append(("bin", 1))
+                    idx += 1
+                self.index_matrix.append(self._quantiles[:, index])
+                self.index_description.append(("quantile", self._qbins.n_bins_[index]))
+                self._quantile_indices.append(idx)
+            else:
+                self.index_matrix.append(self._cat_data[:, index])
+                self.index_description.append(("cat", len(self._oe.categories_[index])))
+                self._cat_indices.append(idx)
+            idx += 1
+        self.index_matrix = np.stack(self.index_matrix, axis=1).astype(np.int32)
+
+        cat_x = self._cat_data
+        num_x = self._num_data
+        y = None
+        if self._target_descr is not None:
+            tdtype, tindex = self._target_descr
+            if tdtype == "cat":
+                y = cat_x[:, tindex].astype(np.int32)
+                cat_x = np.concatenate([cat_x[:, :tindex], cat_x[:, tindex + 1:]], axis=1)
+            else:
+                y = num_x[:, tindex]
+                num_x = np.concatenate([num_x[:, :tindex], num_x[:, tindex + 1:]], axis=1)
+        self.transformed = np.concatenate([cat_x, num_x], axis=1)
+        self.y = y
+
+    def get_index_matrix(self, df: pd.DataFrame) -> np.ndarray:
+        """
+        Transform raw data into index matrix.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            The raw data to be converted.
+
+        Returns
+        -------
+        np.ndarray
+            Transformed data.
+        """
+        cat_data = df[self.cat_columns].values
+        if len(self.cat_columns) > 0:
+            cat_data = self._oe.transform(cat_data)
+        num_data = df[self.num_columns].values
+        if self._kbins is not None:
+            if len(self.num_columns) > 0:
+                bin_data = self._kbins.transform(num_data)
+            else:
+                bin_data = num_data
+        elif self._bin_data is not None:
+            bin_data = np.zeros_like(num_data, dtype=np.int32)
+        else:
+            bin_data = None
+        if len(self.num_columns) > 0:
+            quantiles = self._qbins.transform(num_data)
+        else:
+            quantiles = num_data
+
+        out = []
+        for dtype, index in self._orders:
+            if dtype == "num":
+                if bin_data is not None:
+                    out.append(bin_data[:, index])
+                out.append(quantiles[:, index])
+            else:
+                out.append(cat_data[:, index])
+        return np.stack(out, axis=1).astype(np.int32)
+
+    def recover_index_matrix(self, data: np.ndarray) -> pd.DataFrame:
+        """
+        Inversely transform index matrix to raw data.
+
+        Parameters
+        ----------
+        np.ndarray
+            Transformed data.
+
+        Returns
+        -------
+        pd.DataFrame
+            The raw data recovered.
+        """
+        cat_data = data[:, self._cat_indices]
+        if len(self.cat_columns) > 0:
+            cat_data = self._oe.inverse_transform(cat_data)
+            cat_data = pd.DataFrame(cat_data, columns=self.cat_columns)
+        else:
+            cat_data = pd.DataFrame(index=pd.RangeIndex(data.shape[0]))
+        num_data = data[:, self._quantile_indices]
+        if len(self.num_columns) > 0:
+            num_data = self._qbins.inverse_transform(num_data)
+            num_data = pd.DataFrame(num_data, columns=self.num_columns)
+        else:
+            num_data = pd.DataFrame(index=pd.RangeIndex(data.shape[0]))
+        recovered = pd.concat([cat_data, num_data], axis=1)
+        return recovered[self.columns]
+
+    def transform(self, df: pd.DataFrame) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        """
+        Transform raw data to matrix that is friendly to downstream tasks.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            The raw data to be transformed.
+
+        Returns
+        -------
+        np.ndarray
+            X of transformed data (features).
+        np.ndarray, optional
+            y of transformed data (targets).
+        """
+        cat_data = df[self.cat_columns].values
+        if len(self.cat_columns) > 0:
+            cat_data = self._oe.transform(cat_data)
+        num_data = df[self.num_columns].values
+        if len(self.num_columns) > 0:
+            num_data = self._sc.transform(num_data)
+        if self._target_descr is None:
+            return np.concatenate([cat_data, num_data], axis=1), None
+        tdtype, tindex = self._target_descr
+        if tdtype == "cat":
+            y = cat_data[:, tindex]
+            cat_data = np.concatenate([cat_data[:, :tindex], cat_data[:, tindex + 1:]], axis=1)
+        else:
+            y = num_data[:, tindex]
+            num_data = np.concatenate([num_data[:, :tindex], num_data[:, tindex + 1:]], axis=1)
+        return np.concatenate([cat_data, num_data], axis=1), y
+
+
+class MaskGenerator:
+    def __init__(self, data: Dataset, leaf_dim: int,
+                 tree_mask_ratio: Union[float, Tuple[float, float]],
+                 data_mask_ratio: Union[float, Tuple[float, float]]):
+        """
+        Parameters
+        ----------
+        data : Dataset
+            The tabular dataset.
+        leaf_dim : int
+            The number of dimensions for leaf.
+        tree_mask_ratio : float | (float, float)
+            Mask ratio of tree leaves.
+        data_mask_ratio : float | (float, float)
+            Mask ratio of data values.
+        """
+        # bos_id: 0, eos_id: 1, mask_id: 2
+        self.leaf_dim = leaf_dim
+        self.data_dim = data.index_matrix.shape[-1]
+        if isinstance(tree_mask_ratio, Sequence):
+            self.min_tree_mask_ratio, self.max_tree_mask_ratio = tree_mask_ratio
+        else:
+            self.min_tree_mask_ratio, self.max_tree_mask_ratio = tree_mask_ratio, tree_mask_ratio
+        if isinstance(data_mask_ratio, Sequence):
+            self.min_data_mask_ratio, self.max_data_mask_ratio = data_mask_ratio
+        else:
+            self.min_data_mask_ratio, self.max_data_mask_ratio = data_mask_ratio, data_mask_ratio
+
+        self._bin_indices = []
+        self._quantile_indices = []
+        for i, (itype, _) in enumerate(data.index_description):
+            pos = i + 1 + self.leaf_dim
+            if itype == "quantile":
+                self._quantile_indices.append(pos)
+            elif itype == "bin":
+                self._bin_indices.append(pos)
+        if len(self._quantile_indices) > 0 and len(self._bin_indices) == 0:
+            self._bin_indices = None
+
+    def generate_mask(self, batch_size: int, tree_threshold: Optional[torch.FloatTensor] = None,
+                      data_threshold: Optional[torch.FloatTensor] = None,
+                      prev_mask: Optional[torch.BoolTensor] = None) -> torch.BoolTensor:
+        if tree_threshold is None:
+            tree_threshold = torch.rand(batch_size) * (
+                    self.max_tree_mask_ratio - self.min_tree_mask_ratio
+            ) + self.min_tree_mask_ratio
+        if data_threshold is None:
+            data_threshold = torch.rand(batch_size) * (
+                    self.max_data_mask_ratio - self.min_data_mask_ratio
+            ) + self.min_data_mask_ratio
+        tree_mask = torch.rand(batch_size, self.leaf_dim) < tree_threshold.view(-1, 1)
+        data_mask = torch.rand(batch_size, self.data_dim) < data_threshold.view(-1, 1)
+        if prev_mask is not None:
+            tree_mask = tree_mask.masked_fill(~prev_mask[:, 1:1 + self.leaf_dim], False)
+            data_mask = data_mask.masked_fill(~prev_mask[:, 1 + self.leaf_dim:-1], False)
+        mask = torch.cat([
+            torch.zeros(batch_size, 1).bool(), tree_mask, data_mask, torch.zeros(batch_size, 1).bool()
+        ], dim=-1)
+        if self._bin_indices is not None:
+            bin_mask = mask[:, self._bin_indices]
+            quantile_mask = mask[:, self._quantile_indices]
+            need_to_swap = bin_mask & ~quantile_mask
+            bin_mask[need_to_swap] = False
+            quantile_mask[need_to_swap] = True
+            mask[:, self._bin_indices] = bin_mask
+            mask[:, self._quantile_indices] = quantile_mask
+        return mask
+
+
+class _DataOffsetter:
+    def __init__(self, data: Dataset, leaf_dim: int, max_n_leaves: int):
+        """
+        Parameters
+        ----------
+        data : Dataset
+            The tabular dataset.
+        leaf_dim : int
+            Number of trees.
+        max_n_leaves : int
+            The maximum number of leaves.
+        """
+        self._cat_offset = 3 + max_n_leaves
+        self._bin_offset = self._cat_offset + data.max_n_categories
+        self._quantile_offset = self._bin_offset + data.max_n_bins
+
+        self.offsets = torch.zeros(2 + leaf_dim + data.index_matrix.shape[-1], dtype=torch.long)
+        self.offsets[-1] = 1
+        self.offsets[1:1 + leaf_dim] = 3
+        for i, (itype, _) in enumerate(data.index_description):
+            pos = i + 1 + leaf_dim
+            offset = self._cat_offset if itype == "cat" else self._quantile_offset \
+                if itype == "quantile" else self._bin_offset
+            self.offsets[pos] = offset
+
+
+class TrainingDataCollator:
+    def __init__(self, data: Dataset, leaf_index_matrix: np.ndarray, max_n_leaves: int,
+                 tree_mask_ratio: Union[float, Tuple[float, float]],
+                 data_mask_ratio: Union[float, Tuple[float, float]]):
+        """
+        Parameters
+        ----------
+        data : Dataset
+            The tabular dataset.
+        leaf_index_matrix : np.ndarray
+            The leaf index matrix from tree-based model.
+        max_n_leaves : int
+            The maximum number of leaves.
+        tree_mask_ratio, data_mask_ratio : float | (float, float)
+            Arguments of `MaskGenerator`.
+        """
+        # bos_id: 0, eos_id: 1, mask_id: 2
+        leaf_dim = leaf_index_matrix.shape[1]
+        self._mask_generator = MaskGenerator(
+            data=data, leaf_dim=leaf_dim, tree_mask_ratio=tree_mask_ratio, data_mask_ratio=data_mask_ratio
+        )
+        self._offsetter = _DataOffsetter(data=data, leaf_dim=leaf_dim, max_n_leaves=max_n_leaves)
+
+    def __call__(self, batch: List[Tuple[torch.LongTensor, torch.LongTensor]]) -> BatchEncoding:
+        tokens = torch.cat([
+            torch.zeros(len(batch), 1).long(),
+            torch.stack([a for a, b in batch]), torch.stack([b for a, b in batch]),
+            torch.zeros(len(batch), 1).long()
+        ], dim=1) + self._offsetter.offsets
+        mask = self._mask_generator.generate_mask(len(batch))
+        masked = torch.masked_fill(tokens, mask, 2)
+        return BatchEncoding({
+            "input_ids": masked,
+            "attention_mask": torch.ones_like(masked, dtype=torch.bool),
+            "labels": tokens
+        })
+
+
+class CausalInferenceDataCollator:
+    def __init__(self, leaf_index_matrix: np.ndarray,
+                 tree_mask_ratio: Union[float, Tuple[float, float]]):
+        """
+        Parameters
+        ----------
+        leaf_index_matrix : np.ndarray
+            The leaf index matrix from tree-based model.
+        tree_mask_ratio : float | (float, float)
+            Mask ratio of tree leaves.
+        """
+        # bos_id: 0, eos_id: 1, mask_id: 2
+        self._leaf_dim = leaf_index_matrix.shape[1]
+        if isinstance(tree_mask_ratio, Sequence):
+            self._min_tree_mask_ratio, self._max_tree_mask_ratio = tree_mask_ratio
+        else:
+            self._min_tree_mask_ratio, self._max_tree_mask_ratio = tree_mask_ratio, tree_mask_ratio
+
+    def __call__(self, batch: List[Tuple[torch.LongTensor, ]]) -> BatchEncoding:
+        tokens = torch.cat([
+            torch.zeros(len(batch), 1).long(),
+            torch.stack([a for a, in batch]) + 3,
+        ], dim=1)
+        tree_threshold = torch.rand(len(batch)) * (
+                self._max_tree_mask_ratio - self._min_tree_mask_ratio
+        ) + self._min_tree_mask_ratio
+        tree_mask = torch.rand(len(batch), self._leaf_dim) < tree_threshold.view(-1, 1)
+        mask = torch.cat([
+            torch.zeros(len(batch), 1).bool(), tree_mask,
+        ], dim=-1)
+
+        masked = torch.masked_fill(tokens, mask, 2)
+        return BatchEncoding({
+            "input_ids": masked,
+            "attention_mask": torch.ones_like(masked, dtype=torch.bool),
+        })
+
+
+class _DataLogitsProcessor(LogitsProcessor):
+    def __init__(self, data: Dataset, max_n_leaves: int, leaf_dim: int):
+        super().__init__()
+        self._data = data
+        self._max_n_leaves = max_n_leaves
+        self._leaf_dim = leaf_dim
+        self._main_dim = data.index_matrix.shape[-1]
+        self._cat_offset = 3 + max_n_leaves
+        self._bin_offset = self._cat_offset + data.max_n_categories
+        self._quantile_offset = self._bin_offset + data.max_n_bins
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if input_ids.shape[-1] == 0:
+            valid_tokens = torch.tensor([0])
+        elif 1 <= input_ids.shape[-1] < self._leaf_dim + 1:
+            valid_tokens = torch.arange(3, 3 + self._max_n_leaves)
+        elif self._leaf_dim + 1 <= input_ids.shape[-1] < self._leaf_dim + 1 + self._main_dim:
+            main_index = input_ids.shape[-1] - self._leaf_dim - 1
+            itype, n_classes = self._data.index_description[main_index]
+            offset = self._cat_offset if itype == "cat" else self._bin_offset\
+                if itype == "bin" else self._quantile_offset
+            valid_tokens = torch.arange(offset, offset + n_classes)
+        else:
+            valid_tokens = torch.tensor([1])
+
+        mask = torch.zeros_like(scores, dtype=torch.bool)
+        mask[:, valid_tokens] = True
+        scores = scores.masked_fill(~mask, -1e9)
+        return scores

tabtreeformer/dsml.py

@@ -0,0 +1,141 @@
+"""Downstream ML Models."""
+from abc import ABC, abstractmethod
+from typing import Optional, Type
+
+import numpy as np
+import pandas as pd
+import optuna
+from lightgbm import LGBMClassifier, LGBMRegressor
+from sklearn.base import BaseEstimator
+from sklearn.model_selection import cross_val_score
+
+from .data import Dataset
+
+
+eval_metrics = {
+    "bin": "f1_weighted",
+    "mult": "f1_weighted",
+    "reg": "neg_mean_squared_error"
+}
+
+
+class MLModel(ABC):
+    def __init__(self, data: Dataset, **kwargs):
+        """
+        Parameters
+        ----------
+        data : Dataset
+            The dataset to be used for training and evaluation.
+        **kwargs
+            Hyperparameter tuning search space. Keys are parameter names, and values are dict with key "dtype",
+            where there are 4 dtypes:
+
+            - "const": fixed constant value, with the value in another key "value",
+            - "categorical": categorical value, other keys and values for `optuna.Trial.suggest_categorical`,
+            - "int": integer value, other keys and values for `optuna.Trial.suggest_int`,
+            - "float": float value, other keys and values for `optuna.Trial.suggest_float`.
+        """
+        self.data = data
+        self._kwargs = kwargs
+        self._fixed_kwargs = {}
+        self._model: Type[BaseEstimator] = self._create_model(data)
+        self._best_params = None
+        self._base_best_model: Optional[BaseEstimator] = None
+
+    def _objective(self, trial: optuna.Trial):
+        params = {}
+        for k, v in self._kwargs.items():
+            v = v.copy()
+            dtype = v.pop("dtype")
+            if dtype == "const":
+                params[k] = v["value"]
+            elif dtype == "categorical":
+                params[k] = trial.suggest_categorical(k, **v)
+            elif dtype == "int":
+                params[k] = trial.suggest_int(k, **v)
+            elif dtype == "float":
+                params[k] = trial.suggest_float(k, **v)
+            else:
+                raise ValueError(f"Unrecognized dtype {dtype}")
+        model = self._model(**params, **self._fixed_kwargs)
+        return cross_val_score(
+            model, self.data.transformed, self.data.y,
+            cv=3, scoring=eval_metrics[self.data.ttype]
+        ).mean()
+
+    def fit(self):
+        """
+        Do hyper-parameter tuning, and fit a best-performing model.
+        """
+        study = optuna.create_study(direction="maximize")
+        study.optimize(self._objective, n_trials=50, n_jobs=10)
+        self._best_params = study.best_params
+        self._base_best_model = self._model(**self._fixed_kwargs, **self._best_params)
+        self._base_best_model.fit(self.data.transformed, self.data.y)
+
+    @classmethod
+    @abstractmethod
+    def _create_model(cls, data: Dataset) -> Type[BaseEstimator]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _predict_leaves(self, x: np.ndarray) -> np.ndarray:
+        raise NotImplementedError()
+
+    def apply(self, data: pd.DataFrame) -> np.ndarray:
+        """
+        Apply the model to obtain tree indices.
+
+        Parameters
+        ----------
+        data : pd.DataFrame
+            The data to obtain tree indices from.
+
+        Returns
+        -------
+        np.ndarray
+            Leaf index matrix for data.
+        """
+        x, _ = self.data.transform(data)
+        return self._predict_leaves(x).astype(np.int32)
+
+    @property
+    def n_leaves(self) -> int:
+        """Maximum number of leaves per tree."""
+        raise NotImplementedError()
+
+
+class LightGBMModel(MLModel):
+    def __init__(self, data: Dataset, ):
+        super().__init__(
+            data,
+            learning_rate=dict(dtype="float", low=0.01, high=0.3, log=True),
+            n_estimators=dict(dtype="int", low=50, high=250, step=50),
+            max_depth=dict(dtype="int", low=3, high=10),
+            num_leaves=dict(dtype="int", low=20, high=100, step=5),
+            min_data_in_leaf=dict(dtype="int", low=10, high=50, step=5),
+            feature_fraction=dict(dtype="float", low=0.6, high=1.0),
+            bagging_fraction=dict(dtype="float", low=0.6, high=1.0),
+            lambda_l1=dict(dtype="float", low=0, high=10),
+            lambda_l2=dict(dtype="float", low=0, high=10),
+        )
+        self._fixed_kwargs = {
+            "categorical_feature": [self.data.columns.index(c) for c in self.data.cat_columns],
+            "verbose": -1,
+            "log_level": "error"
+        }
+
+    @classmethod
+    def _create_model(cls, data: Dataset) -> Type[BaseEstimator]:
+        if data.ttype == "reg":
+            return LGBMRegressor
+        else:
+            return LGBMClassifier
+
+    def _predict_leaves(self, x: np.ndarray) -> np.ndarray:
+        return self._base_best_model.predict(x, pred_leaf=True)
+
+    @property
+    def n_leaves(self) -> int:
+        model_dump = self._base_best_model.booster_.dump_model()
+        return max(tree["num_leaves"] for tree in model_dump["tree_info"])

tabtreeformer/model.py

@@ -0,0 +1,347 @@
+import os
+from typing import Literal, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import (
+    AutoConfig, AutoModelForCausalLM, BatchEncoding, PreTrainedModel, PretrainedConfig,
+    Trainer as _Trainer, TrainingArguments,
+)
+
+from .data import (
+    CausalInferenceDataCollator, Dataset, TrainingDataCollator,
+    _DataLogitsProcessor, _DataOffsetter
+)
+from .dsml import MLModel, LightGBMModel
+
+
+def _prepare_config(config: PretrainedConfig,
+                    hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None,
+                    n_layers: Optional[int] = None, n_heads: Optional[int] = None) -> PretrainedConfig:
+    if hidden_size is not None:
+        for k in ["hidden_size", "n_embd"]:
+            if hasattr(config, k):
+                setattr(config, k, hidden_size)
+    if intermediate_size is not None:
+        for k in ["intermediate_size", "n_inner"]:
+            if hasattr(config, k):
+                setattr(config, k, intermediate_size)
+    if n_layers is not None:
+        for k in ["num_hidden_layers", "n_layer"]:
+            if hasattr(config, k):
+                setattr(config, k, n_layers)
+    if n_heads is not None:
+        for k in ["num_attention_heads", "n_head"]:
+            if hasattr(config, k):
+                setattr(config, k, n_heads)
+    config.bos_token_id = 0
+    config.eos_token_id = 1
+    config.masked_token_id = 2
+    config.pad_token_id = 1
+
+    return config
+
+
+def _update_size(config: PretrainedConfig, vocab_size: int, length: int) -> PretrainedConfig:
+    for k in ["max_position_embeddings", "n_positions", "max_length", "n_ctx"]:
+        if hasattr(config, k):
+            setattr(config, k, length)
+    config.vocab_size = vocab_size
+    return config
+
+
+class TokenLoss(nn.Module):
+    def __init__(self, quantile_offset: int, n_quantiles: torch.LongTensor, max_n_quantiles: int, data_offset: int):
+        """
+        Parameters
+        ----------
+        quantile_offset : int
+            Quantile tokens offset (first quantile token's ID).
+        n_quantiles : torch.LongTensor
+            Number of quantiles at each position.
+        max_n_quantiles : int
+            Maximum number of quantiles.
+        data_offset : int
+            The offset position for data values.
+        """
+        super().__init__()
+        max_n_quantiles = max(1, max_n_quantiles)
+        self.quantile_offset = quantile_offset
+        self.n_quantiles = torch.cat([
+            torch.zeros(data_offset - 1, dtype=torch.long),
+            n_quantiles, torch.zeros(1, dtype=torch.long)
+        ], dim=0).view(-1, 1)
+        self.max_n_quantiles = max_n_quantiles
+        self.data_offset = data_offset
+        self.ignore_index = -100
+        self._is_quantile = self.n_quantiles > 0
+        self._quantile_logits_mask = torch.zeros(
+            self.n_quantiles.shape[0], max_n_quantiles, dtype=torch.bool
+        )
+        for i, q in enumerate(self.n_quantiles):
+            self._quantile_logits_mask[i, :q.item()] = True
+        self._quantile_logits_mask = self._quantile_logits_mask.unsqueeze(0).contiguous()
+        self._weight_matrix = torch.zeros(
+            self.n_quantiles.shape[0], max_n_quantiles, max_n_quantiles, dtype=torch.float
+        )
+        se = ((torch.arange(max_n_quantiles) - torch.arange(max_n_quantiles).view(-1, 1)) ** 2).float()
+        for i, q in enumerate(self.n_quantiles):
+            self._weight_matrix[i] = torch.exp(-se / ((q * 0.005) ** 2))
+        self._weight_matrix = 1 + 0.5 - self._weight_matrix
+        self._weight_matrix = self._weight_matrix.masked_fill(torch.isnan(self._weight_matrix), 0)
+        self._weight_matrix = self._weight_matrix.contiguous()
+
+    def __call__(
+            self, model_output: BatchEncoding, labels: torch.LongTensor, shift_labels: bool = False
+    ) -> torch.FloatTensor:
+        logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0]
+        if shift_labels:
+            logits = logits[..., :-1, :].contiguous()
+            labels = labels[..., 1:].contiguous()
+        probs = nn.functional.softmax(logits, dim=-1)
+        log_probs = -torch.log(probs)
+        if labels.dim() == logits.dim() - 1:
+            labels = labels.unsqueeze(-1)
+
+        padding_mask = labels.eq(self.ignore_index)
+        num_active_elements = padding_mask.numel() - padding_mask.long().sum()
+
+        nll_loss = log_probs.gather(dim=-1, index=labels)
+        self.n_quantiles = self.n_quantiles.to(logits.device)
+        self._is_quantile = self._is_quantile.to(logits.device)
+        self._quantile_logits_mask = self._quantile_logits_mask.to(logits.device)
+        self._weight_matrix = self._weight_matrix.to(logits.device)
+        seq_indices = torch.arange(logits.shape[1], device=logits.device)
+        soft_targets = self._weight_matrix[seq_indices, (labels.squeeze(-1) - self.quantile_offset).clip(0), :]
+        masked = ~self._quantile_logits_mask.repeat((soft_targets.shape[0], 1, 1))
+        soft_targets = soft_targets.masked_fill(masked, 0)
+        soft_targets = soft_targets / soft_targets.sum(dim=-1, keepdim=True)
+        soft_targets = soft_targets.masked_fill(~self._is_quantile.unsqueeze(0), 1 / self.max_n_quantiles)
+        quantile_logits = logits[:, :, self.quantile_offset:self.quantile_offset + self.max_n_quantiles]
+        if quantile_logits.shape[-1] > 0:
+            max_logits = quantile_logits.max(dim=-1, keepdim=True).values
+        else:
+            max_logits = torch.zeros_like(quantile_logits[:, :, :1])
+        stabilized_logits = quantile_logits - max_logits
+        weighted_sum = torch.sum(soft_targets * torch.exp(stabilized_logits), dim=-1, keepdim=True)
+
+        log_normalized = torch.log(soft_targets) + quantile_logits
+        if quantile_logits.shape[-1] > 0:
+            soft_nll_loss = (-log_normalized).gather(
+                dim=-1, index=(labels - self.quantile_offset).clamp(0, self.max_n_quantiles)
+            ) + torch.log(weighted_sum) + max_logits
+        else:
+            soft_nll_loss = torch.zeros_like(weighted_sum)
+        soft_nll_loss = soft_nll_loss.masked_fill(~self._is_quantile.unsqueeze(0), 0)
+        is_quantile = self._is_quantile.unsqueeze(0)
+        is_quantile_labels = (~self._is_quantile).long().unsqueeze(0).repeat((logits.shape[0], 1, 1))
+        quantile_logits_mask = torch.zeros((1, *probs.shape[1:]), dtype=torch.bool, device=logits.device)
+        quantile_logits_mask[
+            0, :, self.quantile_offset:self.quantile_offset + self.max_n_quantiles
+        ] = self._quantile_logits_mask
+        quantile_probs = (probs * quantile_logits_mask).sum(dim=-1).masked_fill(~is_quantile[:, :, 0], 0.5)
+        non_quantile_probs = (probs * ~quantile_logits_mask).sum(dim=-1).masked_fill(~is_quantile[:, :, 0], 0.5)
+        quantile_probs = torch.stack([quantile_probs, non_quantile_probs], dim=-1)
+        quantile_log_probs = -torch.log(quantile_probs)
+        quantile_nll_loss = quantile_log_probs.gather(dim=-1, index=is_quantile_labels)
+        nll_loss = nll_loss * (~is_quantile) + (soft_nll_loss + quantile_nll_loss) * is_quantile
+
+        nll_loss.masked_fill_(padding_mask, 0.0)
+        nll_loss = nll_loss.sum() / num_active_elements
+        return nll_loss
+
+
+class Trainer(_Trainer):
+    def __init__(self, *args, quantile_offset: int, n_quantiles: torch.LongTensor, max_n_quantiles: int,
+                 data_offset: int,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+        self.label_smoother = TokenLoss(
+            quantile_offset, n_quantiles, max_n_quantiles, data_offset,
+        )
+
+
+class Transformer:
+    def __init__(self, model: MLModel, backbone: str = "distilgpt2"):
+        """
+        Model of transformer using tree-based model input and data input.
+
+        Parameters
+        ----------
+        model : MLModel
+            The tree-based model.
+        backbone : str
+            The causal LM backbone from huggingface pre-trained model.
+        """
+        self.model = model
+        self.data = self.model.data
+
+        self.config = _prepare_config(
+            AutoConfig.from_pretrained(backbone), 256, None, None, 8
+        )
+
+        self.leaf_index_matrix: Optional[np.ndarray] = None
+        self.lm: Optional[PreTrainedModel] = None
+        self._batch_size = 0
+        self._seq_len = 0
+
+    def train(self, out_dir: str):
+        """
+        Train model.
+
+        Parameters
+        ----------
+        out_dir : str
+            Output directory.
+        """
+        leaf_index_matrix = self.model.apply(self.data.data)
+        self.leaf_index_matrix = leaf_index_matrix
+        vocab_size = int(
+            self.model.n_leaves + 3 + self.data.max_n_categories + self.data.max_n_bins + self.data.max_n_quantiles
+        )
+        seq_len = int(2 + leaf_index_matrix.shape[-1] + self.data.index_matrix.shape[-1])
+        self._seq_len = seq_len
+        self.config = _update_size(self.config, vocab_size, seq_len)
+
+        self.lm = AutoModelForCausalLM.from_config(self.config)
+
+        dataset = TensorDataset(
+            torch.from_numpy(leaf_index_matrix).long(),
+            torch.from_numpy(self.data.index_matrix).long(),
+        )
+        os.makedirs(out_dir, exist_ok=True)
+
+        collator = TrainingDataCollator(
+            self.data, leaf_index_matrix, self.model.n_leaves, (0.3, 0.7), (0.1, 0.4),
+        )
+        training_args = TrainingArguments(
+            output_dir=os.path.join(out_dir, "ckpt"),
+            logging_dir=os.path.join(out_dir, "logs"),
+            per_device_train_batch_size=128,
+            max_steps=30, fp16=True, learning_rate=5e-4, logging_steps=100
+        )
+        trainer = Trainer(
+            model=self.lm,
+            args=training_args,
+            train_dataset=dataset,
+            data_collator=collator,
+            quantile_offset=3 + self.model.n_leaves + self.data.max_n_categories + self.data.max_n_bins,
+            max_n_quantiles=self.data.max_n_quantiles,
+            n_quantiles=torch.tensor(
+                [0 if t != "quantile" else x for t, x in self.data.index_description], dtype=torch.long
+            ),
+            data_offset=1 + leaf_index_matrix.shape[-1],
+        )
+        self._batch_size = 128
+
+        trainer.train()
+
+        self.lm.save_pretrained(os.path.join(out_dir, "final"))
+        self.lm = os.path.join(out_dir, f"final")
+
+    @torch.no_grad()
+    def sample(
+            self, n: int,
+    ) -> torch.LongTensor:
+        """
+        Sample data.
+
+        Parameters
+        ----------
+        n : int
+            Number of rows to be sampled.
+
+        Returns
+        -------
+        torch.LongTensor
+            The generated token IDs.
+        """
+        inference_dataset = TensorDataset(torch.from_numpy(
+            self.leaf_index_matrix[np.random.randint(low=0, high=self.leaf_index_matrix.shape[0], size=(n,))]
+        ).long())
+        dataloader = DataLoader(
+            inference_dataset, collate_fn=CausalInferenceDataCollator(self.leaf_index_matrix, (0.3, 0.7)),
+            batch_size=self._batch_size, shuffle=False
+        )
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if isinstance(self.lm, str):
+            self.lm = AutoModelForCausalLM.from_pretrained(self.lm)
+        self.lm.to(device)
+        self.lm.eval()
+        generated_texts = []
+        logits_processor = _DataLogitsProcessor(
+            self.data, self.model.n_leaves, self.leaf_index_matrix.shape[-1]
+        )
+        for batch in dataloader:
+            batch = batch.to(device)
+            outputs = self.lm.generate(
+                **batch,
+                max_length=self._seq_len,
+                num_return_sequences=1,
+                do_sample=True,
+                eos_token_id=1,
+                bos_token_id=0,
+                pad_token_id=1,
+                logits_processor=[logits_processor],
+                temperature=0.7
+            )
+
+            generated_texts.append(outputs)
+        out = torch.cat(generated_texts, dim=0)
+        return out
+
+
+class TabTreeFormer:
+    def __init__(self):
+        """
+        TabTreeFormer model.
+        """
+        self.data: Optional[Dataset] = None
+        self.ml_model: Optional[MLModel] = None
+        self.transformer: Optional[Transformer] = None
+        self._offsetter: Optional[_DataOffsetter] = None
+
+    def train(self, data: pd.DataFrame, target: str, ttype: Literal["bin", "mult", "reg"], out_dir: str):
+        """
+        Train a TabTreeFormer.
+
+        Parameters
+        ----------
+        data : pd.DataFrame
+            The data to train the model on.
+        target, ttype
+            Arguments for `data.Dataset`.
+        out_dir : str
+            The output directory.
+        """
+        self.data = Dataset(data, target, ttype)
+        self.ml_model = LightGBMModel(data=self.data)
+        self.ml_model.fit()
+        self.transformer = Transformer(model=self.ml_model)
+        self.transformer.train(out_dir)
+        self._offsetter = _DataOffsetter(
+            self.data, self.transformer.leaf_index_matrix.shape[-1], self.ml_model.n_leaves
+        )
+
+    def sample(self, n: int) -> pd.DataFrame:
+        """
+        Sample data by TabTreeFormer.
+
+        Parameters
+        ----------
+        n : int
+            Number of rows to be sampled.
+
+        Returns
+        -------
+        pd.DataFrame
+            Sampled dataset.
+        """
+        out = self.transformer.sample(n)
+        out = out - self._offsetter.offsets.to(out.device)
+        st = 1 + self.transformer.leaf_index_matrix.shape[-1]
+        ed = st + self.data.index_matrix.shape[-1]
+        return self.data.recover_index_matrix(out[:, st:ed].detach().cpu().numpy())