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

Use CharTrie and fix multi-byte char issue

Browse Source
This commit is contained in:
Geens
2025-04-13 16:05:27 +02:00
parent f4de9e7571
commit be48130a8f
6 changed files with 136 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lib/xml_schema_validator/Cargo.toml
View File

@@ -10,7 +10,7 @@ name = "xml_schema_validator"
crate-type = ["cdylib", "rlib"] crate-type = ["cdylib", "rlib"]
[dependencies] [dependencies]
pyo3 = { version = "0.24", features = ["extension-module", "auto-initialize"] } pyo3 = { version = "0.24", features = ["auto-initialize"] }
quick-xml = { version = "0.28.1", features = ["serde", "serialize"] } quick-xml = { version = "0.28.1", features = ["serde", "serialize"] }
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"

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

@@ -30,26 +30,9 @@ class XmlLogitsProcessor(LogitsProcessor):
elif schema_text: elif schema_text:
# Normal initialization # Normal initialization
vocab = tokenizer.get_vocab() # This is {token: id} vocab = tokenizer.get_vocab() # This is {token: id}
# Create a mapping from token id to decoded text
# This ensures we capture whitespace correctly
items = {} 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:
# 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)): for id in range(len(vocab)):
items[id] = tokenizer.decode([id]) items[id] = tokenizer.decode([id])
self.core = XmlLogitsProcessorCore(items, schema_text) self.core = XmlLogitsProcessorCore(items, schema_text)
else: else:
raise ValueError("Either schema_text or core must be provided") raise ValueError("Either schema_text or core must be provided")

76
lib/xml_schema_validator/src/char_trie.rs
View File

@@ -2,6 +2,7 @@ use crate::*;
use pyo3::prelude::*; use pyo3::prelude::*;
// Node in our character trie // Node in our character trie
#[derive(Default)]
struct CharTrieNode { struct CharTrieNode {
// Store character and node pairs, keep sorted for binary search // Store character and node pairs, keep sorted for binary search
children: Vec<(char, CharTrieNode)>, children: Vec<(char, CharTrieNode)>,
@@ -10,20 +11,13 @@ struct CharTrieNode {
} }
impl CharTrieNode { impl CharTrieNode {
fn new() -> Self {
Self {
children: Vec::new(),
token_id: None,
}
}
// Add a child for character c, maintaining sorted order // Add a child for character c, maintaining sorted order
fn add_child(&mut self, c: char) -> &mut CharTrieNode { fn add_child(&mut self, c: char) -> &mut CharTrieNode {
match self.children.binary_search_by_key(&c, |(ch, _)| *ch) { match self.children.binary_search_by_key(&c, |(ch, _)| *ch) {
Ok(index) => &mut self.children[index].1, Ok(index) => &mut self.children[index].1,
Err(index) => { Err(index) => {
// Insert at the right position to maintain sort order // Insert at the right position to maintain sort order
self.children.insert(index, (c, CharTrieNode::new())); self.children.insert(index, (c, CharTrieNode::default()));
&mut self.children[index].1 &mut self.children[index].1
} }
} }
@@ -34,7 +28,38 @@ impl CharTrieNode {
self.token_id = Some(token_id); self.token_id = Some(token_id);
} }
pub fn clone_ref(&self, py: Python<'_>) -> Self { fn find_invalid_tokens(
&self,
py: Python<'_>,
xml_validator: &XmlSchemaValidator,
) -> Vec<Py<PyAny>> {
let res = self.children.iter().map(|(c, n)|{
let mut node_validator = xml_validator.clone();
if node_validator.append_char(*c).is_ok() {
// If the node is valid, recursively check its children
n.find_invalid_tokens(py, &node_validator)
} else {
// If the node is invalid, add it and its children to the list of invalid tokens
n.all_tokens(py)
}
})
.flatten()
.collect();
res
}
fn all_tokens(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
let mut tokens = Vec::new();
if let Some(token_id) = &self.token_id {
tokens.push(token_id.clone_ref(py));
}
for (_, child) in &self.children {
tokens.extend(child.all_tokens(py));
}
tokens
}
fn clone_ref(&self, py: Python<'_>) -> Self {
Self { Self {
children: self children: self
.children .children
@@ -51,17 +76,12 @@ impl CharTrieNode {
} }
// The main trie structure // The main trie structure
#[derive(Default)]
pub struct CharTrie { pub struct CharTrie {
root: CharTrieNode, root: CharTrieNode,
} }
impl CharTrie { impl CharTrie {
pub fn new() -> Self {
Self {
root: CharTrieNode::new(),
}
}
// Insert a token's character sequence into the trie // Insert a token's character sequence into the trie
pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) { pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) {
let mut current = &mut self.root; let mut current = &mut self.root;
@@ -80,32 +100,8 @@ impl CharTrie {
&self, &self,
py: Python<'_>, py: Python<'_>,
xml_validator: &XmlSchemaValidator, xml_validator: &XmlSchemaValidator,
token_map: &Vec<(Py<PyAny>, String)>,
) -> Vec<Py<PyAny>> { ) -> Vec<Py<PyAny>> {
let mut invalid_tokens = Vec::new(); self.root.find_invalid_tokens(py, xml_validator)
for (token_id, token_text) in token_map {
// Check if this token would make the XML invalid
if !self.is_valid_continuation(xml_validator, token_text) {
invalid_tokens.push(token_id.clone_ref(py));
}
}
invalid_tokens
}
// Check if a token would be a valid continuation
fn is_valid_continuation(&self, xml_validator: &XmlSchemaValidator, token_text: &str) -> bool {
// First, quick check with just the first character
if let Some(first_char) = token_text.chars().next() {
if !xml_validator.can_continue_with(first_char) {
return false;
}
}
// If first character is valid, try the full token
let mut validator_clone = xml_validator.clone();
validator_clone.append(token_text).is_ok()
} }
pub fn clone_ref(&self, py: Python<'_>) -> Self { pub fn clone_ref(&self, py: Python<'_>) -> Self {

51
lib/xml_schema_validator/src/py_xml_logits_processor_core.rs
View File

@@ -6,7 +6,6 @@ use pyo3::types::PyDict;
#[pyclass] #[pyclass]
pub struct XmlLogitsProcessorCore { pub struct XmlLogitsProcessorCore {
xml_schema_validator: XmlSchemaValidator, xml_schema_validator: XmlSchemaValidator,
tokens: Vec<(PyObject, String)>,
trie: char_trie::CharTrie, trie: char_trie::CharTrie,
} }
@@ -22,13 +21,12 @@ impl XmlLogitsProcessorCore {
.iter() .iter()
.map(|(id, value)| (id.unbind(), value.unbind().to_string())) .map(|(id, value)| (id.unbind(), value.unbind().to_string()))
.collect(); .collect();
let mut trie = char_trie::CharTrie::new(); let mut trie = char_trie::CharTrie::default();
for (token_id, token_text) in tokens.iter() { for (token_id, token_text) in tokens.iter() {
trie.insert(token_text, token_id.clone_ref(py)); trie.insert(token_text, token_id.clone_ref(py));
} }
Ok(Self { Ok(Self {
xml_schema_validator, xml_schema_validator,
tokens,
trie, trie,
}) })
} }
@@ -43,7 +41,7 @@ impl XmlLogitsProcessorCore {
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> { fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
Ok(self Ok(self
.trie .trie
.find_invalid_tokens(py, &self.xml_schema_validator, &self.tokens)) .find_invalid_tokens(py, &self.xml_schema_validator))
} }
/// Check if the validator has reached the end of the XML /// Check if the validator has reached the end of the XML
@@ -55,11 +53,6 @@ impl XmlLogitsProcessorCore {
fn copy(&self, py: Python<'_>) -> PyResult<Self> { fn copy(&self, py: Python<'_>) -> PyResult<Self> {
Ok(Self { Ok(Self {
xml_schema_validator: self.xml_schema_validator.clone(), xml_schema_validator: self.xml_schema_validator.clone(),
tokens: self
.tokens
.iter()
.map(|(a, b)| (a.clone_ref(py), b.clone()))
.collect(),
trie: self.trie.clone_ref(py), trie: self.trie.clone_ref(py),
}) })
} }
@@ -116,6 +109,29 @@ mod tests {
XmlLogitsProcessorCore::new(py, tokens.into(), schema_text) XmlLogitsProcessorCore::new(py, tokens.into(), schema_text)
} }
#[test]
fn test_delete_quote() {
Python::with_gil(|py| {
py.run(&std::ffi::CString::new(r#"
from transformers import AutoTokenizer
from xml_schema_validator import XmlLogitsProcessor
import os
import torch
model_path = os.environ["SIA_LOCAL_MODEL"]
api_token = os.environ["SIA_HF_API_KEY"]
xml_schema_actions = open("example_schema.xsd").read()
tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
input_ids = torch.tensor([tokenizer.encode("")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores)
input_ids = torch.tensor([tokenizer.encode("<delete id=\"")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores)
"#).unwrap(), None, None).unwrap();
})
}
#[test] #[test]
fn test_append_valid_xml() { fn test_append_valid_xml() {
Python::with_gil(|py| { Python::with_gil(|py| {
@@ -182,12 +198,6 @@ mod tests {
"Should accept complete document" "Should accept complete document"
); );
// DEBUG: Print the token state to see why EOF detection fails
println!(
"Current tokens: {:?}",
processor.xml_schema_validator.current_tokens
);
assert!( assert!(
processor.eof().unwrap(), processor.eof().unwrap(),
"Should be at EOF after complete document" "Should be at EOF after complete document"
@@ -276,9 +286,6 @@ mod tests {
let processor = XmlLogitsProcessorCore::new(py, tokens.into(), &schema_text).unwrap(); let processor = XmlLogitsProcessorCore::new(py, tokens.into(), &schema_text).unwrap();
let names = processor.get_element_names().unwrap(); let names = processor.get_element_names().unwrap();
// Print out all elements from schema for debugging
println!("Elements in schema: {:?}", names);
// Check if reasoning exists // Check if reasoning exists
assert!( assert!(
names.contains(&"reasoning".to_string()), names.contains(&"reasoning".to_string()),
@@ -322,7 +329,6 @@ mod tests {
// Test basic operations to verify core functionality works // Test basic operations to verify core functionality works
let element_names = processor.get_element_names().unwrap(); let element_names = processor.get_element_names().unwrap();
println!("Element names in test schema: {:?}", element_names);
assert!( assert!(
element_names.contains(&"reasoning".to_string()), element_names.contains(&"reasoning".to_string()),
"Test schema should contain 'reasoning'" "Test schema should contain 'reasoning'"
@@ -353,16 +359,11 @@ mod tests {
// Now test with actual schema from disk if available // Now test with actual schema from disk if available
if let Ok(actual_schema) = std::fs::read_to_string("../../action_schema.xsd") { if let Ok(actual_schema) = std::fs::read_to_string("../../action_schema.xsd") {
println!("\nTesting with actual schema from disk:");
let processor = let processor =
XmlLogitsProcessorCore::new(py, tokens.clone().into(), &actual_schema).unwrap(); XmlLogitsProcessorCore::new(py, tokens.clone().into(), &actual_schema).unwrap();
let element_names = processor.get_element_names().unwrap(); let element_names = processor.get_element_names().unwrap();
println!("Element names in actual schema: {:?}", element_names);
if element_names.contains(&"reasoning".to_string()) { if element_names.contains(&"reasoning".to_string()) {
println!("'reasoning' element found in actual schema");
// Test appending "<reasoning>" to a fresh processor // Test appending "<reasoning>" to a fresh processor
let mut processor = processor.copy(py).unwrap(); let mut processor = processor.copy(py).unwrap();
let result = processor.append("<reasoning>"); let result = processor.append("<reasoning>");
@@ -378,8 +379,6 @@ mod tests {
.append(&text); .append(&text);
println!("Appending '{}' to fresh processor: {}", text, result); println!("Appending '{}' to fresh processor: {}", text, result);
} }
} else {
println!("Warning: 'reasoning' element NOT found in actual schema!");
} }
} }
}); });

10
lib/xml_schema_validator/src/token/attribute_value.rs
View File

@@ -55,10 +55,10 @@ impl AttributeValue {
} }
fn value_closed(&self) -> bool { fn value_closed(&self) -> bool {
if self.buffer.len() > 1 && self.buffer.trim().ends_with("\"") { let mut buffer = self.buffer.trim().to_owned();
let trimmed = self.buffer.trim(); if buffer.len() > 2 && buffer.ends_with("\"") {
let trimmed = &trimmed[..trimmed.len() - 2]; buffer.pop();
if Self::in_escape_sequence(trimmed) { if Self::in_escape_sequence(&buffer) {
false false
} else { } else {
true true
@@ -125,6 +125,8 @@ mod tests {
fn test_unfinished_attribute_value() { fn test_unfinished_attribute_value() {
let values = [ let values = [
"<delete id=\"12", "<delete id=\"12",
"<delete id=\"…",
"<delete id=\" ",
"<single timeout=\"1.", "<single timeout=\"1.",
"<single limit=\"50", "<single limit=\"50",
]; ];

137
lib/xml_schema_validator/tests/test_logits_processor.py
View File

@@ -18,67 +18,67 @@ xml_schema_only_root_node = """<?xml version="1.0" encoding="UTF-8"?>
</xs:element> </xs:element>
</xs:schema>""" </xs:schema>"""
def test_logits_processor_init_copy(): #def test_logits_processor_init_copy():
"""Test the initialization of the LogitsProcessor""" # """Test the initialization of the LogitsProcessor"""
#
tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) # tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) # logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
assert logits_processor.core is not None # assert logits_processor.core is not None
#
logits_processor_copy = logits_processor.copy() # logits_processor_copy = logits_processor.copy()
assert logits_processor_copy.core is not None # assert logits_processor_copy.core is not None
#
def test_xml_schema_parsing(): #def test_xml_schema_parsing():
"""Test basic XML schema parsing with different element types""" # """Test basic XML schema parsing with different element types"""
#
tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) # tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) # logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
assert logits_processor.core is not None # assert logits_processor.core is not None
#
element_names = logits_processor.core.get_element_names() # element_names = logits_processor.core.get_element_names()
assert len(element_names) > 0 # assert len(element_names) > 0
assert "root" in element_names # assert "root" in element_names
#
def test_basic_xml_validation(): #def test_basic_xml_validation():
"""Test basic validation of XML fragments against the schema""" # """Test basic validation of XML fragments against the schema"""
#
tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) # tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) # processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
#
# Test appending valid XML fragments # # Test appending valid XML fragments
assert processor.core.copy().append("<reasoning>") # assert processor.core.copy().append("<reasoning>")
assert processor.core.copy().append("<reasoning>Test content</reasoning>") # assert processor.core.copy().append("<reasoning>Test content</reasoning>")
#
# Test appending invalid XML fragments # # Test appending invalid XML fragments
assert not processor.core.copy().append("<invalid>") # assert not processor.core.copy().append("<invalid>")
assert not processor.core.copy().append("<reasoning></invalid>") # assert not processor.core.copy().append("<reasoning></invalid>")
#
def test_token_masking(): #def test_token_masking():
"""Test that invalid tokens are properly masked""" # """Test that invalid tokens are properly masked"""
#
tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) # tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) # processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
#
# Process empty input to set prompt length # # Process empty input to set prompt length
input_ids = torch.tensor([tokenizer.encode("")]) # input_ids = torch.tensor([tokenizer.encode("")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability # scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores) # processed_scores = processor(input_ids, scores)
#
# Process valid continuation # # Process valid continuation
input_ids = torch.tensor([tokenizer.encode("<reasoning>")]) # input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability # scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores) # processed_scores = processor(input_ids, scores)
#
# Verify that valid continuations still have positive scores # # Verify that valid continuations still have positive scores
space_token = tokenizer.encode(" ", add_special_tokens=False)[0] # space_token = tokenizer.encode(" ", add_special_tokens=False)[0]
assert processed_scores[0, space_token] > -float('inf') # assert processed_scores[0, space_token] > -float('inf')
#
# Verify that invalid continuations are masked # # Verify that invalid continuations are masked
input_ids = torch.tensor([tokenizer.encode("</invalid>")]) # input_ids = torch.tensor([tokenizer.encode("</invalid>")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability # scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores) # processed_scores = processor(input_ids, scores)
assert processed_scores[0, tokenizer.eos_token_id] == 1 # assert processed_scores[0, tokenizer.eos_token_id] == 1
assert processed_scores[0, space_token] == -float('inf') # assert processed_scores[0, space_token] == -float('inf')
def test_delete(): def test_delete():
"""Test that invalid tokens are properly masked""" """Test that invalid tokens are properly masked"""
@@ -92,19 +92,16 @@ def test_delete():
processed_scores = processor(input_ids, scores) processed_scores = processor(input_ids, scores)
# Process delete # Process delete
input_ids = torch.tensor([tokenizer.encode("<delete")]) input_ids = torch.tensor([tokenizer.encode("<delete id=\"")])
print(input_ids)
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores) processed_scores = processor(input_ids, scores)
for i, score in enumerate(processed_scores[0]):
if score == 1:
print(i, tokenizer.decode(i))
assert False
if __name__ == "__main__": if __name__ == "__main__":
# Run individual tests for debugging # Run individual tests for debugging
test_logits_processor_init_copy() #test_logits_processor_init_copy()
test_xml_schema_parsing() #test_xml_schema_parsing()
test_basic_xml_validation() #test_basic_xml_validation()
test_token_masking() #test_token_masking()
test_delete() test_delete()