update perf dataset tqdm (#1168)

Author: Yunnglin

evalscope/perf/benchmark.py

@@ -2,7 +2,7 @@
 import json
 import numpy as np
 import sqlite3
-from typing import TYPE_CHECKING, AsyncGenerator, Dict, List, Tuple
+from typing import TYPE_CHECKING, AsyncGenerator, Dict, Tuple
 
 from evalscope.constants import HEARTBEAT_INTERVAL_SEC
 from evalscope.utils.logger import get_logger
@@ -26,48 +26,52 @@
 @exception_handler
 async def get_requests(args: Arguments, api_plugin: 'ApiPluginBase') -> AsyncGenerator[dict, None]:
 
-    async def generate_requests_from_prompt():
+    async def _generate_from_prompt():
+        """Generate requests by repeating a single prompt."""
         prompt = load_prompt(args.prompt)
         messages = [{'role': 'user', 'content': prompt}] if args.apply_chat_template else prompt
         request = api_plugin.build_request(messages)
         for _ in range(args.number):
             yield request
 
-    async def generate_requests_from_dataset():
-        message_generator_class = DatasetRegistry.get_class(args.dataset)
-        message_generator = message_generator_class(args)
-
+    async def _generate_from_dataset():
+        """Generate requests by cycling through a dataset."""
+        message_generator = DatasetRegistry.get_class(args.dataset)(args)
         dataset_messages = []
-        try:
-            for messages in message_generator.build_messages():
+
+        # Load dataset messages into memory (limited by args.number)
+        # We catch StopIteration implicitly via the loop
+        with tqdm(message_generator.build_messages(), desc='Generating datasets', total=args.number, initial=1) as pbar:
+            for messages in pbar:
                 dataset_messages.append(messages)
                 if len(dataset_messages) >= args.number:
                     break
-        except StopIteration:
-            pass
 
         if not dataset_messages:
-            raise Exception('Dataset is empty!')
+            raise ValueError('Dataset is empty!')
 
+        # Yield requests cyclically until total count is reached
         count = 0
         dataset_index = 0
+        num_messages = len(dataset_messages)
 
         while count < args.number:
             messages = dataset_messages[dataset_index]
             request = api_plugin.build_request(messages)
             if request is not None:
                 yield request
                 count += 1
+            dataset_index = (dataset_index + 1) % num_messages
 
-            dataset_index = (dataset_index + 1) % len(dataset_messages)
-
+    # Dispatch based on arguments
     if args.prompt:
-        generator = generate_requests_from_prompt()
+        generator = _generate_from_prompt()
     elif args.dataset:
-        generator = generate_requests_from_dataset()
+        generator = _generate_from_dataset()
     else:
         raise ValueError('Either prompt or dataset is required!')
 
+    # Yield requests with rate limiting
     async for request in generator:
         yield request
         if args.rate != -1:

evalscope/perf/plugin/api/base.py

@@ -1,6 +1,6 @@
 import aiohttp
 from abc import abstractmethod
-from typing import Any, AsyncGenerator, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from evalscope.perf.arguments import Arguments
 from evalscope.perf.utils.benchmark_util import BenchmarkData
@@ -13,7 +13,7 @@ def __init__(self, param: Arguments) -> None:
         self.model_path = param.tokenizer_path
 
     @abstractmethod
-    def build_request(self, messages: List[Dict], param: Arguments = None) -> Dict:
+    def build_request(self, messages: Union[List[Dict], str], param: Optional[Arguments] = None) -> Dict:
         """Build a api request body.
 
         Args:

evalscope/perf/plugin/datasets/base.py

@@ -1,7 +1,6 @@
 import json
-import sys
 from abc import abstractmethod
-from typing import Any, Dict, Iterator, List, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 from evalscope.perf.arguments import Arguments
 
@@ -70,7 +69,7 @@ def dataset_json_list(self, dataset: str) -> Iterator[Dict]:
         for item in data:
             yield item
 
-    def create_message(self, text: str, image_urls: Union[List[str], str] = None, role: str = 'user') -> Dict:
+    def create_message(self, text: str, image_urls: Optional[Union[List[str], str]] = None, role: str = 'user') -> Dict:
         """Create a message with text and optional image URLs.
 
         Args:

evalscope/perf/plugin/datasets/embedding_dataset.py

@@ -6,7 +6,7 @@
 import json
 import numpy as np
 import os
-from typing import Dict, Iterator, List, Union
+from typing import Iterator, List
 
 from evalscope.perf.arguments import Arguments
 from evalscope.perf.plugin.datasets.base import DatasetPluginBase

evalscope/perf/plugin/datasets/random_dataset.py

@@ -18,6 +18,7 @@ def __init__(self, query_parameters: Arguments):
         assert query_parameters.tokenizer_path, 'Tokenizer path is required for random data generation, please provide it with `--tokenizer-path`.'  # noqa: E501
         super().__init__(query_parameters)
 
+        assert self.tokenizer is not None, 'Tokenizer should be initialized for random data generation.'  # noqa: E501
         self.prefix_length = self.query_parameters.prefix_length
         self.number = self.query_parameters.number or 1
         # Use numpy's default_rng for deterministic sampling

evalscope/perf/plugin/datasets/rerank_dataset.py

@@ -156,7 +156,7 @@ def _generate_random_pair(self) -> Dict:
         min_len = max(1, min_len)
         max_len = max(1, max_len)
 
-        query_len = self._rng.integers(min_len, max_len + 1)
+        query_len = int(self._rng.integers(min_len, max_len + 1))
         doc_len = int(query_len * self.document_length_ratio)
         if doc_len < 1:
             doc_len = 1

tests/perf/test_perf.py

@@ -383,7 +383,7 @@ def test_perf_rerank_random(self):
         from evalscope.perf.arguments import Arguments
         task_cfg = Arguments(
             parallel=[1, 2],
-            number=[2, 4],
+            number=[1000, 1000],
             model='qwen3-rerank',
             url='https://dashscope.aliyuncs.com/compatible-api/v1/reranks',
             api_key=env.get('DASHSCOPE_API_KEY'),