52 lines
2.2 KiB
Python
52 lines
2.2 KiB
Python
"""Подсчёт токенов для чанкера и плагина `QwenLocal`."""
|
||
|
||
import logging
|
||
from typing import Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class QwenTokenCounter:
|
||
"""Подсчёт токенов токенизатором модели Qwen2.5-3B-Instruct.
|
||
|
||
Загрузка `tokenizers.Tokenizer` — ленивая (при первом вызове), чтобы
|
||
инстанцирование плагина в API-процессе не тянуло за собой файл
|
||
токенизатора. Если файл отсутствует или повреждён — используется
|
||
эвристический фолбэк `len(text) // 3`, чтобы пайплайн не падал из-за
|
||
отсутствия токенизатора (например, в dev-окружении без volume модели).
|
||
"""
|
||
|
||
def __init__(self, tokenizer_path: str | None = None) -> None:
|
||
self._tokenizer_path = tokenizer_path
|
||
self._tokenizer: Any | None = None
|
||
self._load_attempted = False
|
||
|
||
def _ensure_loaded(self) -> None:
|
||
"""Загрузить токенизатор один раз (синглтон на инстанс counter'а)."""
|
||
if self._load_attempted:
|
||
return
|
||
self._load_attempted = True
|
||
|
||
if not self._tokenizer_path:
|
||
logger.warning("Путь к токенизатору не задан, использую эвристику len(text) // 3")
|
||
return
|
||
|
||
try:
|
||
from tokenizers import Tokenizer
|
||
|
||
self._tokenizer = Tokenizer.from_file(self._tokenizer_path)
|
||
except Exception:
|
||
logger.warning(
|
||
"Не удалось загрузить токенизатор из %s, использую эвристику len(text) // 3",
|
||
self._tokenizer_path,
|
||
)
|
||
self._tokenizer = None
|
||
|
||
def __call__(self, text: str) -> int:
|
||
"""Вернуть число токенов в тексте (или эвристическую оценку)."""
|
||
self._ensure_loaded()
|
||
if self._tokenizer is not None:
|
||
encoded: int = len(self._tokenizer.encode(text).ids)
|
||
return encoded
|
||
return len(text) // 3
|