tabdpt 0.1.0__tar.gz

tabdpt-0.1.0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.3
+Name: tabdpt
+Version: 0.1.0
+Summary: TabDPT
+License: MIT
+Author: Jeremy
+Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: faiss-cpu
+Requires-Dist: gdown
+Requires-Dist: numpy
+Requires-Dist: scikit-learn
+Requires-Dist: torch

tabdpt-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[tool.poetry]
+name = "tabdpt"
+version = "0.1.0"
+description = "TabDPT"
+authors = ["Jeremy", "Valentin"]
+license = "MIT"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+torch = "*"
+numpy = "*"
+scikit-learn = "*"
+faiss-cpu = "*"
+gdown = "*"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"

tabdpt-0.1.0/src/tabdpt/__init__.py

@@ -0,0 +1,4 @@
+from .classifier import TabDPTClassifier
+from .regressor import TabDPTRegressor
+
+__all__ = ['TabDPTClassifier', 'TabDPTRegressor']

tabdpt-0.1.0/src/tabdpt/classifier.py

@@ -0,0 +1,99 @@
+import numpy as np
+import torch
+import math
+from sklearn.base import ClassifierMixin
+from .estimator import TabDPTEstimator
+from .utils import pad_x
+
+
+class TabDPTClassifier(TabDPTEstimator, ClassifierMixin):
+    def __init__(self, path: str = '', inf_batch_size: int = 512, device: str = 'cuda:0', use_flash: bool = True, compile: bool = True):
+        super().__init__(path=path, mode='cls', inf_batch_size=inf_batch_size, device=device, use_flash=use_flash, compile=compile)
+        
+    def fit(self, X, y):
+        super().fit(X, y)
+        self.num_classes = len(np.unique(self.y_train))
+        assert self.num_classes > 1, "Number of classes must be greater than 1"
+        
+    def _predict_large_cls(self, X_train, X_test, y_train):
+        num_digits = math.ceil(math.log(self.num_classes, self.max_num_classes))
+        
+        digit_preds = []
+        for i in range(num_digits):
+            y_train_digit = (y_train // (self.max_num_classes ** i)) % self.max_num_classes
+            pred = self.model(
+                x_src=torch.cat([X_train, X_test], dim=1),
+                y_src=y_train_digit.unsqueeze(-1),
+                task='cls',
+            )
+            digit_preds.append(pred.float())
+
+        full_pred = torch.zeros((X_test.shape[0], X_test.shape[1], self.num_classes), device=X_train.device)
+        for class_idx in range(self.num_classes):
+            class_pred = torch.zeros_like(digit_preds[0][:, :, 0])
+            for digit_idx, digit_pred in enumerate(digit_preds):
+                digit_value = (class_idx // (self.max_num_classes ** digit_idx)) % self.max_num_classes
+                class_pred += digit_pred[:, :, digit_value]
+            full_pred[:, :, class_idx] = class_pred
+
+        return full_pred
+
+    @torch.no_grad()
+    def predict_proba(self, X: np.ndarray, temperature: float = 0.8, context_size: int = 128):
+        train_x, train_y, test_x = self._prepare_prediction(X)
+
+        if context_size >= self.n_instances:
+            X_train = pad_x(train_x[None, :, :], self.max_features).to(self.device)
+            X_test = pad_x(test_x[None, :, :], self.max_features).to(self.device)
+            y_train = train_y[None, :].float()
+            
+            if self.num_classes <= self.max_num_classes:
+                pred = self.model(
+                    x_src=torch.cat([X_train, X_test], dim=1),
+                    y_src=y_train.unsqueeze(-1),
+                    task=self.mode,
+                )
+            else:
+                pred = self._predict_large_cls(X_train, X_test, y_train)
+            
+            pred = pred[..., :self.num_classes] / temperature
+            pred = torch.nn.functional.softmax(pred, dim=-1)
+            return pred.float().squeeze().detach().cpu().numpy()
+        else:
+            pred_list = []
+            for b in range(math.ceil(len(self.X_test) / self.inf_batch_size)):
+                start = b * self.inf_batch_size
+                end = min(len(self.X_test), (b + 1) * self.inf_batch_size)
+
+                indices_nni = self.faiss_knn.get_knn_indices(
+                    self.X_test[start:end], k=context_size
+                )
+                X_nni = train_x[torch.tensor(indices_nni)]
+                y_nni = train_y[torch.tensor(indices_nni)]
+
+                X_nni, y_nni = (
+                    pad_x(torch.Tensor(X_nni), self.max_features).to(self.device),
+                    torch.Tensor(y_nni).to(self.device),
+                )
+                X_eval = test_x[start:end]
+                X_eval = pad_x(X_eval.unsqueeze(1), self.max_features).to(self.device)
+                
+                if self.num_classes <= self.max_num_classes:
+                    pred = self.model(
+                        x_src=torch.cat([X_nni, X_eval], dim=1),
+                        y_src=y_nni.unsqueeze(-1),
+                        task=self.mode,
+                    )
+                else:
+                    pred = self._predict_large_cls(X_nni, X_eval, y_nni)
+
+                pred = pred[..., :self.num_classes].float() / temperature
+                pred = torch.nn.functional.softmax(pred, dim=-1)
+
+                pred_list.append(pred.squeeze())
+
+            
+            return torch.cat(pred_list, dim=0).squeeze().detach().cpu().numpy()
+        
+    def predict(self, X, temperature: float = 0.8, context_size: int = 128):
+        return self.predict_proba(X, temperature=temperature, context_size=context_size).argmax(axis=-1)

tabdpt-0.1.0/src/tabdpt/estimator.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+
+from sklearn.base import BaseEstimator
+from sklearn.utils.validation import check_is_fitted
+from sklearn.impute import SimpleImputer
+from sklearn.preprocessing import StandardScaler
+
+from .model import TabDPTModel
+from .utils import convert_to_torch_tensor, FAISS, download_model
+
+
+class TabDPTEstimator(BaseEstimator):
+    def __init__(self, path: str = '', mode: str = "cls", inf_batch_size: int = 512, device: str = 'cuda:0', use_flash: bool = True, compile: bool = True):
+        self.mode = mode
+        self.inf_batch_size = inf_batch_size
+        self.device = device
+        # automatically download model weight if path is empty
+        if path == '':
+            path = download_model()
+        checkpoint = torch.load(path)
+        self.model = TabDPTModel.load(model_state=checkpoint['model'], config=checkpoint['cfg'], use_flash=use_flash)
+        self.model.eval()
+        self.max_features = self.model.num_features
+        self.max_num_classes = self.model.n_out
+        self.compile = compile
+        assert self.mode in ['cls', 'reg'], "mode must be 'cls' or 'reg'"
+
+    def fit(self, X, y):
+        assert isinstance(X, np.ndarray), "X must be a numpy array"
+        assert isinstance(y, np.ndarray), "y must be a numpy array"
+        assert X.shape[0] == y.shape[0], "X and y must have the same number of samples"
+        assert X.ndim == 2, "X must be a 2D array"
+        assert y.ndim == 1, "y must be a 1D array"
+
+        self.imputer = SimpleImputer(strategy='mean')
+        X = self.imputer.fit_transform(X)
+        self.scaler = StandardScaler()
+        X = self.scaler.fit_transform(X)
+        
+        self.faiss_knn = FAISS(X)
+        self.n_instances, self.n_features = X.shape
+        self.X_train = X
+        self.y_train = y
+        self.is_fitted_ = True
+        if self.compile:
+            self.model = torch.compile(self.model)
+        
+    def _prepare_prediction(self, X: np.ndarray):
+        check_is_fitted(self)
+        self.X_test = self.imputer.transform(X)
+        self.X_test = self.scaler.transform(self.X_test)
+        train_x, train_y, test_x = (
+            convert_to_torch_tensor(self.X_train).to(self.device).float(),
+            convert_to_torch_tensor(self.y_train).to(self.device).float(),
+            convert_to_torch_tensor(self.X_test).to(self.device).float(),
+        )
+
+        # Apply PCA optionally to reduce the number of features
+        if self.n_features > self.max_features:
+            U, S, self.V = torch.pca_lowrank(train_x, q=self.max_features)
+            train_x = train_x @ self.V
+        else:
+            self.V = None
+        
+        test_x = test_x @ self.V if self.V is not None else test_x
+        return train_x, train_y, test_x

tabdpt-0.1.0/src/tabdpt/model.py

@@ -0,0 +1,175 @@
+from typing import Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .utils import normalize_data, clip_outliers, flash_context
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, ff_dim):
+        super().__init__()
+        bias = True  # Set bias=True to match the original model
+        self.embed_dim = embed_dim
+        self.head_dim = embed_dim // num_heads
+        self.num_heads = num_heads
+        self.kv_proj = nn.Linear(embed_dim, 2 * embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.attn_norm = nn.LayerNorm(embed_dim)
+        self.ff_norm = nn.LayerNorm(embed_dim)
+        self.ff = nn.Sequential(
+            nn.Linear(embed_dim, ff_dim),
+            nn.GELU(),
+            nn.Linear(ff_dim, embed_dim)
+        )
+
+    def forward(self, x, eval_pos):
+        B, L, _ = x.size()
+        h = self.attn_norm(x)
+        q = self.q_proj(h)
+        k, v = self.kv_proj(h[:, :eval_pos]).chunk(2, dim=-1)
+        q = q.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, eval_pos, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, eval_pos, self.num_heads, self.head_dim).transpose(1, 2)
+        attn = F.scaled_dot_product_attention(q, k, v).transpose(1, 2)
+        attn = self.out_proj(attn.reshape(B, L, self.embed_dim))
+        x = x + attn
+        x = x + self.ff(self.ff_norm(x))
+        return x
+
+class TabDPTModel(nn.Module):
+    def __init__(self, dropout: float, n_out: int, nhead: int, nhid: int, ninp: int, nlayers: int, norm_first: bool, num_features: int, use_flash: bool):
+        super().__init__()
+        self.n_out = n_out
+        self.num_features = num_features
+        self.encoder = nn.Linear(num_features, ninp)
+        self.y_encoder = nn.Linear(1, ninp)
+        self.cls_head = nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, n_out))
+        self.reg_head = nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, 1))
+        self.task2head = {'cls': self.cls_head, 'reg': self.reg_head}
+        self.transformer_encoder = nn.ModuleList(
+            [
+                TransformerEncoderLayer(embed_dim=ninp, num_heads=nhead, ff_dim=nhid)
+                for _ in range(nlayers)
+            ]
+        )
+        self.use_flash = torch.cuda.is_available() and use_flash
+
+    @flash_context
+    def forward(
+        self,
+        x_src: torch.Tensor,
+        y_src: torch.Tensor,
+        task: Literal["cls", "reg"],  # classification or regression
+    ) -> torch.Tensor:
+        eval_pos = y_src.shape[1]
+        x_src = normalize_data(x_src, -1 if self.training else eval_pos)
+        
+        x_src = clip_outliers(x_src, -1 if self.training else eval_pos, n_sigma=10)
+        if task == "reg":
+            y_src, mean_y, std_y = normalize_data(y_src, return_mean_std=True)
+            y_src = clip_outliers(y_src)
+        
+        x_src = torch.nan_to_num(x_src, nan=0)
+
+        x_src = self.encoder(x_src)
+
+        mean = (x_src**2).mean(dim=-1, keepdim=True)
+        rms = torch.sqrt(mean)
+        x_src = x_src / rms
+
+        y_src = self.y_encoder(y_src)
+        train_x = x_src[:, :eval_pos] + y_src
+        src = torch.cat([train_x, x_src[:, eval_pos:]], 1)
+
+        for layer in self.transformer_encoder:
+            src = layer(src, eval_pos)
+        pred = self.task2head[task](src)
+
+        if task == "reg":
+            pred = pred * std_y + mean_y
+
+        return pred[:, eval_pos:]
+
+    @classmethod
+    def load(cls, model_state, config, use_flash):
+        model = TabDPTModel(
+            dropout=config['training']['dropout'],
+            n_out=config['model']['max_num_classes'],
+            nhead=config['model']['nhead'],
+            nhid=config['model']['emsize'] * config['model']['nhid_factor'],
+            ninp=config['model']['emsize'],
+            nlayers=config['model']['nlayers'],
+            norm_first=config['model']['norm_first'],
+            num_features=config['model']['max_num_features'],
+            use_flash=use_flash
+        )
+
+        # Remove any module prefixes if necessary
+        module_prefix = '_orig_mod.'
+        model_state = {k.replace(module_prefix, ''): v for k, v in model_state.items()}
+
+        # Mapping function to convert state_dict keys
+        def map_state_dict(original_state_dict, model):
+            new_state_dict = {}
+            for key, value in original_state_dict.items():
+                if key.startswith('transformer_encoder.'):
+                    # Handle transformer encoder layers
+                    parts = key.split('.')
+                    layer_idx = parts[1]
+                    sub_module = parts[2]
+                    param_name = '.'.join(parts[3:])
+                    if sub_module == 'self_attn':
+                        if param_name == 'in_proj_weight':
+                            in_proj_weight = value
+                            embed_dim = model.transformer_encoder[int(layer_idx)].embed_dim
+                            q_proj_weight = in_proj_weight[:embed_dim, :]
+                            k_proj_weight = in_proj_weight[embed_dim:2*embed_dim, :]
+                            v_proj_weight = in_proj_weight[2*embed_dim:, :]
+                            kv_proj_weight = torch.cat([k_proj_weight, v_proj_weight], dim=0)
+                            new_state_dict[f'transformer_encoder.{layer_idx}.q_proj.weight'] = q_proj_weight
+                            new_state_dict[f'transformer_encoder.{layer_idx}.kv_proj.weight'] = kv_proj_weight
+                        elif param_name == 'in_proj_bias':
+                            in_proj_bias = value
+                            embed_dim = model.transformer_encoder[int(layer_idx)].embed_dim
+                            q_proj_bias = in_proj_bias[:embed_dim]
+                            k_proj_bias = in_proj_bias[embed_dim:2*embed_dim]
+                            v_proj_bias = in_proj_bias[2*embed_dim:]
+                            kv_proj_bias = torch.cat([k_proj_bias, v_proj_bias], dim=0)
+                            new_state_dict[f'transformer_encoder.{layer_idx}.q_proj.bias'] = q_proj_bias
+                            new_state_dict[f'transformer_encoder.{layer_idx}.kv_proj.bias'] = kv_proj_bias
+                        elif param_name == 'out_proj.weight':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.out_proj.weight'] = value
+                        elif param_name == 'out_proj.bias':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.out_proj.bias'] = value
+                    elif sub_module == 'linear1':
+                        if param_name == 'weight':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff.0.weight'] = value
+                        elif param_name == 'bias':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff.0.bias'] = value
+                    elif sub_module == 'linear2':
+                        if param_name == 'weight':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff.2.weight'] = value
+                        elif param_name == 'bias':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff.2.bias'] = value
+                    elif sub_module == 'norm1':
+                        if param_name == 'weight':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.attn_norm.weight'] = value
+                        elif param_name == 'bias':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.attn_norm.bias'] = value
+                    elif sub_module == 'norm2':
+                        if param_name == 'weight':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff_norm.weight'] = value
+                        elif param_name == 'bias':
+                            new_state_dict[f'transformer_encoder.{layer_idx}.ff_norm.bias'] = value
+                else:
+                    # Copy other parameters directly
+                    new_state_dict[key] = value
+            return new_state_dict
+
+        # Map the state_dict to the new model
+        new_state_dict = map_state_dict(model_state, model)
+        model.load_state_dict(new_state_dict)
+        model.to(config['env']['device'])
+        model.eval()
+        return model

tabdpt-0.1.0/src/tabdpt/regressor.py

@@ -0,0 +1,52 @@
+import math
+import torch
+import numpy as np
+from sklearn.base import RegressorMixin
+from .estimator import TabDPTEstimator
+from .utils import pad_x
+
+class TabDPTRegressor(TabDPTEstimator, RegressorMixin):
+    def __init__(self, path: str = '', inf_batch_size: int = 512, device: str = 'cuda:0', use_flash: bool = True, compile: bool = True):
+        super().__init__(path=path, mode='reg', inf_batch_size=inf_batch_size, device=device, use_flash=use_flash, compile=compile)
+
+    @torch.no_grad()
+    def predict(self, X: np.ndarray, context_size: int = 128):
+        train_x, train_y, test_x = self._prepare_prediction(X)
+        if context_size >= self.n_instances:
+            X_train = pad_x(train_x[None, :, :], self.max_features).to(self.device)
+            X_test = pad_x(test_x[None, :, :], self.max_features).to(self.device)
+            y_train = train_y[None, :].float()
+            pred = self.model(
+                x_src=torch.cat([X_train, X_test], dim=1),
+                y_src=y_train.unsqueeze(-1),
+                task=self.mode,
+            )
+            
+            return pred.float().squeeze().detach().cpu().numpy()
+        else:
+            pred_list = []
+            for b in range(math.ceil(len(self.X_test) / self.inf_batch_size)):
+                start = b * self.inf_batch_size
+                end = min(len(self.X_test), (b + 1) * self.inf_batch_size)
+
+                indices_nni = self.faiss_knn.get_knn_indices(
+                    self.X_test[start:end], k=context_size
+                )
+                X_nni = train_x[torch.tensor(indices_nni)]
+                y_nni = train_y[torch.tensor(indices_nni)]
+
+                X_nni, y_nni = (
+                    pad_x(torch.Tensor(X_nni), self.max_features).to(self.device),
+                    torch.Tensor(y_nni).to(self.device),
+                )
+                X_eval = test_x[start:end]
+                X_eval = pad_x(X_eval.unsqueeze(1), self.max_features).to(self.device)
+                pred = self.model(
+                    x_src=torch.cat([X_nni, X_eval], dim=1),
+                    y_src=y_nni.unsqueeze(-1),
+                    task=self.mode,
+                )
+
+                pred_list.append(pred)
+
+            return torch.cat(pred_list).squeeze().detach().cpu().numpy()

tabdpt-0.1.0/src/tabdpt/utils.py

@@ -0,0 +1,108 @@
+import os
+import random
+from functools import wraps
+import tempfile
+
+import numpy as np
+import torch
+import faiss
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+def download_model():
+    temp_dir = "/tmp/tabdpt_model"
+    model_path = os.path.join(temp_dir, "tabdpt.pth")
+    if not os.path.exists(temp_dir):
+        os.makedirs(temp_dir, exist_ok=True)
+        os.system(f"gdown --id 1v-kAFXMaBWmK1Kk6hLaDDlckdYLTCfV1 -O {model_path}")
+    return model_path
+
+def flash_context(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if getattr(self, "use_flash", False):
+            assert torch.cuda.is_available(), "FlashAttention requires CUDA support"
+            bf_support = torch.cuda.get_device_capability()[0] >= 8
+            dtype = torch.bfloat16 if bf_support else torch.float16
+            with torch.autocast(device_type='cuda', dtype=dtype), sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+                return func(self, *args, **kwargs)
+        else:
+            return func(self, *args, **kwargs)
+    return wrapper
+
+def maskmean(x, mask, dim):
+    x = torch.where(mask, x, 0)
+    return x.sum(dim=dim, keepdim=True) / mask.sum(dim=dim, keepdim=True)
+
+
+def maskstd(x, mask, dim=1):
+    num = mask.sum(dim=dim, keepdim=True)
+    mean = maskmean(x, mask, dim=dim)
+    diffs = torch.where(mask, mean - x, 0)
+    return ((diffs**2).sum(dim=dim, keepdim=True) / (num - 1)) ** 0.5
+
+
+def normalize_data(data, eval_pos=-1, dim=1, return_mean_std: bool = False):
+    X = data[:, :eval_pos] if eval_pos > 0 else data
+    mask = ~torch.isnan(X)
+    mean = maskmean(X, mask, dim=dim)
+    std = maskstd(X, mask, dim=dim) + 1e-6
+    data = (data - mean) / std
+    if return_mean_std:
+        return data, mean, std
+    return data
+
+
+def clip_outliers(data, eval_pos=-1, n_sigma=4, dim=1):
+    X = data[:, :eval_pos] if eval_pos > 0 else data
+    mask = ~torch.isnan(X)
+    mean = maskmean(X, mask, dim=dim)
+    cutoff = n_sigma * maskstd(X, mask, dim=dim)
+    mask &= cutoff >= torch.abs(X - mean)
+    cutoff = n_sigma * maskstd(X, mask, dim=dim)
+    return torch.clip(data, mean - cutoff, mean + cutoff)
+
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+
+
+def convert_to_torch_tensor(input):
+    if isinstance(input, np.ndarray):
+        return torch.from_numpy(input)
+    elif torch.is_tensor(input):
+        return input
+    else:
+        raise TypeError("Input must be a NumPy array or a PyTorch tensor.")
+
+
+def pad_x(X: torch.Tensor, num_features=100):
+    if num_features is None:
+        return X
+    n_features = X.shape[-1]
+    zero_feature_padding = torch.zeros((*X.shape[:-1], num_features - n_features), device=X.device)
+    return torch.cat([X, zero_feature_padding], dim=-1)
+
+
+class FAISS:
+    def __init__(self, X):
+        assert isinstance(X, np.ndarray), "X must be a numpy array"
+        X = np.ascontiguousarray(X)
+        X = X.astype(np.float32)
+        self.index = faiss.IndexFlatL2(X.shape[1])
+        self.index.add(X)
+
+    def get_knn_indices(self, queries, k):
+        if isinstance(queries, torch.Tensor):
+            queries = queries.cpu().numpy()
+        queries = np.ascontiguousarray(queries)
+        assert isinstance(k, int)
+
+        knns = self.index.search(queries, k)
+        indices_Xs = knns[1]
+        return indices_Xs