make room for special tokens in categories embedding table

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'tab-transformer-pytorch',
   packages = find_packages(),
-  version = '0.0.8',
+  version = '0.0.9',
   license='MIT',
   description = 'Tab Transformer - Pytorch',
   author = 'Phil Wang',

tab_transformer_pytorch/tab_transformer_pytorch.py

@@ -115,6 +115,7 @@ def __init__(
         dim_out = 1,
         mlp_hidden_mults = (4, 2),
         mlp_act = None,
+        num_special_tokens = 2,
         continuous_mean_std = None
     ):
         super().__init__()
@@ -125,14 +126,21 @@ def __init__(
         self.num_categories = len(categories)
         self.num_unique_categories = sum(categories)
 
+        # create category embeddings table
+
+        self.num_special_tokens = num_special_tokens
+        total_tokens = self.num_unique_categories + num_special_tokens
+
+        self.categorical_embeds = nn.Embedding(total_tokens, dim)
+
         # for automatically offsetting unique category ids to the correct position in the categories embedding table
 
-        categories_offset = F.pad(torch.tensor(list(categories)), (1, 0), value = 0).cumsum(dim = -1)[:-1]
+        categories_offset = F.pad(torch.tensor(list(categories)), (1, 0), value = num_special_tokens)
+        categories_offset = categories_offset.cumsum(dim = -1)[:-1]
         self.register_buffer('categories_offset', categories_offset)
 
-        self.categorical_embeds = nn.Embedding(self.num_unique_categories, dim)
-
         # continuous
+
         if exists(continuous_mean_std):
             assert continuous_mean_std.shape == (num_continuous, 2), f'continuous_mean_std must have a shape of ({num_continuous}, 2) where the last dimension contains the mean and variance respectively'
             self.register_buffer('continuous_mean_std', continuous_mean_std)