Fixed attribute parsing

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -1,13 +1,12 @@
 from transformers import AutoTokenizer
 from xml_schema_validator import XmlLogitsProcessor
 import os
-import time
 import torch
 
 model_path = "../../model"
 api_token = os.environ["SIA_HF_API_KEY"]
 
-xml_schema_actions = open("../../action_schema.xsd").read()
+xml_schema_actions = open("example_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">
@@ -59,60 +58,27 @@ def test_token_masking():
     
     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)
     
-    # Create dummy input_ids and scores
+    # Process valid continuation
     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')
+    space_token = tokenizer.encode(" ", add_special_tokens=False)[0]
+    assert processed_scores[0, space_token] > -float('inf')
     
     # Verify that invalid continuations are masked
-    assert processed_scores[0, tokenizer.encode("<invalid>", add_special_tokens=False)[0]] == -float('inf')
-
-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')
+    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')
     
 if __name__ == "__main__":
     # Run individual tests for debugging
@@ -120,5 +86,3 @@ if __name__ == "__main__":
     test_xml_schema_parsing()
     test_basic_xml_validation()
     test_token_masking()
-    test_performance()
-    test_subword_tokens()