| from transformers import PreTrainedModel |
| import torch.nn as nn |
| import numpy as np |
| import torch |
| from torch import tensor |
| from torch.nn import TransformerEncoder, TransformerEncoderLayer |
| from .model_config import TabularTransformerConfig |
|
|
| class FleshkaTabularTransformer(PreTrainedModel): |
| config_class = TabularTransformerConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.config = config |
| self.input_proj = nn.Linear(config.input_dim, config.d_model) |
|
|
| encoder_layers = TransformerEncoderLayer( |
| d_model=config.d_model, |
| nhead=config.nhead, |
| dim_feedforward=config.d_model * 4, |
| dropout=config.dropout, |
| activation="gelu", |
| batch_first=True |
| ) |
| self.transformer = TransformerEncoder(encoder_layers, config.num_layers) |
|
|
| self.head = nn.Sequential( |
| nn.LayerNorm(config.d_model), |
| nn.Linear(config.d_model, 1) |
| ) |
|
|
| self._init_weights() |
| |
| def _normalize(self, flist): |
| return (flist - np.array(self.config.mean)) / (np.array(self.config.std) + 1e-8) |
|
|
| def _init_weights(self): |
| for module in self.modules(): |
| if isinstance(module, nn.Linear): |
| nn.init.xavier_normal_(module.weight) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def forward(self, lx: list): |
| out = [] |
| for nx in lx: |
| |
| x = tensor(self._normalize(nx), dtype=self.input_proj.weight.dtype).unsqueeze(0).to(self.device) |
| x = self.input_proj(x) |
| x = x.unsqueeze(1) |
| x = self.transformer(x) |
| x = x.squeeze(1) |
| out.append(self.head(x).item() > 0) |
| return out |
