add mask_first_token

config.json

@@ -52,6 +52,7 @@
   },
   "length_penalty": 2.0,
   "lsh_num_pre_rounds": 1,
+  "mask_first_token": false,
   "max_length": 64,
   "max_position_embeddings": 4096,
   "model_type": "bart",

modeling_lsg_bart.py

@@ -31,6 +31,7 @@ class LSGBartConfig(BartConfig):
         base_model_prefix="lsg",
         block_size=128,
         lsh_num_pre_rounds=1,
+        mask_first_token=False,
         num_global_tokens=1,
         pass_global_tokens_to_decoder=True,
         pool_with_global=True,
@@ -47,6 +48,7 @@ class LSGBartConfig(BartConfig):
         self.base_model_prefix = base_model_prefix
         self.block_size = block_size
         self.lsh_num_pre_rounds = lsh_num_pre_rounds
+        self.mask_first_token = mask_first_token
         self.num_global_tokens = num_global_tokens
         self.pass_global_tokens_to_decoder = pass_global_tokens_to_decoder
         self.pool_with_global = pool_with_global
@@ -711,6 +713,7 @@ class LSGBartEncoder(LSGBartPretrainedModel, BartEncoder):
         assert hasattr(config, "block_size") and hasattr(config, "adaptive")
         self.block_size = config.block_size
         self.adaptive = config.adaptive
+        self.mask_first_token = config.mask_first_token
         self.pool_with_global = config.pool_with_global
         self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
 
@@ -737,7 +740,9 @@ class LSGBartEncoder(LSGBartPretrainedModel, BartEncoder):
 
         if attention_mask is None:
             attention_mask = torch.ones(n, t, device=inputs_.device)
-            
+        if self.mask_first_token:
+            attention_mask[:, 0] = 0
+
         b = self.block_size * 2
         pad = t % self.block_size