unpad

modeling_flexbert.py

@@ -1727,9 +1727,14 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
             
             
         batch_size, seq_len = input_ids.shape[:2]
-        input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = self.unpad_inputs(
-            input_ids, attention_mask, position_ids, None
-        )
+        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,

tokenizer.py

@@ -7,13 +7,53 @@ class ModernDecoderBERTTokenizer(PreTrainedTokenizerFast):
     def _batch_encode_plus(self, *args, **kwargs):
         outputs = super()._batch_encode_plus(*args, **kwargs)
         del outputs["token_type_ids"]
+        
+        # Get the input_ids to check for EOS tokens
+        input_ids = outputs['input_ids']
+        
+        # Function to check if sequence ends with EOS token
+        def ends_with_eos(sequence):
+            if len(sequence) == 0:
+                return False
+            return sequence[-1] == self.eos_token_id
+        
+        # Check for EOS tokens using input_ids only
+        if isinstance(input_ids, torch.Tensor):
+            last_token_is_eos = torch.tensor([
+                ends_with_eos(seq) for seq in input_ids
+            ], dtype=torch.bool)
+        elif isinstance(input_ids, numpy.ndarray):
+            last_token_is_eos = numpy.array([
+                ends_with_eos(seq) for seq in input_ids
+            ], dtype=bool)
+        elif isinstance(input_ids, list):
+            last_token_is_eos = [ends_with_eos(seq) for seq in input_ids]
+        
+        # Use the same last_token_is_eos check for both input_ids and attention_mask
         for key in ['input_ids', 'attention_mask']:
             if isinstance(outputs[key], torch.Tensor):
-                outputs[key] = outputs[key][..., :-1]
-            elif isinstance(outputs[key], numpy.ndarray): 
-                outputs[key] = outputs[key][..., :-1]
+                # Only remove last token where last_token_is_eos is True
+                mask = last_token_is_eos.unsqueeze(-1)
+                outputs[key] = torch.where(
+                    mask,
+                    outputs[key][..., :-1],
+                    outputs[key]
+                )
+            elif isinstance(outputs[key], numpy.ndarray):
+                # Expand dimensions for broadcasting
+                mask = numpy.expand_dims(last_token_is_eos, -1)
+                outputs[key] = numpy.where(
+                    mask,
+                    outputs[key][..., :-1],
+                    outputs[key]
+                )
             elif isinstance(outputs[key], list):
-                outputs[key] = [sequence[:-1] for sequence in outputs[key]]
+                # For lists, use the same last_token_is_eos list for both keys
+                outputs[key] = [
+                    sequence[:-1] if is_eos else sequence
+                    for sequence, is_eos in zip(outputs[key], last_token_is_eos)
+                ]
+        
         return outputs
 
 # Register the class