From fbdfdfee578037604b3d779e96263ef0ef4a9208 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 11:17:12 +0800 Subject: [PATCH] feat(token-counter): add token counting system with multiple implementations --- docs/deprecated.txt | 4 +- reme_ai/core/token_counter/__init__.py | 11 + .../core/token_counter/base_token_counter.py | 59 ++ .../core/token_counter/hf_token_counter.py | 82 +++ .../token_counter/openai_token_counter.py | 58 ++ tests/test_token_counter.py | 511 ++++++++++++++++++ 6 files changed, 724 insertions(+), 1 deletion(-) create mode 100644 reme_ai/core/token_counter/__init__.py create mode 100644 reme_ai/core/token_counter/base_token_counter.py create mode 100644 reme_ai/core/token_counter/hf_token_counter.py create mode 100644 reme_ai/core/token_counter/openai_token_counter.py create mode 100644 tests/test_token_counter.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 6d7cb14f..18618372 100644 --- a/docs/deprecated.txt +++ b/docs/deprecated.txt @@ -1,7 +1,9 @@ from loguru import logger 用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码 -用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范 + +看看代码有什么问题 +用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范，使用list而不是typing.List，不使用typing.Union C0114: Missing module docstring (missing-module-docstring) C0115: Missing class docstring (missing-class-docstring) C0116: Missing function or method docstring (missing-function-docstring) diff --git a/reme_ai/core/token_counter/__init__.py b/reme_ai/core/token_counter/__init__.py new file mode 100644 index 00000000..a9b50826 --- /dev/null +++ b/reme_ai/core/token_counter/__init__.py @@ -0,0 +1,11 @@ +"""token counter""" + +from .base_token_counter import BaseTokenCounter +from .hf_token_counter import HFTokenCounter +from .openai_token_counter import OpenAITokenCounter + +__all__ = [ + "BaseTokenCounter", + "HFTokenCounter", + "OpenAITokenCounter", +] diff --git a/reme_ai/core/token_counter/base_token_counter.py b/reme_ai/core/token_counter/base_token_counter.py new file mode 100644 index 00000000..76ccdb78 --- /dev/null +++ b/reme_ai/core/token_counter/base_token_counter.py @@ -0,0 +1,59 @@ +"""Token counting utility based on character-type rules.""" + +import math +import re +from loguru import logger + +from ..context import C +from ..schema import Message, ToolCall + + +@C.register_token_counter("base") +class BaseTokenCounter: + """A rule-based token counter for Chinese and non-Chinese text.""" + + def __init__(self, model_name: str, **kwargs): + """Initialize with model name and additional parameters.""" + self.model_name = model_name + self.kwargs = kwargs + # Matches Chinese characters including extensions + self._cn_regex = re.compile(r"[\u4e00-\u9fff]") + + def _count_chars(self, text: str) -> tuple[int, int]: + """Count Chinese and other characters in a string.""" + if not text: + return 0, 0 + cn_count = len(self._cn_regex.findall(text)) + return cn_count, len(text) - cn_count + + def count_token( + self, + messages: list[Message], + tools: list[ToolCall] | None = None, + **_kwargs, + ) -> int: + """Calculate total tokens using the 1:2 (CN) and 1:4 (Other) rule.""" + cn_total = 0 + ot_total = 0 + logger.info("Calculating tokens using rule-based estimation.") + + # Extract text from messages + segments = [] + for msg in messages: + content = msg.content + if isinstance(content, bytes): + content = content.decode("utf-8", errors="ignore") + segments.extend([content, msg.reasoning_content]) + + # Extract text from tools + if tools: + for tool in tools: + segments.extend([tool.name, tool.description, tool.arguments]) + + # Process all segments + for text in filter(None, segments): + cn_chars, ot_chars = self._count_chars(text) + cn_total += cn_chars + ot_total += ot_chars + + return math.ceil(cn_total / 2) + math.ceil(ot_total / 4) diff --git a/reme_ai/core/token_counter/hf_token_counter.py b/reme_ai/core/token_counter/hf_token_counter.py new file mode 100644 index 00000000..0bad4c78 --- /dev/null +++ b/reme_ai/core/token_counter/hf_token_counter.py @@ -0,0 +1,82 @@ +"""HuggingFace token counting utilities.""" + +import os + +from loguru import logger + +from .base_token_counter import BaseTokenCounter +from ..context import C +from ..schema import Message, ToolCall + + +@C.register_token_counter("hf") +class HFTokenCounter(BaseTokenCounter): + """Token counter using transformers.AutoTokenizer.apply_chat_template.""" + + def __init__( + self, + model_name: str, + use_fast: bool = False, + trust_remote_code: bool = False, + use_mirror: bool = True, + **kwargs, + ): + """Initialize the counter with model config and lazy tokenizer loading.""" + super().__init__(model_name=model_name, **kwargs) + self.use_fast = use_fast + self.trust_remote_code = trust_remote_code + self.use_mirror = use_mirror + self._tokenizer = None + + def _ensure_tokenizer(self): + """Initialize and cache the HuggingFace tokenizer safely.""" + if self._tokenizer: + return self._tokenizer + + if self.use_mirror: + os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com") + + try: + from transformers import AutoTokenizer + + logger.info("Initializing HuggingFace tokenizer for {}", self.model_name) + + tokenizer = AutoTokenizer.from_pretrained( + self.model_name, + use_fast=self.use_fast, + trust_remote_code=self.trust_remote_code, + **self.kwargs, + ) + + if not hasattr(tokenizer, "chat_template") or tokenizer.chat_template is None: + raise ValueError(f"Model {self.model_name} lacks a chat template.") + + self._tokenizer = tokenizer + return tokenizer + except Exception as e: + logger.error("Failed to load tokenizer {}: {}", self.model_name, e) + raise + + def count_token( + self, + messages: list[Message], + tools: list[ToolCall] | None = None, + **kwargs, + ) -> int: + """Calculate total tokens for messages and tools using the chat template.""" + tokenizer = self._ensure_tokenizer() + + # Serialize inputs for the template + formatted_msgs = [m.simple_dump() for m in messages] + formatted_tools = [t.simple_input_dump() for t in tools] if tools else None + + # Setting tokenize=True and leaving return_tensors=None returns a List[int] + tokens = tokenizer.apply_chat_template( + formatted_msgs, + tools=formatted_tools, + add_generation_prompt=kwargs.pop("add_generation_prompt", False), + tokenize=True, + **kwargs, + ) + + return len(tokens) diff --git a/reme_ai/core/token_counter/openai_token_counter.py b/reme_ai/core/token_counter/openai_token_counter.py new file mode 100644 index 00000000..5793e88a --- /dev/null +++ b/reme_ai/core/token_counter/openai_token_counter.py @@ -0,0 +1,58 @@ +"""Token counting implementation for OpenAI-compatible models.""" + +import json +from loguru import logger +from .base_token_counter import BaseTokenCounter +from ..context import C +from ..schema import Message, ToolCall + + +@C.register_token_counter("openai") +class OpenAITokenCounter(BaseTokenCounter): + """Token counter for OpenAI models using tiktoken.""" + + def __init__(self, model_name: str, **kwargs): + super().__init__(model_name, **kwargs) + self._encoding = None + + @property + def encoding(self): + """Get or initialize the tiktoken encoding for the specified model.""" + if self._encoding is None: + import tiktoken + + try: + self._encoding = tiktoken.encoding_for_model(self.model_name) + except KeyError: + logger.warning(f"Model {self.model_name} not found; falling back to o200k_base.") + self._encoding = tiktoken.get_encoding("o200k_base") + return self._encoding + + def count_token( + self, + messages: list[Message], + tools: list[ToolCall] | None = None, + **_kwargs, + ) -> int: + """Calculate total tokens for a request including messages and tool definitions.""" + enc = self.encoding + total_tokens = 0 + + for msg in messages: + # Every message has <|start|>{role/name}\n{content}<|end|>\n + total_tokens += 3 # Base overhead per message + if msg.content: + total_tokens += len(enc.encode(msg.content)) + + if msg.tool_calls: + for tc in msg.tool_calls: + dump = json.dumps(tc.simple_output_dump(), ensure_ascii=False) + total_tokens += len(enc.encode(dump)) + + if tools: + # Account for tool/function definitions if provided + tool_json = json.dumps([t.simple_input_dump() for t in tools], ensure_ascii=False) + total_tokens += len(enc.encode(tool_json)) + + total_tokens += 3 # Every reply is primed with <|start|>assistant<|message|> + return total_tokens diff --git a/tests/test_token_counter.py b/tests/test_token_counter.py new file mode 100644 index 00000000..3c44a298 --- /dev/null +++ b/tests/test_token_counter.py @@ -0,0 +1,511 @@ +""" +Unit tests for TokenCounter classes covering: +- BaseTokenCounter (rule-based estimation) +- OpenAITokenCounter (tiktoken-based) +- HFTokenCounter (HuggingFace tokenizer-based) + +Usage: + python test_token_counter.py --base # Test BaseTokenCounter only + python test_token_counter.py --openai # Test OpenAITokenCounter only + python test_token_counter.py --hf # Test HFTokenCounter only + python test_token_counter.py --all # Test all token counters +""" + +import argparse +from typing import Type, List + +from reme_ai.core.enumeration import Role +from reme_ai.core.schema import Message, ToolCall +from reme_ai.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter + + +def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter: + """Create and return a token counter instance.""" + default_kwargs = { + "model_name": "gpt-4o", + } + default_kwargs.update(kwargs) + return counter_class(**default_kwargs) + + +def get_test_messages() -> List[Message]: + """Create test messages for token counting.""" + return [ + Message(role=Role.SYSTEM, content="You are a helpful assistant."), + Message(role=Role.USER, content="Hello, how are you today?"), + Message(role=Role.ASSISTANT, content="I'm doing well, thank you for asking! How can I help you?"), + Message(role=Role.USER, content="Can you explain what machine learning is?"), + Message( + role=Role.ASSISTANT, + content="Machine learning is a subset of artificial intelligence that enables computers to " + "learn from data without being explicitly programmed.", + ), + ] + + +def get_chinese_messages() -> List[Message]: + """Create test messages with Chinese content.""" + return [ + Message(role=Role.SYSTEM, content="你是一个有帮助的助手。"), + Message(role=Role.USER, content="你好，今天天气怎么样？"), + Message(role=Role.ASSISTANT, content="今天天气很好，阳光明媚，适合外出活动。"), + Message(role=Role.USER, content="能给我推荐一些好看的电影吗？"), + Message(role=Role.ASSISTANT, content="当然可以！我推荐《肖申克的救赎》、《阿甘正传》和《泰坦尼克号》。"), + ] + + +def get_mixed_messages() -> List[Message]: + """Create test messages with mixed English and Chinese content.""" + return [ + Message(role=Role.SYSTEM, content="You are a bilingual assistant. 你是一个双语助手。"), + Message(role=Role.USER, content="What is AI? 什么是人工智能？"), + Message( + role=Role.ASSISTANT, + content="AI (Artificial Intelligence) 是人工智能的英文缩写，它是计算机科学的一个分支。", + ), + ] + + +def get_messages_with_reasoning() -> List[Message]: + """Create test messages with reasoning content.""" + return [ + Message(role=Role.USER, content="What is 2 + 2?"), + Message( + role=Role.ASSISTANT, + content="The answer is 4.", + reasoning_content="Let me think about this step by step. 2 + 2 " + "equals 4 because addition combines two quantities.", + ), + ] + + +def get_test_tools() -> List[ToolCall]: + """Create test tool calls for token counting.""" + return [ + ToolCall( + **{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a specified location.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and country, e.g., 'Beijing, China'", + }, + "unit": { + "type": "string", + "description": "Temperature unit: 'celsius' or 'fahrenheit'", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + }, + ), + ToolCall( + **{ + "type": "function", + "function": { + "name": "search_web", + "description": "Search the web for information.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query", + }, + "num_results": { + "type": "integer", + "description": "Number of results to return", + }, + }, + "required": ["query"], + }, + }, + }, + ), + ] + + +def get_tool_call_messages() -> List[Message]: + """Create messages with tool call responses.""" + return [ + Message(role=Role.USER, content="What's the weather in Beijing?"), + Message( + role=Role.ASSISTANT, + content="", + tool_calls=[ + ToolCall( + id="call_123", + name="get_weather", + arguments='{"location": "Beijing, China", "unit": "celsius"}', + ), + ], + ), + Message( + role=Role.TOOL, + content='{"temperature": 25, "condition": "sunny", "humidity": 60}', + tool_call_id="call_123", + ), + Message( + role=Role.ASSISTANT, + content="The weather in Beijing is sunny with a temperature of 25°C and 60% humidity.", + ), + ] + + +def test_basic_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test basic token counting with simple messages.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Basic Token Count") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_test_messages() + + print(f"Input: {len(messages)} messages") + for i, msg in enumerate(messages, 1): + content_preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content + print(f" {i}. [{msg.role.value}] {content_preview}") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} basic token count") + + +def test_chinese_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with Chinese content.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Chinese Token Count") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_chinese_messages() + + print(f"Input: {len(messages)} Chinese messages") + for i, msg in enumerate(messages, 1): + content_preview = msg.content[:30] + "..." if len(msg.content) > 30 else msg.content + print(f" {i}. [{msg.role.value}] {content_preview}") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} Chinese token count") + + +def test_mixed_language_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with mixed English and Chinese content.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Mixed Language Token Count") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_mixed_messages() + + print(f"Input: {len(messages)} mixed language messages") + for i, msg in enumerate(messages, 1): + content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content + print(f" {i}. [{msg.role.value}] {content_preview}") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} mixed language token count") + + +def test_reasoning_content_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with reasoning content.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Reasoning Content Token Count") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_messages_with_reasoning() + + print(f"Input: {len(messages)} messages with reasoning content") + for i, msg in enumerate(messages, 1): + print(f" {i}. [{msg.role.value}] content: {msg.content[:30]}...") + if msg.reasoning_content: + print(f" reasoning: {msg.reasoning_content[:30]}...") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} reasoning content token count") + + +def test_token_count_with_tools(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with tool definitions.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Token Count with Tools") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_test_messages()[:2] + tools = get_test_tools() + + print(f"Input: {len(messages)} messages, {len(tools)} tools") + for tool in tools: + print(f" Tool: {tool.name} - {tool.description[:40]}...") + + token_count = counter.count_token(messages, tools=tools) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + # Token count with tools should be higher than without + token_count_no_tools = counter.count_token(messages) + assert token_count > token_count_no_tools, f"{counter_name}: Token count with tools should be higher" + + print(f"\n✓ Token count without tools: {token_count_no_tools}") + print(f"✓ Token count with tools: {token_count}") + print(f"✓ Tools added {token_count - token_count_no_tools} tokens") + print(f"✓ PASSED: {counter_name} token count with tools") + + +def test_tool_call_messages_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with messages containing tool calls.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Tool Call Messages Token Count") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = get_tool_call_messages() + + print(f"Input: {len(messages)} messages with tool calls") + for i, msg in enumerate(messages, 1): + if msg.tool_calls: + print(f" {i}. [{msg.role.value}] tool_calls: {[tc.name for tc in msg.tool_calls]}") + else: + content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content + print(f" {i}. [{msg.role.value}] {content_preview}") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} tool call messages token count") + + +def test_empty_messages(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with empty message list.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Empty Messages") + print(f"{'=' * 60}") + + # HFTokenCounter does not support empty message list (apply_chat_template requires at least one message) + if counter_class == HFTokenCounter: + print("⊘ SKIPPED: HFTokenCounter does not support empty message list") + return + + counter = get_token_counter(counter_class, **kwargs) + messages: List[Message] = [] + + print("Input: 0 messages") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + + # OpenAITokenCounter adds 3 tokens for reply priming even with empty messages + if counter_class == OpenAITokenCounter: + assert token_count == 3, f"{counter_name}: Empty messages should have 3 tokens (reply priming)" + print(f"\n✓ Token count: {token_count} (includes 3 tokens for reply priming)") + else: + assert token_count == 0, f"{counter_name}: Empty messages should have 0 tokens" + print(f"\n✓ Token count: {token_count}") + + print(f"✓ PASSED: {counter_name} empty messages") + + +def test_single_message(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with a single message.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Single Message") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + messages = [Message(role=Role.USER, content="Hello!")] + + print(f"Input: 1 message - '{messages[0].content}'") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + print(f"\n✓ Token count: {token_count}") + print(f"✓ PASSED: {counter_name} single message") + + +def test_long_content(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Test token counting with long content.""" + print(f"\n{'=' * 60}") + print(f"Testing {counter_name}: Long Content") + print(f"{'=' * 60}") + + counter = get_token_counter(counter_class, **kwargs) + + # Create a long message + long_text = "This is a test sentence. " * 100 + messages = [Message(role=Role.USER, content=long_text)] + + print(f"Input: 1 message with {len(long_text)} characters") + + token_count = counter.count_token(messages) + + assert token_count is not None, f"{counter_name}: Token count is None" + assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" + assert token_count > 0, f"{counter_name}: Token count should be positive" + + # Long content should have more tokens + short_messages = [Message(role=Role.USER, content="This is a test sentence.")] + short_token_count = counter.count_token(short_messages) + assert token_count > short_token_count, f"{counter_name}: Long content should have more tokens" + + print(f"\n✓ Short content token count: {short_token_count}") + print(f"✓ Long content token count: {token_count}") + print(f"✓ PASSED: {counter_name} long content") + + +def run_all_tests_for_counter(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): + """Run all tests for a specific token counter class.""" + print(f"\n\n{'#' * 60}") + print(f"# Running all tests for: {counter_name}") + print(f"{'#' * 60}") + + test_basic_token_count(counter_class, counter_name, **kwargs) + test_chinese_token_count(counter_class, counter_name, **kwargs) + test_mixed_language_token_count(counter_class, counter_name, **kwargs) + test_reasoning_content_token_count(counter_class, counter_name, **kwargs) + test_token_count_with_tools(counter_class, counter_name, **kwargs) + test_tool_call_messages_token_count(counter_class, counter_name, **kwargs) + test_empty_messages(counter_class, counter_name, **kwargs) + test_single_message(counter_class, counter_name, **kwargs) + test_long_content(counter_class, counter_name, **kwargs) + + print(f"\n{'=' * 60}") + print(f"✓ All tests passed for {counter_name}!") + print(f"{'=' * 60}") + + +def main(): + """Main entry point for running tests.""" + parser = argparse.ArgumentParser( + description="Run token counter tests", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python test_token_counter.py --base # Test BaseTokenCounter only + python test_token_counter.py --openai # Test OpenAITokenCounter only + python test_token_counter.py --hf # Test HFTokenCounter only + python test_token_counter.py --all # Test all token counters + """, + ) + parser.add_argument( + "--base", + action="store_true", + help="Test BaseTokenCounter (rule-based)", + ) + parser.add_argument( + "--openai", + action="store_true", + help="Test OpenAITokenCounter (tiktoken-based)", + ) + parser.add_argument( + "--hf", + action="store_true", + help="Test HFTokenCounter (HuggingFace tokenizer-based)", + ) + parser.add_argument( + "--hf-model", + type=str, + default="Qwen/Qwen2.5-0.5B-Instruct", + help="HuggingFace model name for HFTokenCounter (default: Qwen/Qwen2.5-0.5B-Instruct)", + ) + parser.add_argument( + "--all", + action="store_true", + help="Run tests for all available token counters", + ) + + args = parser.parse_args() + + # Determine which counters to test + counters_to_test = [] + + if args.all: + counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {})) + counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {})) + counters_to_test.append( + (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), + ) + else: + if args.base: + counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {})) + if args.openai: + counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {})) + if args.hf: + counters_to_test.append( + (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), + ) + + if not counters_to_test: + # Default to all counters if no argument provided + counters_to_test = [ + (BaseTokenCounter, "BaseTokenCounter", {}), + (OpenAITokenCounter, "OpenAITokenCounter", {}), + (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), + ] + print("No counter specified, defaulting to test all counters") + print("Use --base/--openai/--hf to test specific ones\n") + + # Run tests for each counter + for counter_class, counter_name, kwargs in counters_to_test: + try: + run_all_tests_for_counter(counter_class, counter_name, **kwargs) + except Exception as e: + print(f"\n✗ FAILED: {counter_name} tests failed with error:") + print(f" {type(e).__name__}: {e}") + raise + + # Final summary + print(f"\n\n{'#' * 60}") + print("# TEST SUMMARY") + print(f"{'#' * 60}") + print(f"✓ All tests passed for {len(counters_to_test)} token counter(s):") + for _, counter_name, _ in counters_to_test: + print(f" - {counter_name}") + print(f"{'#' * 60}\n") + + +if __name__ == "__main__": + main()