diff --git a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py index 22d78ac..a8a20f1 100644 --- a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py +++ b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py @@ -1,130 +1,4 @@ -import torch -from transformers import LogitsProcessor -from transformers import AutoTokenizer -from xml_schema_validator._rs import XmlLogitsProcessorCore +from .transformers_logits_processor import TransformersLogitsProcessor +from .llama_cpp_logits_processor import LlamaCppLogitsProcessor -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: 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 \ No newline at end of file +__all__ = ["TransformersLogitsProcessor", "LlamaCppLogitsProcessor"] \ No newline at end of file diff --git a/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py b/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py new file mode 100644 index 0000000..b67022d --- /dev/null +++ b/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py @@ -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 \ No newline at end of file diff --git a/lib/xml_schema_validator/python/xml_schema_validator/transformers_logits_processor.py b/lib/xml_schema_validator/python/xml_schema_validator/transformers_logits_processor.py new file mode 100644 index 0000000..8f38b35 --- /dev/null +++ b/lib/xml_schema_validator/python/xml_schema_validator/transformers_logits_processor.py @@ -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 \ No newline at end of file