From c2bf4223630d9c55e4921096669e77e898aa7077 Mon Sep 17 00:00:00 2001
From: Geens <niels.geens@gmail.com>
Date: Sat, 12 Apr 2025 19:59:25 +0200
Subject: [PATCH] Fixed whitespace processing

---
 .../python/xml_schema_validator/__init__.py   | 27 ++++++++++++++++---
 lib/xml_schema_validator/src/lib.rs           |  2 +-
 .../tests/test_logits_processor.py            | 22 +++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
index c7cc733..e30e3bc 100644
--- a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
+++ b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
@@ -30,13 +30,34 @@ 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():
-                items[id] = token
+                # 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")
-            
+        
         self.prompt_length = None
         self.is_first_call = True
 
diff --git a/lib/xml_schema_validator/src/lib.rs b/lib/xml_schema_validator/src/lib.rs
index 795e6b9..37584d8 100644
--- a/lib/xml_schema_validator/src/lib.rs
+++ b/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());
diff --git a/lib/xml_schema_validator/tests/test_logits_processor.py b/lib/xml_schema_validator/tests/test_logits_processor.py
index 5ddcd18..b7c6016 100644
--- a/lib/xml_schema_validator/tests/test_logits_processor.py
+++ b/lib/xml_schema_validator/tests/test_logits_processor.py
@@ -79,6 +79,27 @@ def test_token_masking():
     processed_scores = processor(input_ids, scores)
     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
@@ -86,3 +107,4 @@ if __name__ == "__main__":
     test_xml_schema_parsing()
     test_basic_xml_validation()
     test_token_masking()
+    test_delete()