Optimized hot-loop for logits processor

lib/xml_schema_validator/tests/conftest.py

@@ -0,0 +1,9 @@
+import subprocess
+import sys
+
+def pytest_sessionstart(session):
+    """
+    Build the Rust extension module before running tests.
+    """
+    print("Installing test dependencies...")
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "test_requirements.txt"])

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -0,0 +1,71 @@
+from xml_logits_processor import XmlLogitsProcessor
+from threading import Thread
+from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
+import os
+import sys
+
+def test_logits_processor():
+    print("test_logits_processor")
+
+    model_name = "../../model"
+    xml_schema_text = """<?xml version="1.0" encoding="UTF-8"?>
+<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+  <xs:element name="root">
+    <xs:complexType mixed="true">
+      <xs:sequence>
+        <xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
+      </xs:sequence>
+    </xs:complexType>
+  </xs:element>
+</xs:schema>"""
+
+    model = "google/gemma-3-1b-it"
+
+    pipline = pipeline(
+        "text-generation",
+        model=model,
+        token=os.environ["SIA_HF_API_KEY"],
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model,
+        token=os.environ["SIA_HF_API_KEY"],
+    )
+
+    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)
+                        
+    messages = [
+        {"role": "system", "content": "Always respond with <root>message by the user</root>. So if the user says 'hello world', you response <root>hello world</root>"},
+        {"role": "user", "content": "hello, I am the user"}
+    ]
+
+    text_inputs = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True
+    )
+
+    generation_kwargs = {
+        "text_inputs": text_inputs,
+        "max_new_tokens": 20,
+        "streamer": streamer,
+        #"logits_processor": [logits_processor],
+    }
+
+    generation_thread = Thread(
+        target=pipline,
+        kwargs=generation_kwargs
+    )
+
+    generation_thread.start()
+
+    for text in streamer:
+        print(text, end="", file=sys.stderr, flush=True)
+
+    generation_thread.join()
+
+if __name__ == "__main__":
+    test_logits_processor()