Fix char_trie to store all token ids

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -3,7 +3,7 @@ from xml_schema_validator import XmlLogitsProcessor
 import os
 import torch
 
-model_path = "../../model"
+model_path = "/root/models/current"
 api_token = os.environ["SIA_HF_API_KEY"]
 
 xml_schema_actions = open("example_schema.xsd").read()
@@ -58,12 +58,32 @@ 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)
     
+    # Check if '<' is the only allowed token
+    less_than_token = tokenizer.encode("<", add_special_tokens=False)[0]
+    
+    # Count how many tokens have scores not equal to -inf
+    valid_tokens = (processed_scores[0] > -float('inf')).sum().item()
+    
+    # Check if the '<' token is allowed (not masked)
+    is_less_than_allowed = processed_scores[0, less_than_token] > -float('inf')
+    
+    print(f"Number of allowed tokens: {valid_tokens}")
+    print(f"Is '<' token allowed: {is_less_than_allowed}")
+    
+    # Always print the first 10 allowed tokens
+    allowed_token_ids = torch.where(processed_scores[0] > -float('inf'))[0].tolist()
+    # Take only the first 10 tokens (or all if less than 10)
+    allowed_token_ids_sample = allowed_token_ids[:min(10, len(allowed_token_ids))]
+    allowed_tokens = [tokenizer.decode([token_id]) for token_id in allowed_token_ids_sample]
+    print(f"First 10 allowed token IDs: {allowed_token_ids_sample}")
+    print(f"First 10 allowed tokens: {allowed_tokens}")
+    
     # Process valid continuation
     input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
     scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability