fix tok

tokenizer.py

@@ -22,40 +22,50 @@ class ModernDecoderBERTTokenizer(PreTrainedTokenizerFast):
             last_token_is_eos = torch.tensor([
                 ends_with_eos(seq) for seq in input_ids
             ], dtype=torch.bool)
+            
+            if last_token_is_eos.any():
+                # Process each sequence individually
+                batch_size = input_ids.shape[0]
+                for i in range(batch_size):
+                    if last_token_is_eos[i]:
+                        for key in ['input_ids', 'attention_mask']:
+                            # Remove last token and add padding at start for this sequence
+                            truncated = outputs[key][i, :-1]
+                            outputs[key][i] = torch.cat([
+                                torch.zeros_like(truncated[:1]),
+                                truncated
+                            ])
+                
         elif isinstance(input_ids, numpy.ndarray):
             last_token_is_eos = numpy.array([
                 ends_with_eos(seq) for seq in input_ids
             ], dtype=bool)
+            
+            if last_token_is_eos.any():
+                batch_size = input_ids.shape[0]
+                for i in range(batch_size):
+                    if last_token_is_eos[i]:
+                        for key in ['input_ids', 'attention_mask']:
+                            # Remove last token and add padding at start for this sequence
+                            truncated = outputs[key][i, :-1]
+                            outputs[key][i] = numpy.concatenate([
+                                numpy.zeros_like(truncated[:1]),
+                                truncated
+                            ])
+                    
         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):
-                # 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):
-                # 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)
-                ]
+            
+            if any(last_token_is_eos):
+                for key in ['input_ids', 'attention_mask']:
+                    outputs[key] = [
+                        [0] + sequence[:-1] if is_eos else sequence
+                        for sequence, is_eos in zip(outputs[key], last_token_is_eos)
+                    ]
         
         return outputs
 
+        
 # Register the class
 from transformers import AutoTokenizer
 AutoTokenizer.register(ModernDecoderBERTTokenizer, fast_tokenizer_class=ModernDecoderBERTTokenizer)