63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
import torch
|
|
from transformers import LogitsProcessor
|
|
from xml_schema_validator import XmlSchemaValidator
|
|
|
|
class XmlLogitsProcessor(LogitsProcessor):
|
|
"""
|
|
A LogitsProcessor that enforces valid XML according to a schema.
|
|
|
|
This processor masks tokens that would lead to invalid XML
|
|
by setting their logits to negative infinity.
|
|
"""
|
|
|
|
def __init__(self, tokenizer, schema_text: str):
|
|
"""
|
|
Initialize the processor with a schema and tokenizer.
|
|
|
|
Args:
|
|
tokenizer: The tokenizer to use for decoding tokens
|
|
schema_text: The XSD schema text to validate against
|
|
"""
|
|
self.tokenizer = tokenizer
|
|
self.schema_validator = XmlSchemaValidator(schema_text)
|
|
|
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
|
|
"""
|
|
Process logits to mask invalid XML tokens.
|
|
|
|
Args:
|
|
input_ids: Current input token IDs
|
|
scores: Current scores (logits) for next token prediction
|
|
|
|
Returns:
|
|
Processed scores with invalid tokens masked
|
|
"""
|
|
batch_size, _ = input_ids.shape
|
|
|
|
# For each sequence in the batch
|
|
for batch_idx in range(batch_size):
|
|
# Get the current text generated so far
|
|
current_ids = input_ids[batch_idx]
|
|
current_text = self.tokenizer.decode(current_ids)
|
|
|
|
# Get all possible next tokens
|
|
vocab_size = scores.shape[-1]
|
|
|
|
# Create a mask to track which tokens are valid
|
|
valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
|
|
|
|
# For each possible next token
|
|
for token_idx in range(vocab_size):
|
|
# Create a copy of the validator to test this token
|
|
validator_copy = self.schema_validator.copy()
|
|
|
|
# Decode the token and test if appending it would be valid
|
|
token_text = self.tokenizer.decode([token_idx])
|
|
if validator_copy.append(token_text):
|
|
valid_tokens_mask[token_idx] = True
|
|
|
|
# Mask out invalid tokens by setting their scores to negative infinity
|
|
invalid_tokens_mask = ~valid_tokens_mask
|
|
scores[batch_idx, invalid_tokens_mask] = float('-inf')
|
|
|
|
return scores |