done

config.json

@@ -70,7 +70,7 @@
   "num_hidden_layers": 22,
   "num_initial_layers": 1,
   "pad_logits": true,
-  "pad_token_id": 50283,
+  "pad_token_id": null,
   "padding": "unpadded",
   "pooling_type": "cls",
   "position_embedding_type": "absolute",

modeling_flexbert.py

@@ -1713,36 +1713,14 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         self,
         input_ids: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
         **kwargs
     ) -> dict:
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
 
-        # Calculate sequence-local positions
-        seqlens = attention_mask.sum(dim=-1)  # Get length of each sequence
-        position_ids = torch.zeros_like(input_ids)
-        for i in range(len(seqlens)):
-            position_ids[i, :seqlens[i]] = torch.arange(seqlens[i], device=input_ids.device)
-            
-            
-        batch_size, seq_len = input_ids.shape[:2]
-        if self.unpad_embeddings:
-            input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = self.unpad_inputs(
-                input_ids, attention_mask, position_ids, None
-            )
-        else:
-            indices = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).repeat(batch_size, 1)
-            cu_seqlens = None
-            max_seqlen = None
         return {
             "input_ids": input_ids,
             "attention_mask": attention_mask,
-            "position_ids": position_ids,
-            "indices": indices,
-            "cu_seqlens": cu_seqlens,
-            "max_seqlen": max_seqlen,
-            "batch_size": batch_size,
         }
 
     def get_number_parameters(self, count_embeddings: bool = True, trainable: bool = True) -> int:

tokenizer.py

@@ -23,7 +23,11 @@ class ModernDecoderBERTTokenizer(PreTrainedTokenizerFast):
                 ends_with_eos(seq) for seq in input_ids
             ], dtype=torch.bool)
             
-            if last_token_is_eos.any():
+            if last_token_is_eos.all():
+                # If all sequences have EOS, just truncate all
+                for key in ['input_ids', 'attention_mask']:
+                    outputs[key] = outputs[key][..., :-1]
+            elif last_token_is_eos.any():
                 # Process each sequence individually
                 batch_size = input_ids.shape[0]
                 for i in range(batch_size):
@@ -41,7 +45,11 @@ class ModernDecoderBERTTokenizer(PreTrainedTokenizerFast):
                 ends_with_eos(seq) for seq in input_ids
             ], dtype=bool)
             
-            if last_token_is_eos.any():
+            if last_token_is_eos.all():
+                # If all sequences have EOS, just truncate all
+                for key in ['input_ids', 'attention_mask']:
+                    outputs[key] = outputs[key][..., :-1]
+            elif last_token_is_eos.any():
                 batch_size = input_ids.shape[0]
                 for i in range(batch_size):
                     if last_token_is_eos[i]:
@@ -56,7 +64,11 @@ class ModernDecoderBERTTokenizer(PreTrainedTokenizerFast):
         elif isinstance(input_ids, list):
             last_token_is_eos = [ends_with_eos(seq) for seq in input_ids]
             
-            if any(last_token_is_eos):
+            if all(last_token_is_eos):
+                # If all sequences have EOS, just truncate all
+                for key in ['input_ids', 'attention_mask']:
+                    outputs[key] = [sequence[:-1] for sequence in outputs[key]]
+            elif any(last_token_is_eos):
                 for key in ['input_ids', 'attention_mask']:
                     outputs[key] = [
                         [0] + sequence[:-1] if is_eos else sequence