Handle attributes and self-closing tags

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -1,14 +1,14 @@
-from xml_logits_processor import XmlLogitsProcessor
-from threading import Thread
-from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer
+from xml_schema_validator import XmlLogitsProcessor
 import os
-import sys
+import time
+import torch
 
-def test_logits_processor():
-    print("test_logits_processor")
+model_path = "../../model"
+api_token = os.environ["SIA_HF_API_KEY"]
 
-    model_name = "../../model"
-    xml_schema_text = """<?xml version="1.0" encoding="UTF-8"?>
+xml_schema_actions = open("../../action_schema.xsd").read()
+xml_schema_only_root_node = """<?xml version="1.0" encoding="UTF-8"?>
 <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
   <xs:element name="root">
     <xs:complexType mixed="true">
@@ -19,53 +19,106 @@ def test_logits_processor():
   </xs:element>
 </xs:schema>"""
 
-    model = "google/gemma-3-1b-it"
+def test_logits_processor_init_copy():
+    """Test the initialization of the LogitsProcessor"""
 
-    pipline = pipeline(
-        "text-generation",
-        model=model,
-        token=os.environ["SIA_HF_API_KEY"],
-    )
+    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
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        model,
-        token=os.environ["SIA_HF_API_KEY"],
-    )
+    logits_processor_copy = logits_processor.copy()
+    assert logits_processor_copy.core is not None
 
-    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)
-                        
-    messages = [
-        {"role": "system", "content": "Always respond with <root>message by the user</root>. So if the user says 'hello world', you response <root>hello world</root>"},
-        {"role": "user", "content": "hello, I am the user"}
-    ]
+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
 
-    text_inputs = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
+    element_names = logits_processor.core.get_element_names()
+    assert len(element_names) > 0
+    assert "root" in element_names
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True
-    )
+def test_basic_xml_validation():
+    """Test basic validation of XML fragments against the schema"""
 
-    generation_kwargs = {
-        "text_inputs": text_inputs,
-        "max_new_tokens": 20,
-        "streamer": streamer,
-        #"logits_processor": [logits_processor],
-    }
+    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>")
 
-    generation_thread = Thread(
-        target=pipline,
-        kwargs=generation_kwargs
-    )
+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)
+    
+    # Create dummy input_ids and scores
+    input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
+    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+    
+    # Process the scores
+    processed_scores = processor(input_ids, scores)
+    
+    # Verify that valid continuations still have positive scores
+    assert processed_scores[0, tokenizer.encode(" ", add_special_tokens=False)[0]] > -float('inf')
+    
+    # Verify that invalid continuations are masked
+    assert processed_scores[0, tokenizer.encode("<invalid>", add_special_tokens=False)[0]] == -float('inf')
 
-    generation_thread.start()
-
-    for text in streamer:
-        print(text, end="", file=sys.stderr, flush=True)
-
-    generation_thread.join()
+def test_performance():
+    """Test performance with larger XML documents"""
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+    
+    # Generate a larger XML document
+    large_xml = "<reasoning>" + "Test content. " * 100 + "</reasoning>"
+    
+    # Measure time to process
+    start_time = time.time()
+    processor.core.append(large_xml)
+    processing_time = time.time() - start_time
+    
+    print(f"Processing time for large XML: {processing_time:.4f} seconds")
+    # You might want to assert that processing time is below a threshold
 
+def test_subword_tokens():
+    """Test handling of subword tokens that might split XML tags"""
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    schema_text = open("../../action_schema.xsd").read()
+    
+    processor = XmlLogitsProcessor(tokenizer, schema_text)
+    
+    # Find some XML tag that gets split into multiple tokens by your tokenizer
+    tag = "<reasoning>"
+    tokens = tokenizer.encode(tag, add_special_tokens=False)
+    
+    if len(tokens) > 1:
+        print(f"Tag '{tag}' is split into {len(tokens)} tokens")
+        
+        # Test that the processor can handle these split tokens
+        import torch
+        input_ids = torch.tensor([[tokens[0]]])
+        scores = torch.ones((1, len(tokenizer)))
+        
+        processed_scores = processor(input_ids, scores)
+        
+        # The next token in the sequence should have a high score
+        assert processed_scores[0, tokens[1]] > -float('inf')
+    
 if __name__ == "__main__":
-    test_logits_processor()
+    # Run individual tests for debugging
+    test_logits_processor_init_copy()
+    test_xml_schema_parsing()
+    test_basic_xml_validation()
+    test_token_masking()
+    test_performance()
+    test_subword_tokens()