Revert "Final code"

This reverts commit 5b4bf3bc57df685ce7894faa33d20baca68215fc.

app.py

@@ -17,7 +17,7 @@ def compute(table: np.array):
     vfunc = np.vectorize(lambda s: len(str(s)))
     non_empty_row_mask = (vfunc(table).sum(1) != 0)
     table = table[non_empty_row_mask]
-    empty_mask = table == ''
+    empty_mask = table == '(predict)'
     empty_inds = np.where(empty_mask)
     if table.shape[0] > 1024:
         return "⚠️ **ERROR: TabPFN is not made for datasets with a trainingsize > 1024.**", None, None

transformer.py

@@ -0,0 +1,232 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module, TransformerEncoder
+
+from layer import TransformerEncoderLayer, _get_activation_fn
+from utils import SeqBN, bool_mask_to_att_mask
+
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, encoder, n_out, ninp, nhead, nhid, nlayers, dropout=0.0, style_encoder=None, y_encoder=None,
+                 pos_encoder=None, decoder=None, input_normalization=False, init_method=None, pre_norm=False,
+                 activation='gelu', recompute_attn=False, num_global_att_tokens=0, full_attention=False,
+                 all_layers_same_init=False, efficient_eval_masking=True):
+        super().__init__()
+        self.model_type = 'Transformer'
+        encoder_layer_creator = lambda: TransformerEncoderLayer(ninp, nhead, nhid, dropout, activation=activation,
+                                                                pre_norm=pre_norm, recompute_attn=recompute_attn)
+        self.transformer_encoder = TransformerEncoder(encoder_layer_creator(), nlayers)\
+            if all_layers_same_init else TransformerEncoderDiffInit(encoder_layer_creator, nlayers)
+        self.ninp = ninp
+        self.encoder = encoder
+        self.y_encoder = y_encoder
+        self.pos_encoder = pos_encoder
+        self.decoder = decoder(ninp, nhid, n_out) if decoder is not None else nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, n_out))
+        self.input_ln = SeqBN(ninp) if input_normalization else None
+        self.style_encoder = style_encoder
+        self.init_method = init_method
+        if num_global_att_tokens is not None:
+            assert not full_attention
+        self.global_att_embeddings = nn.Embedding(num_global_att_tokens, ninp) if num_global_att_tokens else None
+        self.full_attention = full_attention
+        self.efficient_eval_masking = efficient_eval_masking
+
+        self.n_out = n_out
+        self.nhid = nhid
+
+        self.init_weights()
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self.__dict__.setdefault('efficient_eval_masking', False)
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_D_q_matrix(sz, query_size):
+        train_size = sz-query_size
+        mask = torch.zeros(sz,sz) == 0
+        mask[:,train_size:].zero_()
+        mask |= torch.eye(sz) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_query_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        sz = seq_len + num_global_att_tokens
+        mask = torch.zeros(num_query_tokens, sz) == 0
+        mask[:,train_size:].zero_()
+        mask[:,train_size:] |= torch.eye(num_query_tokens) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_trainset_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        trainset_size = seq_len - num_query_tokens
+        mask = torch.zeros(trainset_size, num_global_att_tokens) == 0
+        #mask[:,num_global_att_tokens:].zero_()
+        #mask[:,num_global_att_tokens:] |= torch.eye(trainset_size) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_globaltokens_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        mask = torch.zeros(num_global_att_tokens, num_global_att_tokens+seq_len-num_query_tokens) == 0
+        return bool_mask_to_att_mask(mask)
+
+    def init_weights(self):
+        initrange = 1.
+        # if isinstance(self.encoder,EmbeddingEncoder):
+        #    self.encoder.weight.data.uniform_(-initrange, initrange)
+        # self.decoder.bias.data.zero_()
+        # self.decoder.weight.data.uniform_(-initrange, initrange)
+        if self.init_method is not None:
+            self.apply(self.init_method)
+        for layer in self.transformer_encoder.layers:
+            nn.init.zeros_(layer.linear2.weight)
+            nn.init.zeros_(layer.linear2.bias)
+            attns = layer.self_attn if isinstance(layer.self_attn, nn.ModuleList) else [layer.self_attn]
+            for attn in attns:
+                nn.init.zeros_(attn.out_proj.weight)
+                nn.init.zeros_(attn.out_proj.bias)
+
+    def forward(self, src, src_mask=None, single_eval_pos=None):
+        assert isinstance(src, tuple), 'inputs (src) have to be given as (x,y) or (style,x,y) tuple'
+
+        if len(src) == 2: # (x,y) and no style
+            src = (None,) + src
+
+        style_src, x_src, y_src = src
+        x_src = self.encoder(x_src)
+        y_src = self.y_encoder(y_src.unsqueeze(-1) if len(y_src.shape) < len(x_src.shape) else y_src)
+        style_src = self.style_encoder(style_src).unsqueeze(0) if self.style_encoder else \
+            torch.tensor([], device=x_src.device)
+        global_src = torch.tensor([], device=x_src.device) if self.global_att_embeddings is None else \
+            self.global_att_embeddings.weight.unsqueeze(1).repeat(1, x_src.shape[1], 1)
+
+        if src_mask is not None: assert self.global_att_embeddings is None or isinstance(src_mask, tuple)
+        if src_mask is None:
+            if self.global_att_embeddings is None:
+                full_len = len(x_src) + len(style_src)
+                if self.full_attention:
+                    src_mask = bool_mask_to_att_mask(torch.ones((full_len, full_len), dtype=torch.bool)).to(x_src.device)
+                elif self.efficient_eval_masking:
+                    src_mask = single_eval_pos + len(style_src)
+                else:
+                    src_mask = self.generate_D_q_matrix(full_len, len(x_src) - single_eval_pos).to(x_src.device)
+            else:
+                src_mask_args = (self.global_att_embeddings.num_embeddings,
+                                 len(x_src) + len(style_src),
+                                 len(x_src) + len(style_src) - single_eval_pos)
+                src_mask = (self.generate_global_att_globaltokens_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_trainset_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_query_matrix(*src_mask_args).to(x_src.device))
+
+        train_x = x_src[:single_eval_pos] + y_src[:single_eval_pos]
+        src = torch.cat([global_src, style_src, train_x, x_src[single_eval_pos:]], 0)
+
+        if self.input_ln is not None:
+            src = self.input_ln(src)
+
+        if self.pos_encoder is not None:
+            src = self.pos_encoder(src)
+
+        output = self.transformer_encoder(src, src_mask)
+        output = self.decoder(output)
+        return output[single_eval_pos+len(style_src)+(self.global_att_embeddings.num_embeddings if self.global_att_embeddings else 0):]
+
+    @torch.no_grad()
+    def init_from_small_model(self, small_model):
+        assert isinstance(self.decoder, nn.Linear) and isinstance(self.encoder, (nn.Linear, nn.Sequential)) \
+               and isinstance(self.y_encoder, (nn.Linear, nn.Sequential))
+
+        def set_encoder_weights(my_encoder, small_model_encoder):
+            my_encoder_linear, small_encoder_linear = (my_encoder, small_model_encoder) \
+                if isinstance(my_encoder, nn.Linear) else (my_encoder[-1], small_model_encoder[-1])
+            small_in_dim = small_encoder_linear.out_features
+            my_encoder_linear.weight.zero_()
+            my_encoder_linear.bias.zero_()
+            my_encoder_linear.weight[:small_in_dim] = small_encoder_linear.weight
+            my_encoder_linear.bias[:small_in_dim] = small_encoder_linear.bias
+
+        set_encoder_weights(self.encoder, small_model.encoder)
+        set_encoder_weights(self.y_encoder, small_model.y_encoder)
+
+        small_in_dim = small_model.decoder.in_features
+
+        self.decoder.weight[:, :small_in_dim] = small_model.decoder.weight
+        self.decoder.bias = small_model.decoder.bias
+
+        for my_layer, small_layer in zip(self.transformer_encoder.layers, small_model.transformer_encoder.layers):
+            small_hid_dim = small_layer.linear1.out_features
+            my_in_dim = my_layer.linear1.in_features
+
+            # packed along q,k,v order in first dim
+            my_in_proj_w = my_layer.self_attn.in_proj_weight
+            small_in_proj_w = small_layer.self_attn.in_proj_weight
+
+            my_in_proj_w.view(3, my_in_dim, my_in_dim)[:, :small_in_dim, :small_in_dim] = small_in_proj_w.view(3,
+                                                                                                               small_in_dim,
+                                                                                                               small_in_dim)
+            my_layer.self_attn.in_proj_bias.view(3, my_in_dim)[:,
+            :small_in_dim] = small_layer.self_attn.in_proj_bias.view(3, small_in_dim)
+
+            my_layer.self_attn.out_proj.weight[:small_in_dim, :small_in_dim] = small_layer.self_attn.out_proj.weight
+            my_layer.self_attn.out_proj.bias[:small_in_dim] = small_layer.self_attn.out_proj.bias
+
+            my_layer.linear1.weight[:small_hid_dim, :small_in_dim] = small_layer.linear1.weight
+            my_layer.linear1.bias[:small_hid_dim] = small_layer.linear1.bias
+
+            my_layer.linear2.weight[:small_in_dim, :small_hid_dim] = small_layer.linear2.weight
+            my_layer.linear2.bias[:small_in_dim] = small_layer.linear2.bias
+
+            my_layer.norm1.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm1.weight
+            my_layer.norm2.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm2.weight
+
+            my_layer.norm1.bias[:small_in_dim] = small_layer.norm1.bias
+            my_layer.norm2.bias[:small_in_dim] = small_layer.norm2.bias
+
+
+class TransformerEncoderDiffInit(Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer_creator: a function generating objects of TransformerEncoderLayer class without args (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer_creator, num_layers, norm=None):
+        super().__init__()
+        self.layers = nn.ModuleList([encoder_layer_creator() for _ in range(num_layers)])
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            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.
+        """
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output