Fixed whitespace processing

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()