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

Validator works for start

Browse Source
This commit is contained in:
Niels Geens
2025-04-07 14:31:56 +00:00
parent 50539f6d03
commit 2e08be4b75
16 changed files with 108 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lib/xml_schema_validator/src/lib.rs
View File

@@ -14,7 +14,7 @@ use token::*;
/// An XML validator that checks XML fragments against an XSD schema
#[derive(Clone, Debug)]
pub struct XmlSchemaValidator {
current_tokens: Vec<Box<Token>>,
current_tokens: Vec<Token>,
}
/// The Result type used throughout this crate
@@ -27,7 +27,7 @@ impl XmlSchemaValidator {
pub fn new(schema_text: &str) -> Result<Self> {
let schema: schema::XsSchema = quick_xml::de::from_str(schema_text)?;
let current_tokens = schema.xs_element.iter().map(|element| {
Box::new(Token::StartOfFile(StartOfFile::new(Box::new(element.clone()))))
Token::StartOfFile(StartOfFile::new(element.clone()))
}).collect();
Ok(Self {
@@ -45,7 +45,6 @@ impl XmlSchemaValidator {
/// Append a single character to the validator
fn append_char(&mut self, c: char) -> Result<()> {
println!("Appending char: {}", c);
let mut new_tokens = vec![];
for token in self.current_tokens.clone() {
new_tokens.append(&mut token.append(c));

10
lib/xml_schema_validator/src/python.rs
View File

@@ -2,7 +2,7 @@ use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// Python wrapper for XmlValidator
#[pyclass(unsendable)]
#[pyclass]
struct XmlSchemaValidator {
validator: crate::XmlSchemaValidator,
}
@@ -17,9 +17,11 @@ impl XmlSchemaValidator {
Ok(Self { validator })
}
fn append(&mut self, fragment: &str) -> PyResult<()> {
self.validator.append(fragment)
.map_err(|e| PyValueError::new_err(format!("Validation error: {}", e)))
fn append(&mut self, fragment: &str) -> PyResult<(bool, Option<String>)> {
match self.validator.append(fragment) {
Ok(()) => PyResult::Ok((true, None)),
Err(e) => PyResult::Ok((false, Some(e.to_string()))),
}
}
fn copy(&self) -> PyResult<Self> {

4
lib/xml_schema_validator/src/token/element_close_end.rs
View File

@@ -11,9 +11,9 @@ impl ElementCloseEnd {
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
if '>' == c {
vec![Box::new(Token::EndOfFile(EndOfFile::new()))]
vec![Token::EndOfFile(EndOfFile::new())]
} else {
vec![]
}

20
lib/xml_schema_validator/src/token/element_close_name.rs
View File

@@ -5,21 +5,21 @@ use crate::*;
#[derive(Clone, Debug)]
pub struct ElementCloseName {
/// The element that was previously opened and is now being closed
element: Box<schema::XsElement>,
element: schema::XsElement,
/// Current buffer of characters processed for the element name
buffer: String,
}
impl ElementCloseName {
/// Create a new instance with the given buffer
pub fn new(element: Box<schema::XsElement>, buffer: String) -> Self {
pub fn new(element: schema::XsElement, buffer: String) -> Self {
Self {
element,
buffer,
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
let new_buffer = self.buffer.clone() + &c.to_string();
// If still building the element name
@@ -29,19 +29,19 @@ impl ElementCloseName {
// Now we're expecting '>' to close the tag
if c.is_whitespace() {
// Handle whitespace after the element name
vec![Box::new(Token::Whitespace(Whitespace::new(
vec![Box::new(Token::ElementCloseEnd(ElementCloseEnd::new()))]
)))]
vec![Token::Whitespace(Whitespace::new(
vec![Token::ElementCloseEnd(ElementCloseEnd::new())]
))]
} else {
// Directly transition to the closed state
vec![Box::new(Token::ElementCloseEnd(ElementCloseEnd::new()))]
vec![Token::ElementCloseEnd(ElementCloseEnd::new())]
}
} else {
// Still building the element name
vec![Box::new(Token::ElementCloseName(ElementCloseName::new(
Box::new((*self.element).clone()),
vec![Token::ElementCloseName(ElementCloseName::new(
self.element.clone(),
new_buffer,
)))]
))]
}
} else {
// The character doesn't match what we expect for this element name

18
lib/xml_schema_validator/src/token/element_close_start.rs
View File

@@ -4,30 +4,30 @@ use crate::*;
#[derive(Clone, Debug)]
pub struct ElementCloseStart {
/// The element that was previously opened and is now being closed
element: Box<schema::XsElement>,
element: schema::XsElement,
has_slash: bool,
}
impl ElementCloseStart {
pub fn new(element: Box<schema::XsElement>) -> Self {
pub fn new(element: schema::XsElement) -> Self {
Self {
element,
has_slash: false,
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
if !self.has_slash && '/' == c {
vec![Box::new(Token::ElementCloseStart(Self {
element: Box::new((*self.element).clone()),
vec![Token::ElementCloseStart(Self {
element: self.element.clone(),
has_slash: true,
}))]
})]
} else {
if self.element.name.starts_with(c) {
vec![Box::new(Token::ElementCloseName(ElementCloseName::new(
Box::new((*self.element).clone()),
vec![Token::ElementCloseName(ElementCloseName::new(
self.element.clone(),
c.to_string(),
)))]
))]
} else {
vec![]
}

14
lib/xml_schema_validator/src/token/element_open_end.rs
View File

@@ -3,17 +3,17 @@ use crate::*;
/// Represents a fully opened XML element, containing the schema element reference
#[derive(Clone, Debug)]
pub struct ElementOpenEnd {
element: Box<schema::XsElement>,
element: schema::XsElement,
}
impl ElementOpenEnd {
pub fn new(element: Box<schema::XsElement>) -> Self {
pub fn new(element: schema::XsElement) -> Self {
Self {
element,
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
// Check if the element allows text content (mixed="true")
let allows_text = self.element.xs_complex_type.mixed
.as_ref()
@@ -22,11 +22,11 @@ impl ElementOpenEnd {
if allows_text {
if c.is_whitespace() {
vec![Box::new(Token::Whitespace(Whitespace::new(
vec![Box::new(Token::TextContent(TextContent::new(Box::new((*self.element).clone()))))]
)))]
vec![Token::Whitespace(Whitespace::new(
vec![Token::TextContent(TextContent::new(self.element.clone()))]
))]
} else {
vec![Box::new(Token::TextContent(TextContent::new(Box::new((*self.element).clone()))))]
vec![Token::TextContent(TextContent::new(self.element.clone()))]
}
} else {
vec![]

20
lib/xml_schema_validator/src/token/element_open_name.rs
View File

@@ -4,11 +4,11 @@ use crate::*;
#[derive(Clone, Debug)]
pub struct ElementOpenName {
buffer: String,
element: Box<schema::XsElement>,
element: schema::XsElement,
}
impl ElementOpenName {
pub fn new(element: Box<schema::XsElement>, buffer: String) -> Result<Self> {
pub fn new(element: schema::XsElement, buffer: String) -> Result<Self> {
if element.name.starts_with(buffer.as_str()) {
Ok(Self {
buffer,
@@ -19,19 +19,19 @@ impl ElementOpenName {
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
let buffer = self.buffer.clone() + &c.to_string();
if self.element.name.starts_with(buffer.as_str()) {
vec![Box::new(Token::ElementName(ElementOpenName {
vec![Token::ElementName(ElementOpenName {
buffer,
element: Box::new((*self.element).clone()),
}))]
element: self.element.clone(),
})]
} else if c.is_whitespace() {
vec![Box::new(Token::Whitespace(Whitespace::new(
vec![Box::new(Token::ElementOpen(ElementOpenEnd::new(Box::new((*self.element).clone()))))]
)))]
vec![Token::Whitespace(Whitespace::new(
vec![Token::ElementOpen(ElementOpenEnd::new(self.element.clone()))]
))]
} else if '>' == c {
vec![Box::new(Token::ElementOpen(ElementOpenEnd::new(Box::new((*self.element).clone()))))]
vec![Token::ElementOpen(ElementOpenEnd::new(self.element.clone()))]
} else {
vec![]
}

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

@@ -3,20 +3,20 @@ use crate::*;
/// Represents an XML element that is in the process of being opened, containing the text buffer and schema element reference
#[derive(Clone, Debug)]
pub struct ElementOpenStart {
element: Box<schema::XsElement>,
element: schema::XsElement,
}
impl ElementOpenStart {
pub fn new(element: Box<schema::XsElement>) -> Self {
pub fn new(element: schema::XsElement) -> Self {
Self {
element,
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
let element_name = ElementOpenName::new(Box::new((*self.element).clone()), c.to_string());
pub fn append(&self, c: char) -> Vec<Token> {
let element_name = ElementOpenName::new(self.element.clone(), c.to_string());
if let Ok(element_name) = element_name {
vec![Box::new(Token::ElementName(element_name))]
vec![Token::ElementName(element_name)]
} else {
vec![]
}

2
lib/xml_schema_validator/src/token/end_of_file.rs
View File

@@ -10,7 +10,7 @@ impl EndOfFile {
Self {}
}
pub fn append(&self, _c: char) -> Vec<Box<Token>> {
pub fn append(&self, _c: char) -> Vec<Token> {
vec![]
}
}

2
lib/xml_schema_validator/src/token/mod.rs
View File

@@ -36,7 +36,7 @@ pub enum Token {
}
impl Token {
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
match self {
Token::ElementCloseEnd(element_close_end) => element_close_end.append(c),
Token::ElementCloseName(element_close_name) => element_close_name.append(c),

8
lib/xml_schema_validator/src/token/start_of_file.rs
View File

@@ -3,17 +3,17 @@ use crate::*;
/// Represents the initial state at the start of an XML file, containing a reference to the root element schema
#[derive(Clone, Debug)]
pub struct StartOfFile {
element: Box<schema::XsElement>,
element: schema::XsElement,
}
impl StartOfFile {
pub fn new(element: Box<schema::XsElement>) -> Self {
pub fn new(element: schema::XsElement) -> Self {
Self { element }
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
if '<' == c {
vec![Box::new(Token::ElementStart(ElementOpenStart::new(Box::new((*self.element).clone()))))]
vec![Token::ElementStart(ElementOpenStart::new(self.element.clone()))]
} else {
vec![]
}

16
lib/xml_schema_validator/src/token/text_content.rs
View File

@@ -3,12 +3,12 @@ use crate::*;
/// Represents the text content within an XML element
#[derive(Clone, Debug)]
pub struct TextContent {
element: Box<schema::XsElement>,
element: schema::XsElement,
buffer: String,
}
impl TextContent {
pub fn new(element: Box<schema::XsElement>) -> Self {
pub fn new(element: schema::XsElement) -> Self {
Self {
element,
buffer: String::new(),
@@ -16,21 +16,21 @@ impl TextContent {
}
/// Create a new TextContent with the given initial buffer
pub fn with_buffer(element: Box<schema::XsElement>, buffer: String) -> Self {
pub fn with_buffer(element: schema::XsElement, buffer: String) -> Self {
Self {
element,
buffer,
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
if c == '<' {
vec![Box::new(Token::ElementCloseStart(ElementCloseStart::new(Box::new((*self.element).clone()))))]
vec![Token::ElementCloseStart(ElementCloseStart::new(self.element.clone()))]
} else {
vec![Box::new(Token::TextContent(TextContent::with_buffer(
Box::new((*self.element).clone()),
vec![Token::TextContent(TextContent::with_buffer(
self.element.clone(),
self.buffer.clone() + &c.to_string(),
)))]
))]
}
}
}

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

@@ -3,21 +3,21 @@ use crate::*;
/// Represents one or more whitespace characters
#[derive(Clone, Debug)]
pub struct Whitespace {
next: Vec<Box<Token>>
next: Vec<Token>
}
impl Whitespace {
pub fn new(next: Vec<Box<Token>>) -> Self {
pub fn new(next: Vec<Token>) -> Self {
Self {
next
}
}
pub fn append(&self, c: char) -> Vec<Box<Token>> {
pub fn append(&self, c: char) -> Vec<Token> {
if c.is_whitespace() {
vec![Box::new(Token::Whitespace(Self {
vec![Token::Whitespace(Self {
next: self.next.clone()
}))]
})]
} else {
let mut tokens = vec![];
for token in self.next.iter() {

1
sia/__main__.py
View File

@@ -67,6 +67,7 @@ class Main:
config.qwq_model,
config.qwq_temperature,
config.qwq_token_limit,
self._action_schema,
)
if not self._llms:

9
sia/llm_engine/qwq_llm_engine.py
View File

@@ -29,8 +29,6 @@ class QwQLlmEngine(LlmEngine):
"""
self._temperature = temperature
self._token_limit = token_limit
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
tokenizer_config = json.load(f)
@@ -62,6 +60,11 @@ class QwQLlmEngine(LlmEngine):
return_full_text=False,
)
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
else:
self._logits_processor = None
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
@@ -100,7 +103,7 @@ class QwQLlmEngine(LlmEngine):
}
if self._logits_processor:
generation_kwargs["logits_processor"] = self.logits_processor
generation_kwargs["logits_processor"] = [self._logits_processor]
generation_thread = Thread(
target=self._pipline,

39
sia/llm_engine/xml_logits_processor.py
View File

@@ -20,6 +20,8 @@ class XmlLogitsProcessor(LogitsProcessor):
"""
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:
"""
@@ -34,29 +36,48 @@ class XmlLogitsProcessor(LogitsProcessor):
"""
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 the current text generated so far
# Get only the generated portion of the text
current_ids = input_ids[batch_idx]
current_text = self.tokenizer.decode(current_ids)
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()
# For each possible next token
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):
# 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_validator = batch_validator.copy()
token_text = self.tokenizer.decode([token_idx])
if validator_copy.append(token_text):
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
# 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')