Fix import for xml_schema_validator

This commit is contained in:
2025-04-09 10:56:20 +00:00
parent 8fb30c8ed0
commit e1da61fd71
3 changed files with 5 additions and 3 deletions

View File

@@ -13,4 +13,6 @@ requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[tool.maturin]
python-source = "python"
module-name = "xml_schema_validator._rs"
features = ["pyo3/extension-module"]

View File

@@ -2,7 +2,7 @@ import torch
import sys
from transformers import LogitsProcessor
from transformers import AutoTokenizer
import xml_schema_validator
from xml_schema_validator._rs import XmlLogitsProcessorCore
class XmlLogitsProcessor(LogitsProcessor):
"""
@@ -26,7 +26,7 @@ class XmlLogitsProcessor(LogitsProcessor):
items = dict()
for token, id in vocab.items():
items[id] = token
self.core = xml_schema_validator.XmlLogitsProcessorCore(items, schema_text)
self.core = XmlLogitsProcessorCore(items, schema_text)
self.prompt_length = None
self.is_first_call = True

View File

@@ -2,7 +2,7 @@ use pyo3::prelude::*;
/// Python module for XML Schema validation
#[pymodule]
fn xml_schema_validator(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
fn _rs(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
crate::py_xml_logits_processor_core::register_module(py, m)?;
crate::py_xml_schema_validator::register_module(py, m)?;
Ok(())