Fixed token cache for whitespace bug

This commit is contained in:
2025-04-13 10:03:14 +02:00
parent ff82173f2a
commit f4de9e7571

View File

@@ -34,25 +34,21 @@ class XmlLogitsProcessor(LogitsProcessor):
# Create a mapping from token id to decoded text # Create a mapping from token id to decoded text
# This ensures we capture whitespace correctly # This ensures we capture whitespace correctly
items = {} items = {}
for token, id in vocab.items(): #for token, id in vocab.items():
# Properly decode each token to get its actual string representation # # Properly decode each token to get its actual string representation
# This preserves whitespace characters # # This preserves whitespace characters
if isinstance(token, str): # if isinstance(token, str):
# For string tokens, use them directly # # For string tokens, use them directly
items[id] = token # items[id] = token
else: # else:
# For byte tokens, ensure they're decoded properly # print("b")
items[id] = tokenizer.convert_ids_to_tokens([id])[0] # # For byte tokens, ensure they're decoded properly
# items[id] = tokenizer.convert_ids_to_tokens([id])[0]
# Additional special handling for whitespace tokens # Additional special handling for whitespace tokens
# Some tokenizers have special tokens for whitespace # Some tokenizers have special tokens for whitespace
for id in range(len(vocab)): for id in range(len(vocab)):
if id in items: items[id] = tokenizer.decode([id])
# 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) self.core = XmlLogitsProcessorCore(items, schema_text)
else: else: