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 <jinli.yl@alibaba-inc.com>

.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

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]",

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",

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()

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"]

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 "<EFBFBD>" 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()