autogluon.tabular 1.2.1b20250407__py3-none-any.whl → 1.2.1b20250409__py3-none-any.whl

autogluon/tabular/models/tab_transformer/modified_transformer.py

@@ -1,494 +0,0 @@
-"""
-This code is a modification of the official PyTorch Transformer code found at:
-https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/transformer.py
-
-The modification allows the option of fixing the attention map
-"""
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-
-# Needed for pytorch 1.7 and 1.2
-try:
-    from torch.overrides import handle_torch_function, has_torch_function
-# Needed for pytorch 1.6
-except ImportError:
-    from torch._overrides import handle_torch_function, has_torch_function
-
-from torch.nn import Module, init
-from torch.nn.functional import dropout, linear, softmax
-from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
-from torch.nn.modules.dropout import Dropout
-from torch.nn.modules.normalization import LayerNorm
-from torch.nn.parameter import Parameter
-
-
-def multi_head_attention_forward(
-    self,
-    query,  # type: Tensor
-    key,  # type: Tensor
-    value,  # type: Tensor
-    embed_dim_to_check,  # type: int
-    num_heads,  # type: int
-    in_proj_weight,  # type: Tensor
-    in_proj_bias,  # type: Tensor
-    bias_k,  # type: Optional[Tensor]
-    bias_v,  # type: Optional[Tensor]
-    add_zero_attn,  # type: bool
-    dropout_p,  # type: float
-    out_proj_weight,  # type: Tensor
-    out_proj_bias,
-    fixed_k=None,  # type: Tensor
-    fixed_q=None,  # type: Tensor
-    training=True,  # type: bool
-    key_padding_mask=None,  # type: Optional[Tensor]
-    need_weights=True,  # type: bool
-    attn_mask=None,  # type: Optional[Tensor]
-    use_separate_proj_weight=False,  # type: bool
-    q_proj_weight=None,  # type: Optional[Tensor]
-    k_proj_weight=None,  # type: Optional[Tensor]
-    v_proj_weight=None,  # type: Optional[Tensor]
-    static_k=None,  # type: Optional[Tensor]
-    static_v=None,  # type: Optional[Tensor]
-):
-    # type: (...) -> tuple[Tensor, Optional[Tensor]]
-    """
-    Args:
-        query, key, value: map a query and a set of key-value pairs to an output.
-            See "Attention Is All You Need" for more details.
-        embed_dim_to_check: total dimension of the model.
-        num_heads: parallel attention heads.
-        in_proj_weight, in_proj_bias: input projection weight and bias.
-        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        dropout_p: probability of an element to be zeroed.
-        out_proj_weight, out_proj_bias: the output projection weight and bias.
-        training: apply dropout if is ``True``.
-        key_padding_mask: if provided, specified padding elements in the key will
-            be ignored by the attention. This is an binary mask. When the value is True,
-            the corresponding value on the attention layer will be filled with -inf.
-        need_weights: output attn_output_weights.
-        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-        use_separate_proj_weight: the function accept the proj. weights for query, key,
-            and value in different forms. If false, in_proj_weight will be used, which is
-            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
-        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
-        static_k, static_v: static key and value used for attention operators.
-    Shape:
-        Inputs:
-        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
-          will be unchanged. If a BoolTensor is provided, the positions with the
-          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
-          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
-          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
-          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-          is provided, it will be added to the attention weight.
-        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-        Outputs:
-        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-          E is the embedding dimension.
-        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-          L is the target sequence length, S is the source sequence length.
-    """
-
-    if not torch.jit.is_scripting():
-        tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
-        if any([type(t) is not torch.Tensor for t in tens_ops]) and has_torch_function(tens_ops):
-            return handle_torch_function(
-                self.multi_head_attention_forward,
-                tens_ops,
-                query,
-                key,
-                value,
-                embed_dim_to_check,
-                num_heads,
-                in_proj_weight,
-                in_proj_bias,
-                bias_k,
-                bias_v,
-                add_zero_attn,
-                dropout_p,
-                out_proj_weight,
-                out_proj_bias,
-                training=training,
-                key_padding_mask=key_padding_mask,
-                need_weights=need_weights,
-                attn_mask=attn_mask,
-                use_separate_proj_weight=use_separate_proj_weight,
-                q_proj_weight=q_proj_weight,
-                k_proj_weight=k_proj_weight,
-                v_proj_weight=v_proj_weight,
-                static_k=static_k,
-                static_v=static_v,
-            )
-    tgt_len, bsz, embed_dim = query.size()
-
-    assert embed_dim == embed_dim_to_check
-    # allow MHA to have different sizes for the feature dimension
-    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
-
-    head_dim = embed_dim // num_heads
-    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
-    scaling = float(head_dim) ** -0.5
-
-    # self-attention
-    # q = linear(query, in_proj_weight, in_proj_bias) #.chunk(1, dim=-1)
-    # we assume we are in the case key==value==query
-    v = linear(query, in_proj_weight, in_proj_bias)  # .chunk(2, dim=-1)
-
-    k = torch.cat([fixed_k.unsqueeze(1) for _ in range(key.shape[1])], dim=1)
-    q = torch.cat([fixed_q.unsqueeze(1) for _ in range(key.shape[1])], dim=1)
-    q = q * scaling
-
-    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
-    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-
-    src_len = k.size(1)
-
-    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
-
-    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
-
-    attn_output_weights = softmax(attn_output_weights, dim=-1)
-    attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training)
-    attn_output = torch.bmm(attn_output_weights, v)
-
-    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
-    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
-    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
-
-    if need_weights:
-        # average attention weights over heads
-        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-        return attn_output, attn_output_weights.sum(dim=1) / num_heads
-    else:
-        return attn_output, None
-
-
-class MultiheadAttention(Module):
-    """Allows the model to jointly attend to information
-    from different representation subspaces.
-    See reference: Attention Is All You Need
-
-    .. math::
-        text{MultiHead}(Q, K, V) = text{Concat}(head_1,dots,head_h)W^O
-        text{where} head_i = text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
-
-    Args:
-        embed_dim: total dimension of the model.
-        num_heads: parallel attention heads.
-        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
-        bias: add bias as module parameter. Default: True.
-        add_bias_kv: add bias to the key and value sequences at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        kdim: total number of features in key. Default: None.
-        vdim: total number of features in value. Default: None.
-
-        Note: if kdim and vdim are None, they will be set to embed_dim such that
-        query, key, and value have the same number of features.
-
-    Examples::
-
-        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
-        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
-    """
-
-    __annotations__ = {
-        "bias_k": torch._jit_internal.Optional[torch.Tensor],
-        "bias_v": torch._jit_internal.Optional[torch.Tensor],
-    }
-
-    def __init__(self, embed_dim, n_cat_embeddings, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
-        super().__init__()
-
-        self.embed_dim = embed_dim
-        self.kdim = kdim if kdim is not None else embed_dim
-        self.vdim = vdim if vdim is not None else embed_dim
-        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
-
-        self.num_heads = num_heads
-        self.dropout = dropout
-        self.head_dim = embed_dim // num_heads
-        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
-
-        if self._qkv_same_embed_dim is False:
-            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
-            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
-            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
-            self.register_parameter("in_proj_weight", None)
-            self.register_parameter("fixed_k", None)
-        else:
-            self.in_proj_weight = Parameter(torch.empty(embed_dim, embed_dim))
-            # self.in_proj_weight = Parameter(torch.empty(2 * embed_dim, embed_dim))
-            self.fixed_k = Parameter(torch.empty(n_cat_embeddings, embed_dim))
-            self.fixed_q = Parameter(torch.empty(n_cat_embeddings, embed_dim))
-            self.register_parameter("q_proj_weight", None)
-            self.register_parameter("k_proj_weight", None)
-            self.register_parameter("v_proj_weight", None)
-
-        if bias:
-            self.in_proj_bias = Parameter(torch.empty(embed_dim))
-        else:
-            self.register_parameter("in_proj_bias", None)
-        self.out_proj = _LinearWithBias(embed_dim, embed_dim)
-
-        if add_bias_kv:
-            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
-        else:
-            self.bias_k = self.bias_v = None
-
-        self.add_zero_attn = add_zero_attn
-
-        self._reset_parameters()
-
-    def _reset_parameters(self):
-        if self._qkv_same_embed_dim:
-            xavier_uniform_(self.in_proj_weight)
-            xavier_uniform_(self.fixed_k)
-            xavier_uniform_(self.fixed_q)
-        else:
-            xavier_uniform_(self.q_proj_weight)
-            xavier_uniform_(self.k_proj_weight)
-            xavier_uniform_(self.v_proj_weight)
-
-        if self.in_proj_bias is not None:
-            constant_(self.in_proj_bias, 0.0)
-            constant_(self.out_proj.bias, 0.0)
-        if self.bias_k is not None:
-            xavier_normal_(self.bias_k)
-        if self.bias_v is not None:
-            xavier_normal_(self.bias_v)
-
-    def __setstate__(self, state):
-        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
-        if "_qkv_same_embed_dim" not in state:
-            state["_qkv_same_embed_dim"] = True
-
-        super().__setstate__(state)
-
-    def forward(self, query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None):
-        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> tuple[Tensor, Optional[Tensor]]
-        """
-        Args:
-            query, key, value: map a query and a set of key-value pairs to an output.
-                See "Attention Is All You Need" for more details.
-            key_padding_mask: if provided, specified padding elements in the key will
-                be ignored by the attention. When given a binary mask and a value is True,
-                the corresponding value on the attention layer will be ignored. When given
-                a byte mask and a value is non-zero, the corresponding value on the attention
-                layer will be ignored
-            need_weights: output attn_output_weights.
-            attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-                the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-
-        Shape:
-            - Inputs:
-            - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-              the embedding dimension.
-            - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-              the embedding dimension.
-            - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-              the embedding dimension.
-            - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-              If a ByteTensor is provided, the non-zero positions will be ignored while the position
-              with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
-              value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-            - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
-              3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
-              S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
-              positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-              while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-              is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-              is provided, it will be added to the attention weight.
-
-            - Outputs:
-            - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-              E is the embedding dimension.
-            - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-              L is the target sequence length, S is the source sequence length.
-        """
-
-        return multi_head_attention_forward(
-            self,
-            query=query,
-            key=key,
-            value=value,
-            embed_dim_to_check=self.embed_dim,
-            num_heads=self.num_heads,
-            in_proj_weight=self.in_proj_weight,
-            in_proj_bias=self.in_proj_bias,
-            bias_k=self.bias_k,
-            bias_v=self.bias_v,
-            add_zero_attn=self.add_zero_attn,
-            dropout_p=self.dropout,
-            out_proj_weight=self.out_proj.weight,
-            out_proj_bias=self.out_proj.bias,
-            fixed_k=self.fixed_k,
-            fixed_q=self.fixed_q,
-            training=self.training,
-            key_padding_mask=key_padding_mask,
-            need_weights=need_weights,
-            attn_mask=attn_mask,
-        )
-
-
-class Linear(Module):
-    r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`
-    Args:
-        in_features: size of each input sample
-        out_features: size of each output sample
-        bias: If set to ``False``, the layer will not learn an additive bias.
-            Default: ``True``
-    Shape:
-        - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of
-          additional dimensions and :math:`H_{in} = \text{in\_features}`
-        - Output: :math:`(N, *, H_{out})` where all but the last dimension
-          are the same shape as the input and :math:`H_{out} = \text{out\_features}`.
-    Attributes:
-        weight: the learnable weights of the module of shape
-            :math:`(\text{out\_features}, \text{in\_features})`. The values are
-            initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
-            :math:`k = \frac{1}{\text{in\_features}}`
-        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
-                If :attr:`bias` is ``True``, the values are initialized from
-                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
-                :math:`k = \frac{1}{\text{in\_features}}`
-    Examples::
-        >>> m = nn.Linear(20, 30)
-        >>> input = torch.randn(128, 20)
-        >>> output = m(input)
-        >>> print(output.size())
-        torch.Size([128, 30])
-    """
-
-    __constants__ = ["in_features", "out_features"]
-    in_features: int
-    out_features: int
-    weight: torch.Tensor
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
-        super().__init__()
-
-        self.in_features = in_features
-        self.out_features = out_features
-        self.weight = Parameter(torch.Tensor(out_features, in_features))
-        if bias:
-            self.bias = Parameter(torch.Tensor(out_features))
-        else:
-            self.register_parameter("bias", None)
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-        if self.bias is not None:
-            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
-            bound = 1 / math.sqrt(fan_in)
-            init.uniform_(self.bias, -bound, bound)
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        return F.linear(input, self.weight, self.bias)
-
-    def extra_repr(self) -> str:
-        return "in_features={}, out_features={}, bias={}".format(self.in_features, self.out_features, self.bias is not None)
-
-
-# This class exists solely for Transformer; it has an annotation stating
-# that bias is never None, which appeases TorchScript
-class _LinearWithBias(Linear):
-    bias: torch.Tensor
-
-    def __init__(self, in_features: int, out_features: int) -> None:
-        super().__init__(in_features, out_features, bias=True)
-
-
-class TransformerEncoderLayerModified(Module):
-    """TransformerEncoderLayer is made up of self-attn and feedforward network.
-    This standard encoder layer is based on the paper "Attention Is All You Need".
-    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
-    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
-    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
-    in a different way during application.
-
-    Args:
-        d_model: the number of expected features in the input (required).
-        nhead: the number of heads in the multiheadattention models (required).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-        activation: the activation function of intermediate layer, relu or gelu (default=relu).
-
-    Examples::
-        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
-        >>> src = torch.rand(10, 32, 512)
-        >>> out = encoder_layer(src)
-    """
-
-    def __init__(self, d_model, n_cat_embeddings, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
-        super().__init__()
-
-        self.self_attn = MultiheadAttention(d_model, n_cat_embeddings, nhead, dropout=dropout)
-
-        # Implementation of Feedforward model
-        self.linear1 = Linear(d_model, dim_feedforward)
-        self.dropout = Dropout(dropout)
-        self.linear2 = Linear(dim_feedforward, d_model)
-
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.dropout1 = Dropout(dropout)
-        self.dropout2 = Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-    def __setstate__(self, state):
-        if "activation" not in state:
-            state["activation"] = F.relu
-        super().__setstate__(state)
-
-    def forward(self, src: torch.Tensor, src_mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        """Pass the input through the encoder layer.
-
-        Args:
-            src: the sequence to the encoder layer (required).
-            src_mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-
-        Shape:
-            see the docs in Transformer class.
-        """
-        src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
-        src = src + self.dropout1(src2)
-        src = self.norm1(src)
-        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
-        src = src + self.dropout2(src2)
-        src = self.norm2(src)
-        return src
-
-
-def _get_activation_fn(activation):
-    if activation == "relu":
-        return F.relu
-    elif activation == "gelu":
-        return F.gelu
-
-    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

autogluon/tabular/models/tab_transformer/pretexts.py

@@ -1,150 +0,0 @@
-from copy import deepcopy
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-from autogluon.core.constants import REGRESSION
-
-"""
-possible TODO: although there is a supervised pretext option below, i.e. pretrain using
-labeled data presumably from a different but related task, currently this cannot be handled 
-by the AG API. One simple way to handle it would be to have an optional "pretext label" 
-argument to task.fit that tells the pretraining module to use the supervised pretext task
-with that particular column as the label.
-"""
-
-
-class SupervisedPretext(nn.Module):
-    # holder class to handle supervised pretraining.
-    def __init__(self, problem_type, device):
-        super().__init__()
-        self.device = device
-        self.loss_funct = nn.MSELoss() if problem_type == REGRESSION else nn.CrossEntropyLoss()
-
-    def forward(self, out, target):
-        loss = self.loss_funct(out, target)
-        pred = out.max(dim=1, keepdim=True)[1]
-        correct = pred.eq(target.view_as(pred)).sum()
-        correct = correct.float()
-        correct = correct / out.shape[0]
-        return loss, correct
-
-    def get(self, data, target):
-        data = data.to(self.device, non_blocking=True)
-        target = target.to(self.device, non_blocking=True)
-
-        return data, target
-
-
-class BERTPretext(nn.Module):
-    """
-    This is the current default pretext task module.
-
-    Functionality:
-
-        self.get:
-            inputs: (data, target) (target will often be None)
-            outputs: (pretext_data, pretext_target)
-
-            called before the forward pass through the TabTransformer to create the
-            input and label for a BERT-style pretext task.
-
-        self.forward:
-            inputs: out (embedding for TabTransformer), target (pretext task label)
-            outputs: loss, % accuracy on pretext task
-
-            given the embedding it passes it through a classifier which learns to
-            predict the pretext label.
-    """
-
-    def __init__(self, cat_feat_origin_cards, device, hidden_dim, replacement_noise="random", p_replace=0.3):
-        super().__init__()
-        self.cat_feat_origin_cards = cat_feat_origin_cards
-        self.device = device
-        self.hidden_dim = hidden_dim
-        self.loss_funct = nn.CrossEntropyLoss()
-        self.p_replace = p_replace
-        self.predicters = nn.ModuleList()
-        self.n_cat_feats = len(cat_feat_origin_cards)
-        self.replacement_noise = replacement_noise
-
-        for col in range(self.n_cat_feats):
-            lin = nn.Linear(self.hidden_dim, 2)
-            self.predicters.append(lin)
-
-    def forward(self, out, target):
-        prob = torch.cat([self.predicters[col](out[:, col, :]).unsqueeze(1) for col in range(self.n_cat_feats)], dim=1)
-        prob = prob.view(-1, 2)
-        target = target.view(-1)
-
-        loss = self.loss_funct(prob, target)
-        pred = prob.max(dim=1, keepdim=True)[1]
-        correct = pred.eq(target.view_as(pred)).sum()
-        correct = correct.float()
-        correct = correct / pred.shape[0]
-
-        return loss, correct
-
-    def get(self, data, target):
-        cat_feats = data
-
-        orig_cat_feats = deepcopy(cat_feats.detach())
-
-        if self.replacement_noise == "swap":
-            n_cat = cat_feats.shape[1]
-            cols_to_shuffle = np.random.choice(n_cat, int(self.p_replace * n_cat), replace=False)
-            for col in cols_to_shuffle:
-                cat_feats[:, col] = cat_feats[:, col][torch.randperm(cat_feats.shape[0])]
-
-        elif self.replacement_noise == "random":
-            locs_to_replace = torch.empty_like(cat_feats, dtype=float).uniform_() < self.p_replace
-            col_cardinalities = torch.LongTensor([i[1] for i in self.cat_feat_origin_cards]).to(cat_feats)
-            col_cardinalities = col_cardinalities.unsqueeze(0).expand_as(cat_feats)
-
-            unif = torch.rand(cat_feats.shape, device=col_cardinalities.device)
-            random_feats = (unif * col_cardinalities).floor().to(torch.int64) + 1  # + 1 since 0 is the padding value
-
-            extra_replace = torch.mul((cat_feats == random_feats).to(int), locs_to_replace.to(int)).to(torch.bool)
-            cat_feats[locs_to_replace] = random_feats[locs_to_replace]
-
-            assert torch.all(cat_feats[extra_replace] == orig_cat_feats[extra_replace]).item() is True
-            extra_plus1 = cat_feats[extra_replace] + 1
-            extra_minus1 = cat_feats[extra_replace] - 1
-            extra_zero_pad_idx = extra_minus1 == 0
-            extra_minus1[extra_zero_pad_idx] = extra_plus1[extra_zero_pad_idx]
-
-            cat_feats[extra_replace] = extra_minus1
-            assert torch.all(~(cat_feats[extra_replace] == orig_cat_feats[extra_replace])).item() is True
-
-        elif self.replacement_noise == "low_rank":
-            assert self.p_replace + 0.2 <= 1, "p_replace too big, lower it!"
-            weights = torch.tensor([self.p_replace, 0.1, 0.9 - self.p_replace], dtype=torch.float)  # 0=pad, 1=replace with random value, 2=dont change
-
-            locs_to_change = torch.multinomial(weights, np.prod(cat_feats.shape), replacement=True).view(cat_feats.shape)
-            col_cardinalities = torch.LongTensor([i[1] for i in self.cat_feat_origin_cards]).to(cat_feats)
-            col_cardinalities = col_cardinalities.unsqueeze(0).expand_as(cat_feats)
-
-            unif = torch.rand(cat_feats.shape, device=col_cardinalities.device)
-            random_feats = (unif * col_cardinalities).floor().to(torch.int64) + 1  # + 1 since 0 is the padding value
-
-            extra_replace = torch.mul((cat_feats == random_feats).to(int), (locs_to_change == 1).to(int)).to(torch.bool)
-            cat_feats[locs_to_change == 1] = random_feats[locs_to_change == 1]
-            cat_feats[locs_to_change == 0] = 0
-
-            assert torch.all(cat_feats[extra_replace] == orig_cat_feats[extra_replace]).item() is True
-            extra_plus1 = cat_feats[extra_replace] + 1
-            extra_minus1 = cat_feats[extra_replace] - 1
-            extra_zero_pad_idx = extra_minus1 == 0
-            extra_minus1[extra_zero_pad_idx] = extra_plus1[extra_zero_pad_idx]
-
-            cat_feats[extra_replace] = extra_minus1
-            assert torch.all(~(cat_feats[extra_replace] == orig_cat_feats[extra_replace])).item() is True
-
-        pretext_label = (cat_feats != orig_cat_feats).long()
-        pretext_data = cat_feats
-
-        pretext_data = pretext_data.to(self.device, non_blocking=True)
-        pretext_label = pretext_label.to(self.device, non_blocking=True)
-
-        return pretext_data, pretext_label