Use CharTrie and fix multi-byte char issue

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -18,67 +18,67 @@ xml_schema_only_root_node = """<?xml version="1.0" encoding="UTF-8"?>
   </xs:element>
 </xs:schema>"""
 
-def test_logits_processor_init_copy():
-    """Test the initialization of the LogitsProcessor"""
-
-    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
-    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
-    assert logits_processor.core is not None
-
-    logits_processor_copy = logits_processor.copy()
-    assert logits_processor_copy.core is not None
-
-def test_xml_schema_parsing():
-    """Test basic XML schema parsing with different element types"""
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
-    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
-    assert logits_processor.core is not None
-
-    element_names = logits_processor.core.get_element_names()
-    assert len(element_names) > 0
-    assert "root" in element_names
-
-def test_basic_xml_validation():
-    """Test basic validation of XML fragments against the schema"""
-
-    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
-    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
-    
-    # Test appending valid XML fragments
-    assert processor.core.copy().append("<reasoning>")
-    assert processor.core.copy().append("<reasoning>Test content</reasoning>")
-    
-    # Test appending invalid XML fragments
-    assert not processor.core.copy().append("<invalid>")
-    assert not processor.core.copy().append("<reasoning></invalid>")
-
-def test_token_masking():
-    """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 valid continuation
-    input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
-    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
-    processed_scores = processor(input_ids, scores)
-    
-    # Verify that valid continuations still have positive scores
-    space_token = tokenizer.encode(" ", add_special_tokens=False)[0]
-    assert processed_scores[0, space_token] > -float('inf')
-    
-    # Verify that invalid continuations are masked
-    input_ids = torch.tensor([tokenizer.encode("</invalid>")])
-    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
-    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_logits_processor_init_copy():
+#    """Test the initialization of the LogitsProcessor"""
+#
+#    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+#    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
+#    assert logits_processor.core is not None
+#
+#    logits_processor_copy = logits_processor.copy()
+#    assert logits_processor_copy.core is not None
+#
+#def test_xml_schema_parsing():
+#    """Test basic XML schema parsing with different element types"""
+#    
+#    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+#    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
+#    assert logits_processor.core is not None
+#
+#    element_names = logits_processor.core.get_element_names()
+#    assert len(element_names) > 0
+#    assert "root" in element_names
+#
+#def test_basic_xml_validation():
+#    """Test basic validation of XML fragments against the schema"""
+#
+#    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+#    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+#    
+#    # Test appending valid XML fragments
+#    assert processor.core.copy().append("<reasoning>")
+#    assert processor.core.copy().append("<reasoning>Test content</reasoning>")
+#    
+#    # Test appending invalid XML fragments
+#    assert not processor.core.copy().append("<invalid>")
+#    assert not processor.core.copy().append("<reasoning></invalid>")
+#
+#def test_token_masking():
+#    """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 valid continuation
+#    input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
+#    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+#    processed_scores = processor(input_ids, scores)
+#    
+#    # Verify that valid continuations still have positive scores
+#    space_token = tokenizer.encode(" ", add_special_tokens=False)[0]
+#    assert processed_scores[0, space_token] > -float('inf')
+#    
+#    # Verify that invalid continuations are masked
+#    input_ids = torch.tensor([tokenizer.encode("</invalid>")])
+#    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+#    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"""
@@ -92,19 +92,16 @@ def test_delete():
     processed_scores = processor(input_ids, scores)
     
     # Process delete
-    input_ids = torch.tensor([tokenizer.encode("<delete")])
+    input_ids = torch.tensor([tokenizer.encode("<delete id=\"")])
+    print(input_ids)
     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_logits_processor_init_copy()
+    #test_xml_schema_parsing()
+    #test_basic_xml_validation()
+    #test_token_masking()
     test_delete()