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