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

Fix char_trie to store all token ids

Browse Source
This commit is contained in:
Niels Geens
2025-04-19 11:49:09 +00:00
parent 33aaaf4455
commit 5ac701316e
6 changed files with 60 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,4 @@
import torch
import sys
from transformers import LogitsProcessor
from transformers import AutoTokenizer
from xml_schema_validator._rs import XmlLogitsProcessorCore
@@ -79,20 +78,6 @@ class XmlLogitsProcessor(LogitsProcessor):
# Find the last occurrence of the assistant start marker
last_marker_pos = full_text.rfind(self.assistant_start_marker)
if last_marker_pos == -1:
# If primary marker not found, try looking for a second common marker as fallback
fallback_markers = ["<|assistant|>", "\nAssistant: ", "\nA: "]
for fallback in fallback_markers:
last_marker_pos = full_text.rfind(fallback)
if last_marker_pos != -1:
# Found a fallback marker
self.assistant_start_marker = fallback # Update for future calls
break
# If still not found, we can't determine where the assistant content starts
if last_marker_pos == -1:
continue
# 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:]

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

@@ -7,8 +7,8 @@ use rayon::prelude::*;
struct CharTrieNode {
// Store character and node pairs, keep sorted for binary search
children: Vec<(char, CharTrieNode)>,
// Store token IDs that are valid at this point
token_id: Option<Arc<Py<PyAny>>>,
// Store token IDs that are valid at this point - CHANGED to Vec<u32>
token_ids: Vec<u32>,
}
impl CharTrieNode {
@@ -24,15 +24,18 @@ impl CharTrieNode {
}
}
// Add a token id to this node
fn add_token(&mut self, token_id: Arc<Py<PyAny>>) {
self.token_id = Some(token_id);
// Add a token id to this node - CHANGED to handle Vec<u32>
fn add_token(&mut self, token_id: u32) {
// Only add if not already present (avoid duplicates)
if !self.token_ids.contains(&token_id) {
self.token_ids.push(token_id);
}
}
fn find_invalid_tokens(
&self,
xml_validator: &XmlSchemaValidator,
) -> Vec<Arc<Py<PyAny>>> {
) -> Vec<u32> {
let res = self
.children
.iter()
@@ -55,7 +58,7 @@ impl CharTrieNode {
fn find_invalid_tokens_parallel(
&self,
xml_validator: &XmlSchemaValidator,
) -> Vec<Arc<Py<PyAny>>> {
) -> Vec<u32> {
// Convert children to a parallel iterator
self.children
.par_iter() // Use parallel iterator from rayon
@@ -76,11 +79,13 @@ impl CharTrieNode {
.collect()
}
fn all_tokens(&self) -> Vec<Arc<Py<PyAny>>> {
// CHANGED to handle Vec<u32>
fn all_tokens(&self) -> Vec<u32> {
let mut tokens = Vec::new();
if let Some(token_id) = &self.token_id {
tokens.push(token_id.clone());
}
// Add all token IDs stored at this node
tokens.extend_from_slice(&self.token_ids);
// Recursively collect token IDs from all children
for (_, child) in &self.children {
tokens.extend(child.all_tokens());
}
@@ -96,7 +101,7 @@ pub struct CharTrie {
impl CharTrie {
// 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: u32) {
let mut current = &mut self.root;
// Add each character in the token text
@@ -105,15 +110,20 @@ impl CharTrie {
}
// Mark this node as representing token_id
current.add_token(Arc::new(token_id));
current.add_token(token_id);
}
// Find all tokens that could be invalid given the current XML state
pub fn find_invalid_tokens(
&self,
xml_validator: &XmlSchemaValidator,
) -> Vec<Arc<Py<PyAny>>> {
) -> Vec<u32> {
// Use the parallel version starting at the root
self.root.find_invalid_tokens_parallel(xml_validator)
}
// NEW: Return all token IDs registered in the trie
pub fn all_registered_tokens(&self) -> Vec<u32> {
self.root.all_tokens()
}
}

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

@@ -8,6 +8,7 @@ use pyo3::types::PyDict;
pub struct XmlLogitsProcessorCore {
xml_schema_validator: XmlSchemaValidator,
trie: char_trie::CharTrie,
tokens: Vec<(u32, String)>,
}
#[pymethods]
@@ -18,17 +19,21 @@ impl XmlLogitsProcessorCore {
Ok(xml_schema_validator) => xml_schema_validator,
Err(e) => return Err(pyo3::exceptions::PyValueError::new_err(e.to_string())),
};
let tokens: Vec<(Py<PyAny>, String)> = tokens
let tokens: Vec<(u32, String)> = tokens
.iter()
.map(|(id, value)| (id.unbind(), value.unbind().to_string()))
.map(|(id, value)| {
let id: u32 = id.extract().unwrap();
(id, value.unbind().to_string())
})
.collect();
let mut trie = char_trie::CharTrie::default();
for (token_id, token_text) in tokens.iter() {
trie.insert(token_text, token_id.clone_ref(py));
trie.insert(token_text, token_id.clone());
}
Ok(Self {
xml_schema_validator,
trie,
tokens,
})
}
@@ -39,13 +44,10 @@ impl XmlLogitsProcessorCore {
}
/// Get a list of tokens that would lead to invalid XML
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<u32>> {
Ok(self
.trie
.find_invalid_tokens(&self.xml_schema_validator)
.iter()
.map(|token_id| token_id.clone_ref(py))
.collect())
.find_invalid_tokens(&self.xml_schema_validator))
}
/// Check if the validator has reached the end of the XML

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

@@ -49,6 +49,9 @@ pub enum Token {
impl Token {
pub fn append(self, c: char) -> Vec<Token> {
if !c.is_ascii_graphic() && ! c.is_ascii_whitespace() {
return vec![];
}
match self {
Token::AttributeEquals(attribute_equals) => attribute_equals.append(c),
Token::AttributeName(attribute_name) => attribute_name.append(c),

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

@@ -3,7 +3,7 @@ from xml_schema_validator import XmlLogitsProcessor
import os
import torch
model_path = "../../model"
model_path = "/root/models/current"
api_token = os.environ["SIA_HF_API_KEY"]
xml_schema_actions = open("example_schema.xsd").read()
@@ -64,6 +64,26 @@ def test_token_masking():
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
processed_scores = processor(input_ids, scores)
# Check if '<' is the only allowed token
less_than_token = tokenizer.encode("<", add_special_tokens=False)[0]
# Count how many tokens have scores not equal to -inf
valid_tokens = (processed_scores[0] > -float('inf')).sum().item()
# Check if the '<' token is allowed (not masked)
is_less_than_allowed = processed_scores[0, less_than_token] > -float('inf')
print(f"Number of allowed tokens: {valid_tokens}")
print(f"Is '<' token allowed: {is_less_than_allowed}")
# Always print the first 10 allowed tokens
allowed_token_ids = torch.where(processed_scores[0] > -float('inf'))[0].tolist()
# Take only the first 10 tokens (or all if less than 10)
allowed_token_ids_sample = allowed_token_ids[:min(10, len(allowed_token_ids))]
allowed_tokens = [tokenizer.decode([token_id]) for token_id in allowed_token_ids_sample]
print(f"First 10 allowed token IDs: {allowed_token_ids_sample}")
print(f"First 10 allowed tokens: {allowed_tokens}")
# Process valid continuation
input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability

2
sia/web/response_websocket.py
View File

@@ -1,4 +1,4 @@
from aiohttp import web
from aiohttp import web, WSMsgType
from ..web_agent import WebAgent
from .util import wrap_async