84 lines
3.3 KiB
Python
84 lines
3.3 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)
|
|
self.prompt_length = None
|
|
self.is_first_call = True
|
|
|
|
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
|
|
|
|
# If this is the first call, store the prompt length
|
|
if self.is_first_call:
|
|
self.prompt_length = input_ids.shape[1]
|
|
self.is_first_call = False
|
|
|
|
# For each sequence in the batch
|
|
for batch_idx in range(batch_size):
|
|
# Get only the generated portion of the text
|
|
current_ids = input_ids[batch_idx]
|
|
generated_ids = current_ids[self.prompt_length:]
|
|
generated_text = self.tokenizer.decode(generated_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)
|
|
batch_validator = self.schema_validator.copy()
|
|
|
|
print("evaluate tokens continuing:", generated_text)
|
|
|
|
if generated_text:
|
|
valid, msg = batch_validator.append(generated_text)
|
|
if not valid:
|
|
print("current generated text invalid, only accept eos_token_id")
|
|
eos_token_id = self.tokenizer.eos_token_id
|
|
if eos_token_id is not None:
|
|
valid_tokens_mask[eos_token_id] = True
|
|
invalid_tokens_mask = ~valid_tokens_mask
|
|
scores[batch_idx, invalid_tokens_mask] = float('-inf')
|
|
continue
|
|
|
|
# Rest of the method remains the same
|
|
for token_idx in range(vocab_size):
|
|
token_validator = batch_validator.copy()
|
|
token_text = self.tokenizer.decode([token_idx])
|
|
valid, msg = token_validator.append(token_text)
|
|
#print("token:", token_text, "valid:", valid)
|
|
if valid:
|
|
#print("token:", generated_text + token_text)
|
|
valid_tokens_mask[token_idx] = True
|
|
|
|
invalid_tokens_mask = ~valid_tokens_mask
|
|
scores[batch_idx, invalid_tokens_mask] = float('-inf')
|
|
|
|
return scores |