llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added schema validator for llama.cpp

Browse Source
This commit is contained in:
geens
2025-05-06 18:46:28 +02:00
parent f61c9ab15f
commit 44fd60a6be
3 changed files with 262 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
lib/xml_schema_validator/python/xml_schema_validator/__init__.py
View File

@@ -1,130 +1,4 @@
import torch from .transformers_logits_processor import TransformersLogitsProcessor
from transformers import LogitsProcessor from .llama_cpp_logits_processor import LlamaCppLogitsProcessor
from transformers import AutoTokenizer
from xml_schema_validator._rs import XmlLogitsProcessorCore
class XmlLogitsProcessor(LogitsProcessor): __all__ = ["TransformersLogitsProcessor", "LlamaCppLogitsProcessor"]
"""
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: AutoTokenizer, schema_text: str = None, core: XmlLogitsProcessorCore=None):
"""
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
core: An existing core processor (for internal use in copy())
"""
self.eos_token_id = tokenizer.eos_token_id
self.tokenizer = tokenizer
if core is not None:
# Used for copy() operation
self.core = core
elif schema_text:
# Normal initialization
vocab = tokenizer.get_vocab() # This is {token: id}
items = {}
for id in range(len(vocab)):
items[id] = tokenizer.decode([id])
self.core = XmlLogitsProcessorCore(items, schema_text)
else:
raise ValueError("Either schema_text or core must be provided")
# Find the assistant start marker in the chat template
# You can also manually override this by setting a specific marker if needed
self.assistant_start_marker = self._get_assistant_start_marker(tokenizer)
def _get_assistant_start_marker(self, tokenizer):
"""
Extract the assistant start marker from the tokenizer's chat template.
Args:
tokenizer: The tokenizer to extract the marker from
Returns:
The marker string that indicates the start of the assistant's response
"""
return tokenizer.apply_chat_template(
[{"role": "assistant", "content": ""},],
tokenize=False,
add_generation_prompt=False,
)
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):
# Decode the entire sequence so far
current_ids = input_ids[batch_idx]
full_text = self.tokenizer.decode(current_ids)
# Find the last occurrence of the assistant start marker
last_marker_pos = full_text.rfind(self.assistant_start_marker)
# Extract the assistant's response by taking text after the marker plus its length
start_pos = last_marker_pos + len(self.assistant_start_marker)
generated_text = full_text[start_pos:]
# Create a processor copy to track which tokens are valid
batch_processor = self.core.copy()
if generated_text:
# Try to append the current text to the validator
append_result = batch_processor.append(generated_text)
# If we've reached a valid EOF state, only allow the EOS token
if batch_processor.eof():
vocab_size = scores.shape[-1]
valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
valid_tokens_mask[self.eos_token_id] = True
invalid_tokens_mask = ~valid_tokens_mask
scores[batch_idx, invalid_tokens_mask] = float('-inf')
continue
# If the validation failed altogether, this is an invalid path
if not append_result:
# Allow only EOS token if validation fails
scores[batch_idx, :] = float('-inf')
scores[batch_idx, self.eos_token_id] = 0
continue
# Get tokens that would lead to invalid XML and mask them
invalid_tokens = batch_processor.get_invalid_tokens()
for token in invalid_tokens:
if token < scores.shape[1]:
scores[batch_idx, token] = float('-inf')
return scores
def copy(self):
"""
Create a copy of the processor.
Returns:
A new instance of the processor
"""
# Create a new instance using the existing core
cloned_core = self.core.copy()
cloned = XmlLogitsProcessor(self.tokenizer, core=cloned_core)
# Copy all state
cloned.assistant_start_marker = self.assistant_start_marker
return cloned

123
lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py Normal file
View File

@@ -0,0 +1,123 @@
from transformers import AutoTokenizer
from xml_schema_validator._rs import XmlLogitsProcessorCore
class LlamaCppLogitsProcessor:
"""
A LogitsProcessor for llama-cpp-python 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: AutoTokenizer, schema_text: str = None, core: XmlLogitsProcessorCore=None):
"""
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
core: An existing core processor (for internal use)
"""
self.eos_token_id = tokenizer.eos_token_id
self.tokenizer = tokenizer
if core is not None:
# Used for internal operations
self.core = core
elif schema_text:
# Normal initialization
vocab = tokenizer.get_vocab() # This is {token: id}
items = {}
for id in range(len(vocab)):
items[id] = tokenizer.decode([id])
self.core = XmlLogitsProcessorCore(items, schema_text)
else:
raise ValueError("Either schema_text or core must be provided")
# Find the assistant start marker in the chat template
# You can also manually override this by setting a specific marker if needed
self.assistant_start_marker = self._get_assistant_start_marker(tokenizer)
def _get_assistant_start_marker(self, tokenizer):
"""
Extract the assistant start marker from the tokenizer's chat template.
Args:
tokenizer: The tokenizer to extract the marker from
Returns:
The marker string that indicates the start of the assistant's response
"""
template_text = tokenizer.apply_chat_template(
[
{"role": "user", "content": "XML_LOGITS_PROCESSOR_MARKER"},
],
tokenize=False,
add_generation_prompt=True,
)
marker_start = template_text.find("XML_LOGITS_PROCESSOR_MARKER") + len("XML_LOGITS_PROCESSOR_MARKER")
start_marker = template_text[marker_start:]
print("schema validator start marker:", start_marker)
return start_marker
def get_processor(self):
"""
Returns a function that can be used with llama-cpp-python's LogitsProcessorList.
Returns:
A function that processes logits for llama-cpp-python
"""
def process_logits(tokens, logits):
"""
Process logits to mask invalid XML tokens.
Args:
tokens: List of token IDs generated so far
logits: List of logits for the next token prediction
Returns:
Processed logits with invalid tokens masked
"""
# Create a processor copy to track which tokens are valid for this specific call
batch_processor = self.core.copy()
# Decode the entire sequence so far
full_text = self.tokenizer.decode(tokens)
# Find the last occurrence of the assistant start marker
last_marker_pos = full_text.rfind(self.assistant_start_marker)
# Extract the assistant's response by taking text after the marker plus its length
start_pos = last_marker_pos + len(self.assistant_start_marker)
generated_text = full_text[start_pos:]
if generated_text:
# Try to append the current text to the validator
append_result = batch_processor.append(generated_text)
# If we've reached a valid EOF state, only allow the EOS token
if batch_processor.eof():
vocab_size = len(logits)
processed_logits = [float('-inf')] * vocab_size
processed_logits[self.eos_token_id] = 0.0
return processed_logits
# If the validation failed altogether, this is an invalid path
if not append_result:
# Allow only EOS token if validation fails
processed_logits = [float('-inf')] * len(logits)
processed_logits[self.eos_token_id] = 0.0
return processed_logits
# Make a copy of the logits to modify
processed_logits = list(logits)
# Get tokens that would lead to invalid XML and mask them
invalid_tokens = batch_processor.get_invalid_tokens()
for token in invalid_tokens:
if token < len(processed_logits):
processed_logits[token] = float('-inf')
return processed_logits
return process_logits

136
lib/xml_schema_validator/python/xml_schema_validator/transformers_logits_processor.py Normal file
View File

@@ -0,0 +1,136 @@
import torch
from transformers import LogitsProcessor
from transformers import AutoTokenizer
from xml_schema_validator._rs import XmlLogitsProcessorCore
class TransformersLogitsProcessor(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: AutoTokenizer, schema_text: str = None, core: XmlLogitsProcessorCore=None):
"""
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
core: An existing core processor (for internal use in copy())
"""
self.eos_token_id = tokenizer.eos_token_id
self.tokenizer = tokenizer
if core is not None:
# Used for copy() operation
self.core = core
elif schema_text:
# Normal initialization
vocab = tokenizer.get_vocab() # This is {token: id}
items = {}
for id in range(len(vocab)):
items[id] = tokenizer.decode([id])
self.core = XmlLogitsProcessorCore(items, schema_text)
else:
raise ValueError("Either schema_text or core must be provided")
# Find the assistant start marker in the chat template
# You can also manually override this by setting a specific marker if needed
self.assistant_start_marker = self._get_assistant_start_marker(tokenizer)
def _get_assistant_start_marker(self, tokenizer):
"""
Extract the assistant start marker from the tokenizer's chat template.
Args:
tokenizer: The tokenizer to extract the marker from
Returns:
The marker string that indicates the start of the assistant's response
"""
template_text = tokenizer.apply_chat_template(
[
{"role": "user", "content": "XML_LOGITS_PROCESSOR_MARKER"},
],
tokenize=False,
add_generation_prompt=True,
)
marker_start = template_text.find("XML_LOGITS_PROCESSOR_MARKER") + len("XML_LOGITS_PROCESSOR_MARKER")
start_marker = template_text[marker_start:]
print("schema validator start marker:", start_marker)
return start_marker
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):
# Decode the entire sequence so far
current_ids = input_ids[batch_idx]
full_text = self.tokenizer.decode(current_ids)
# Find the last occurrence of the assistant start marker
last_marker_pos = full_text.rfind(self.assistant_start_marker)
# Extract the assistant's response by taking text after the marker plus its length
start_pos = last_marker_pos + len(self.assistant_start_marker)
generated_text = full_text[start_pos:]
# Create a processor copy to track which tokens are valid
batch_processor = self.core.copy()
if generated_text:
# Try to append the current text to the validator
append_result = batch_processor.append(generated_text)
# If we've reached a valid EOF state, only allow the EOS token
if batch_processor.eof():
vocab_size = scores.shape[-1]
valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
valid_tokens_mask[self.eos_token_id] = True
invalid_tokens_mask = ~valid_tokens_mask
scores[batch_idx, invalid_tokens_mask] = float('-inf')
continue
# If the validation failed altogether, this is an invalid path
if not append_result:
# Allow only EOS token if validation fails
scores[batch_idx, :] = float('-inf')
scores[batch_idx, self.eos_token_id] = 0
continue
# Get tokens that would lead to invalid XML and mask them
invalid_tokens = batch_processor.get_invalid_tokens()
for token in invalid_tokens:
if token < scores.shape[1]:
scores[batch_idx, token] = float('-inf')
return scores
def copy(self):
"""
Create a copy of the processor.
Returns:
A new instance of the processor
"""
# Create a new instance using the existing core
cloned_core = self.core.copy()
cloned = TransformersLogitsProcessor(self.tokenizer, core=cloned_core)
# Copy all state
cloned.assistant_start_marker = self.assistant_start_marker
return cloned