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#""#;
+ let input = r#""#;
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("