From 514ca643a17cebf85d99da514da6b11b339197f2 Mon Sep 17 00:00:00 2001 From: xyf2020 <75460675+xyf2020@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:58:35 +0800 Subject: [PATCH] feat(token counter): add Qwen token counting utility with fallback (#616) * feat: add Qwen token counting utility with fallback * fix: preserve Qwen tokenizer text prefixes and isolate loading * fix: isolate injected Qwen tokenizers and refresh local sources * perf: avoid unnecessary Qwen tokenizer initialization and locking --------- Co-authored-by: jinli.yl --- .github/workflows/ci-python-tests.yml | 10 +- pyproject.toml | 4 + reme/utils/__init__.py | 2 + reme/utils/qwen_token_tools.py | 147 +++++++++++++++ tests/unit/test_package_versions.py | 10 + tests/unit/test_qwen_token_tools.py | 259 ++++++++++++++++++++++++++ 6 files changed, 431 insertions(+), 1 deletion(-) create mode 100644 reme/utils/qwen_token_tools.py create mode 100644 tests/unit/test_qwen_token_tools.py diff --git a/.github/workflows/ci-python-tests.yml b/.github/workflows/ci-python-tests.yml index 1d3971cb..82b9be2b 100644 --- a/.github/workflows/ci-python-tests.yml +++ b/.github/workflows/ci-python-tests.yml @@ -38,7 +38,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip setuptools wheel - pip install -e reme_studio -e ".[dev,core,classifier-jev]" + pip install -e reme_studio -e ".[dev,core,classifier-jev,qwen-tokenizer]" pip install --no-deps -e plugins/auto-fin pip install -e plugins/daily_paper pip install -e plugins/lme -e plugins/beam @@ -46,6 +46,7 @@ jobs: - name: Run unit tests run: | + python -c "import tokenizers; print('tokenizers', tokenizers.__version__)" coverage run -m pytest tests/unit plugins/auto-fin plugins/daily_paper plugins/lme plugins/beam \ -v \ --tb=long \ @@ -54,3 +55,10 @@ jobs: - name: Generate coverage report run: coverage report -m + + - name: Verify minimum tokenizer version + if: matrix.python-version == '3.11' + run: | + python -m pip install "tokenizers==0.21.0" + python -c "import tokenizers; print('tokenizers', tokenizers.__version__)" + python -m pytest tests/unit/test_qwen_token_tools.py tests/unit/test_package_versions.py -q diff --git a/pyproject.toml b/pyproject.toml index 6adfc24f..97cc3588 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,9 @@ core = [ "polars>=1.43.0", "reme_studio", ] +qwen-tokenizer = [ + "tokenizers>=0.21,<1", +] image-heif = [ "pillow>=10.0.0", "pillow-heif>=0.13.0", @@ -77,6 +80,7 @@ dev = [ ] full = [ "reme-ai[core]", + "reme-ai[qwen-tokenizer]", "reme-ai[classifier-jev]", "reme-ai[dev]", "reme-ai[image-heif]", diff --git a/reme/utils/__init__.py b/reme/utils/__init__.py index 60b68968..5cdbfafc 100644 --- a/reme/utils/__init__.py +++ b/reme/utils/__init__.py @@ -22,6 +22,7 @@ from .service_utils import ( ) from .similarity_utils import cosine_similarity, batch_cosine_similarity from .token_utils import estimate_token_count +from .qwen_token_tools import QwenTokenTools from .web_static import REME_WEB_STATIC_DIR, resolve_web_static_dir from .agent_state_io import AsStateHandler from .counter import ( @@ -55,6 +56,7 @@ __all__ = [ "cosine_similarity", "batch_cosine_similarity", "estimate_token_count", + "QwenTokenTools", "REME_WEB_STATIC_DIR", "resolve_web_static_dir", "AsStateHandler", diff --git a/reme/utils/qwen_token_tools.py b/reme/utils/qwen_token_tools.py new file mode 100644 index 00000000..72cd1ae8 --- /dev/null +++ b/reme/utils/qwen_token_tools.py @@ -0,0 +1,147 @@ +"""Local Qwen token counting and truncation utilities.""" + +from __future__ import annotations + +import warnings +from functools import lru_cache +from pathlib import Path +from threading import Lock +from typing import Any + +from .token_utils import estimate_token_count + +DEFAULT_TOKENIZER = "Qwen/Qwen3.8-27B" + + +def _estimated_count(text: str) -> int: + return max(1, estimate_token_count(text)) if text else 0 + + +class QwenTokenTools: + """Count and truncate text with a lazily loaded Hugging Face tokenizer. + + Install ``reme-ai[qwen-tokenizer]`` for exact counts. The default source + downloads tokenizer data only; select the inference model's source when + model-specific token counts are needed. Truncation preserves a prefix of + the original string, including when the tokenizer normalizes its input. + Injected Hugging Face tokenizers are copied with truncation and padding disabled. + Local sources are read on first use for each instance, so new instances see file updates. + """ + + def __init__(self, tokenizer_source: str = DEFAULT_TOKENIZER, *, tokenizer: Any | None = None) -> None: + if not isinstance(tokenizer_source, str) or not tokenizer_source.strip(): + raise ValueError("tokenizer_source must be a non-empty string") + self.tokenizer_source = tokenizer_source.strip() + if tokenizer is not None: + try: + from tokenizers import Tokenizer + from tokenizers.implementations import BaseTokenizer + except ImportError: + pass # Custom tokenizer implementations do not require this package. + else: + if isinstance(tokenizer, (Tokenizer, BaseTokenizer)): + tokenizer = Tokenizer.from_str(tokenizer.to_str()) + tokenizer.no_truncation() + tokenizer.no_padding() + self._tokenizer = tokenizer + self._tokenizer_lock = Lock() + self._tokenizer_unavailable = False + + @property + def tokenizer(self) -> Any | None: + """Load once under a lock; warn once and return None if unavailable.""" + with self._tokenizer_lock: + if self._tokenizer_unavailable: + return None + try: + if self._tokenizer is None: + self._tokenizer = _load_tokenizer(self.tokenizer_source) + return self._tokenizer + except (ImportError, RuntimeError) as exc: + self._tokenizer_unavailable = True + reason = "missing 'tokenizers' package" if isinstance(exc, ImportError) else "load failed" + warnings.warn( + f"Qwen tokenizer is unavailable ({reason}); token count is estimated from UTF-8 byte length. " + "Install exact tokenization support with: pip install 'reme-ai[qwen-tokenizer]'.", + RuntimeWarning, + stacklevel=3, + ) + return None + + def count(self, text: str) -> int: + """Count tokenizer IDs, or estimate from bytes when the tokenizer is unavailable.""" + if not isinstance(text, str): + raise TypeError("text must be a string") + if not text: + return 0 + tokenizer = self.tokenizer + if tokenizer is None: + return _estimated_count(text) + return len(tokenizer.encode(text, add_special_tokens=False).ids) + + def truncate(self, text: str, max_tokens: int) -> tuple[str, int, bool]: + """Return an original-text prefix, its token count, and whether text changed.""" + if not isinstance(text, str): + raise TypeError("text must be a string") + if not isinstance(max_tokens, int) or isinstance(max_tokens, bool) or max_tokens < 1: + raise ValueError("max_tokens must be a positive integer") + if not text: + return "", 0, False + tokenizer = self.tokenizer + if tokenizer is None: + token_count = _estimated_count(text) + if token_count <= max_tokens: + return text, token_count, False + low, high = 0, len(text) + while low < high: + middle = (low + high + 1) // 2 + if _estimated_count(text[:middle]) <= max_tokens: + low = middle + else: + high = middle - 1 + truncated = text[:low] + return truncated, _estimated_count(truncated), True + + encoding = tokenizer.encode(text, add_special_tokens=False) + token_count = len(encoding.ids) + if token_count <= max_tokens: + return text, token_count, False + + # Byte-level tokens can share character offsets. Re-encode the original + # prefix to check the budget instead of decoding a partial character. + end = encoding.offsets[max_tokens - 1][1] + truncated = text[:end] + token_count = len(tokenizer.encode(truncated, add_special_tokens=False).ids) + while truncated and token_count > max_tokens: + truncated = truncated[:-1] + token_count = len(tokenizer.encode(truncated, add_special_tokens=False).ids) + return truncated, token_count, True + + +def _load_tokenizer(source: str) -> Any: + """Read local files afresh; build independent instances from cached remote data.""" + try: + from tokenizers import Tokenizer + except ImportError as exc: + raise ImportError("Install 'reme-ai[qwen-tokenizer]' for exact tokenization") from exc + try: + local_path = Path(source) + if local_path.is_dir(): + tokenizer = Tokenizer.from_file(str(local_path / "tokenizer.json")) + elif local_path.is_file(): + tokenizer = Tokenizer.from_file(str(local_path)) + else: + tokenizer = Tokenizer.from_str(_load_remote_tokenizer_json(source)) + tokenizer.no_truncation() + tokenizer.no_padding() + return tokenizer + except Exception as exc: + raise RuntimeError(f"Unable to load Qwen tokenizer {source!r}: {exc}") from exc + + +@lru_cache(maxsize=4) +def _load_remote_tokenizer_json(source: str) -> str: + """Cache remote serialized data, never local files or mutable tokenizer objects.""" + from tokenizers import Tokenizer + + return Tokenizer.from_pretrained(source).to_str() diff --git a/tests/unit/test_package_versions.py b/tests/unit/test_package_versions.py index 7f451cad..07855130 100644 --- a/tests/unit/test_package_versions.py +++ b/tests/unit/test_package_versions.py @@ -238,3 +238,13 @@ def test_studio_package_preparation_preserves_static_gitignore(monkeypatch, tmp_ assert (static_dir / "index.html").is_file() assert (static_dir / ".gitignore").read_text(encoding="utf-8") == package_studio.STATIC_GITIGNORE + + +def test_qwen_tokenizer_dependency_is_optional() -> None: + """Exact Qwen tokenization is opt-in and included in the full install.""" + config = tomllib.loads((REPOSITORY / "pyproject.toml").read_text(encoding="utf-8")) + extras = config["project"]["optional-dependencies"] + assert extras["qwen-tokenizer"] == ["tokenizers>=0.21,<1"] + assert "tokenizers" not in {Requirement(value).name for value in config["project"]["dependencies"]} + assert "tokenizers" not in {Requirement(value).name for value in extras["core"]} + assert "reme-ai[qwen-tokenizer]" in extras["full"] diff --git a/tests/unit/test_qwen_token_tools.py b/tests/unit/test_qwen_token_tools.py new file mode 100644 index 00000000..1a8e1844 --- /dev/null +++ b/tests/unit/test_qwen_token_tools.py @@ -0,0 +1,259 @@ +"""Qwen tokenizer utility tests without external downloads.""" + +import sys +import warnings +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from time import sleep +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from reme.utils import QwenTokenTools, estimate_token_count +from reme.utils import qwen_token_tools + + +class _Tokenizer: + def encode(self, text, add_special_tokens): + """Return one fake token per character.""" + assert add_special_tokens is False + return SimpleNamespace(ids=list(text), offsets=[(index, index + 1) for index in range(len(text))]) + + +def test_custom_tokenizer_without_optional_dependency(monkeypatch): + """Custom counting and truncation do not require the optional package.""" + monkeypatch.setitem(sys.modules, "tokenizers", None) + tools = QwenTokenTools(tokenizer=_Tokenizer()) + + assert tools.count("你好 world") == 8 + assert tools.truncate("你好 world", 3) == ("你好 ", 3, True) + assert tools.truncate("hi", 3) == ("hi", 2, False) + + +@pytest.fixture(name="word_tokenizer") +def _word_tokenizer(): + """Build a small tokenizer for loading and isolation regressions.""" + tokenizers = pytest.importorskip("tokenizers") + tokenizer = tokenizers.Tokenizer( + tokenizers.models.WordLevel( + {"[UNK]": 0, "[PAD]": 1, "one": 2, "two": 3, "three": 4, "four": 5}, + unk_token="[UNK]", + ), + ) + tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Whitespace() + return tokenizer + + +@pytest.mark.parametrize("source_kind", ["native", "wrapped", "file", "directory", "remote"]) +def test_tokenizer_sources_disable_limits_and_isolate_instances(word_tokenizer, tmp_path, monkeypatch, source_kind): + """Every source removes saved limits and isolates caller and instance mutations.""" + tokenizers = pytest.importorskip("tokenizers") + tokenizer = word_tokenizer + if source_kind == "wrapped": + tokenizer = tokenizers.implementations.BaseTokenizer(tokenizer) + tokenizer.enable_truncation(max_length=1) + tokenizer.enable_padding(length=8, pad_id=1, pad_token="[PAD]") + original = tokenizer.to_str() + if source_kind in ("native", "wrapped"): + kwargs = {"tokenizer": tokenizer} + else: + path = tmp_path / "tokenizer.json" + tokenizer.save(str(path)) + source = str(path if source_kind == "file" else tmp_path) + if source_kind == "remote": + source = f"test/{tmp_path.name}" + load = Mock(return_value=tokenizer) + monkeypatch.setattr(tokenizers.Tokenizer, "from_pretrained", load) + kwargs = {"tokenizer_source": source} + first, second = QwenTokenTools(**kwargs), QwenTokenTools(**kwargs) + text = "one two three four" + for tools in (first, second): + assert tools.count(text) == 4 + assert tools.count("") == 0 + assert tools.truncate(text, 1) == ("one", 1, True) + assert tools.tokenizer is not tokenizer + assert first.tokenizer is not second.tokenizer + assert tokenizer.to_str() == original + + tokenizer.enable_truncation(max_length=2) + tokenizer.enable_padding(length=10, pad_id=1, pad_token="[PAD]") + first.tokenizer.enable_truncation(max_length=1) + first.tokenizer.enable_padding(length=8, pad_id=1, pad_token="[PAD]") + assert tokenizer.truncation["max_length"] == 2 + assert tokenizer.padding["length"] == 10 + assert second.count(text) == 4 + assert second.count("") == 0 + assert second.truncate(text, 1) == ("one", 1, True) + if source_kind == "remote": + load.assert_called_once_with(source) + + +@pytest.mark.parametrize( + "failure, reason", + [ + ("missing_dependency", "missing 'tokenizers' package"), + ("invalid_json", "load failed"), + ("remote", "load failed"), + ], +) +def test_unavailable_tokenizer_warns_once_and_estimates(tmp_path, monkeypatch, failure, reason): + """Real import, local JSON, and remote failures warn once and fall back.""" + source = f"test/{tmp_path.name}" + if failure == "missing_dependency": + monkeypatch.setitem(sys.modules, "tokenizers", None) + else: + tokenizers = pytest.importorskip("tokenizers") + if failure == "invalid_json": + path = tmp_path / "tokenizer.json" + path.write_text("invalid JSON", encoding="utf-8") + source = str(path) + else: + load = Mock(side_effect=OSError("offline")) + monkeypatch.setattr(tokenizers.Tokenizer, "from_pretrained", load) + tools = QwenTokenTools(source) + + with pytest.warns(RuntimeWarning, match=reason) as caught: + assert tools.count("你好 world") == estimate_token_count("你好 world") + assert tools.count("a") == 1 + assert tools.count("") == 0 + assert tools.truncate("abcdefghij", 1) == ("abcde", 1, True) + assert tools.truncate("你好", 1) == ("你", 1, True) + assert tools.truncate("abc", 1) == ("abc", 1, False) + assert tools.truncate("", 1) == ("", 0, False) + assert tools.tokenizer is None + assert len(caught) == 1 + assert "pip install 'reme-ai[qwen-tokenizer]'" in str(caught[0].message) + if failure == "remote": + load.assert_called_once_with(source) + + +@pytest.mark.parametrize("failure", [None, RuntimeError]) +def test_concurrent_initialization_loads_and_warns_once(monkeypatch, failure): + """Concurrent property and count access share one initialization attempt.""" + load_attempts = [] + tokenizer = _Tokenizer() + ready = Barrier(8) + + def load(source): + load_attempts.append(source) + sleep(0.01) # Allow competing callers to reach the initialization lock. + if failure is not None: + raise failure("unavailable") + return tokenizer + + def access(index): + ready.wait(timeout=5) + return tools.tokenizer if index % 2 else tools.count("hello") + + monkeypatch.setattr(qwen_token_tools, "_load_tokenizer", load) + tools = QwenTokenTools() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(access, range(8))) + assert tools.tokenizer is (None if failure else tokenizer) + assert load_attempts == [tools.tokenizer_source] + assert results[::2] == [1 if failure else 5] * 4 + assert all(result is (None if failure else tokenizer) for result in results[1::2]) + assert len(caught) == (1 if failure else 0) + + +@pytest.fixture(name="byte_tokenizer") +def _byte_tokenizer(): + """Build a byte-level tokenizer locally, including partial Unicode tokens.""" + tokenizers = pytest.importorskip("tokenizers") + alphabet = tokenizers.pre_tokenizers.ByteLevel.alphabet() + tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE({char: index for index, char in enumerate(alphabet)}, [])) + tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) + tokenizer.decoder = tokenizers.decoders.ByteLevel() + return tokenizer + + +@pytest.mark.parametrize("text", ["🦄abc", "𠮷abc", "e\u0301abc"]) +def test_unicode_truncation_preserves_original_prefix(byte_tokenizer, text): + """Never decode incomplete UTF-8 tokens or change the source spelling.""" + tools = QwenTokenTools(tokenizer=byte_tokenizer) + for budget in range(1, tools.count(text) + 1): + truncated, count, changed = tools.truncate(text, budget) + assert text.startswith(truncated) + assert "�" not in truncated + assert count == tools.count(truncated) + assert count <= budget + assert changed == (truncated != text) + if text[0] != "e": + assert tools.truncate(text, 1) == ("", 0, True) + assert tools.truncate(text, 4) == (text[0], 4, True) + + +def test_normalized_truncation_preserves_original_prefix(): + """NFC output must never replace the original decomposed spelling.""" + tokenizers = pytest.importorskip("tokenizers") + tokenizer = tokenizers.Tokenizer( + tokenizers.models.WordPiece({"[UNK]": 0, "é": 1, "##a": 2, "##b": 3, "##c": 4}), + ) + tokenizer.normalizer = tokenizers.normalizers.NFC() + tools = QwenTokenTools(tokenizer=tokenizer) + text = "e\u0301abc" + assert tools.count(text) == 4 + truncated, count, changed = tools.truncate(text, 1) + assert truncated + assert text.startswith(truncated) + assert count == tools.count(truncated) == 1 + assert changed is True + + +def test_special_tokens_are_counted_without_automatic_wrapping(byte_tokenizer): + """Explicit special tokens count, while the post-processor adds nothing.""" + tokenizers = pytest.importorskip("tokenizers") + byte_tokenizer.add_special_tokens([tokenizers.AddedToken("[SPECIAL]", lstrip=True, rstrip=True)]) + byte_tokenizer.post_processor = tokenizers.processors.TemplateProcessing( + single="[SPECIAL] $A [SPECIAL]", + special_tokens=[("[SPECIAL]", 256)], + ) + tools = QwenTokenTools(tokenizer=byte_tokenizer) + assert tools.count("") == 0 + assert tools.count("[SPECIAL]abc") == 4 + assert tools.truncate(" [SPECIAL] abc", 1) == (" [SPECIAL] ", 1, True) + + +@pytest.mark.parametrize("source", ["tokenizer.json", "."]) +def test_local_sources_follow_file_updates_and_current_directory(word_tokenizer, tmp_path, monkeypatch, source): + """New instances read updated files and resolve relative paths in the current directory.""" + tokenizers = pytest.importorskip("tokenizers") + monkeypatch.chdir(tmp_path) + word_tokenizer.save("tokenizer.json") + first = QwenTokenTools(source) + assert first.count("one two") == 2 + + word_tokenizer.normalizer = tokenizers.normalizers.Replace(" ", "") + word_tokenizer.save("tokenizer.json") + assert QwenTokenTools(source).count("one two") == 1 + assert first.count("one two") == 2 + + other = tmp_path / "other" + other.mkdir() + monkeypatch.chdir(other) + word_tokenizer.normalizer = None + word_tokenizer.save("tokenizer.json") + assert QwenTokenTools(source).count("one two") == 2 + + +def test_invalid_inputs_do_not_load_tokenizer(monkeypatch): + """Handle empty text and reject invalid inputs before loading or downloading.""" + load = Mock(side_effect=AssertionError("unexpected tokenizer load")) + monkeypatch.setattr(qwen_token_tools, "_load_tokenizer", load) + tools = QwenTokenTools() + assert tools.count("") == 0 + assert tools.truncate("", 1) == ("", 0, False) + for text in (None, b"hello"): + with pytest.raises(TypeError, match="text must be a string"): + tools.count(text) + with pytest.raises(TypeError, match="text must be a string"): + tools.truncate(text, 1) + for budget in (True, 0, -1, 1.5): + with pytest.raises(ValueError, match="max_tokens must be a positive integer"): + tools.truncate("hello", budget) + with pytest.raises(ValueError, match="max_tokens must be a positive integer"): + tools.truncate("", budget) + load.assert_not_called()