Coverage for src\genesis\inference\token.py: 100%

1from __future__ import annotations 

2 

3import unicodedata 

4from typing import Callable 

5 

6# CJK 每字符保守 token 数（T9 整改：cl100k/4字符1token 对中文/日文严重低估） 

7# 日文/中文实际每字符 1~3 token，取保守 1.5 防止上下文裁剪失效导致 API 超限 

8_CJK_TOKENS_PER_CHAR = 1.5 

9 

10 

11def _is_cjk_char(ch: str) -> bool: 

12 """判断字符是否属于 CJK 密集区（中文/日文假名/韩文/全角符号）。""" 

13 cp = ord(ch) 

14 return ( 

15 0x4E00 <= cp <= 0x9FFF # CJK 统一表意文字 

16 or 0x3040 <= cp <= 0x30FF # 平假名/片假名 

17 or 0xAC00 <= cp <= 0xD7AF # 韩文音节 

18 or 0xF900 <= cp <= 0xFAFF # CJK 兼容表意文字 

19 or 0xFF00 <= cp <= 0xFFEF # 全角形式（全角标点/字母） 

20 or 0x3400 <= cp <= 0x4DBF # CJK 扩展 A 

21 ) 

22 

23 

24def approximate_token_count(text: str) -> int: 

25 """内置近似估算（T9 CJK 保守）：CJK 字符按 1.5 token/字符， 

26 其余字符按 4 字符 ≈ 1 token；最少 1 token（无外部依赖，可离线）。""" 

27 cjk_chars = sum(1 for ch in text if _is_cjk_char(ch)) 

28 other_chars = len(text) - cjk_chars 

29 tokens = cjk_chars * _CJK_TOKENS_PER_CHAR + (other_chars + 3) // 4 

30 return max(1, int(tokens)) 

31 

32 

33def _tiktoken_estimator(text: str) -> int | None: 

34 """tiktoken 编码估算；tiktoken 未安装时返回 None。""" 

35 try: 

36 import tiktoken 

37 except ImportError: 

38 return None 

39 try: 

40 enc = tiktoken.get_encoding("cl100k_base") 

41 return len(enc.encode(text)) 

42 except Exception: 

43 return None 

44 

45 

46def make_estimator(backend: str = "tiktoken") -> Callable[[str], int]: 

47 """按配置选择估算器：backend="tiktoken"（默认）优先 tiktoken， 

48 缺失或异常回落内置 approximate；backend="approximate" 直接用近似。""" 

49 if backend == "approximate": 

50 return approximate_token_count 

51 return lambda text: _tiktoken_estimator(text) or approximate_token_count(text)