Fixed whitespace processing

lib/xml_schema_validator/python/xml_schema_validator/__init__.py

@@ -30,9 +30,30 @@ class XmlLogitsProcessor(LogitsProcessor):
         elif schema_text:
             # Normal initialization
             vocab = tokenizer.get_vocab()  # This is {token: id}
-            items = dict()
+            
+            # Create a mapping from token id to decoded text
+            # This ensures we capture whitespace correctly
+            items = {}
             for token, id in vocab.items():
+                # Properly decode each token to get its actual string representation
+                # This preserves whitespace characters
+                if isinstance(token, str):
+                    # For string tokens, use them directly
                     items[id] = token
+                else:
+                    # For byte tokens, ensure they're decoded properly
+                    items[id] = tokenizer.convert_ids_to_tokens([id])[0]
+            
+            # Additional special handling for whitespace tokens
+            # Some tokenizers have special tokens for whitespace
+            for id in range(len(vocab)):
+                if id in items:
+                    # Try decoding single tokens to find whitespace
+                    decoded = tokenizer.decode([id])
+                    if decoded and (decoded.isspace() or decoded.startswith(' ')):
+                        # This token represents whitespace - ensure it's preserved
+                        items[id] = decoded
+            
             self.core = XmlLogitsProcessorCore(items, schema_text)
         else:
             raise ValueError("Either schema_text or core must be provided")

lib/xml_schema_validator/src/lib.rs

@@ -146,7 +146,7 @@ mod tests {
     #[test]
     fn test_delete_with_attribute() {
         let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
-        let input = r#"<delete id="1234567890"/>"#;
+        let input = r#"<delete id="20250411_170104_361"/>"#;
         let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
         validator.append(input).unwrap();
         assert!(validator.eof());

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -80,9 +80,31 @@ def test_token_masking():
     assert processed_scores[0, tokenizer.eos_token_id] == 1
     assert processed_scores[0, space_token] == -float('inf')
 
+def test_delete():
+    """Test that invalid tokens are properly masked"""
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+
+    # Process empty input to set prompt length
+    input_ids = torch.tensor([tokenizer.encode("")])
+    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+    processed_scores = processor(input_ids, scores)
+    
+    # Process delete
+    input_ids = torch.tensor([tokenizer.encode("<delete")])
+    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+    processed_scores = processor(input_ids, scores)
+    for i, score in enumerate(processed_scores[0]):
+        if score == 1:
+            print(i, tokenizer.decode(i))
+    assert False
+    
+    
 if __name__ == "__main__":
     # Run individual tests for debugging
     test_logits_processor_init_copy()
     test_xml_schema_parsing()
     test_basic_xml_validation()
     test_token_masking()
+    test_delete()