evalscope 0.13.2__py3-none-any.whl → 0.15.0__py3-none-any.whl

evalscope/models/local_model.py

@@ -1,7 +1,8 @@
-import torch
+import importlib
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Optional
 
-from evalscope.constants import DEFAULT_MODEL_CACHE_DIR, DEFAULT_MODEL_REVISION, EvalType
+from evalscope.constants import DEFAULT_MODEL_CACHE_DIR, DEFAULT_MODEL_REVISION, EvalType, ModelTask
 from evalscope.utils.logger import get_logger
 from evalscope.utils.model_utils import get_device
 
@@ -11,31 +12,55 @@ if TYPE_CHECKING:
 logger = get_logger()
 
 
-class LocalModel:
+class LocalModel(ABC):
 
     def __init__(self,
                  model_id: str,
-                 model_revision: str = DEFAULT_MODEL_REVISION,
-                 device_map: str = 'auto',
+                 model_revision: str = None,
+                 device_map: str = None,
                  torch_dtype: str = 'auto',
                  cache_dir: str = None,
                  **kwargs):
-        from modelscope import AutoModelForCausalLM, AutoTokenizer
 
-        model_cache_dir = cache_dir or DEFAULT_MODEL_CACHE_DIR
+        self.model_id = model_id
+        self.model_revision = model_revision or DEFAULT_MODEL_REVISION
+        self.device = device_map or get_device()
+        self.cache_dir = cache_dir or DEFAULT_MODEL_CACHE_DIR
+        self.kwargs = kwargs
+        self.model = None
+        self.tokenizer = None
 
         if isinstance(torch_dtype, str) and torch_dtype != 'auto':
+            import torch
             torch_dtype = eval(torch_dtype)
+        self.torch_dtype = torch_dtype
+
+        self.model_cfg = {
+            'model_id': self.model_id,
+            'device_map': self.device,
+            'torch_dtype': str(self.torch_dtype),
+        }
+
+    @abstractmethod
+    def load_model(self):
+        pass
 
-        self.model_id = model_id
-        self.model_revision = model_revision
-        self.device = device_map
+
+class LocalChatModel(LocalModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def load_model(self):
+        from modelscope import AutoModelForCausalLM, AutoTokenizer
+
+        logger.info(f'Loading model {self.model_id} ...')
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_id,
-            revision=model_revision,
+            revision=self.model_revision,
             trust_remote_code=True,
-            cache_dir=model_cache_dir,
+            cache_dir=self.cache_dir,
         )
 
         # Fix no padding
@@ -44,18 +69,45 @@ class LocalModel:
 
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_id,
-            revision=model_revision,
-            device_map=device_map,
+            revision=self.model_revision,
+            device_map=self.device,
             trust_remote_code=True,
-            torch_dtype=torch_dtype,
-            cache_dir=model_cache_dir,
+            torch_dtype=self.torch_dtype,
+            cache_dir=self.cache_dir,
         )
 
-        self.model_cfg = {
-            'model_id': model_id,
-            'device_map': device_map,
-            'torch_dtype': str(torch_dtype),
-        }
+
+class LocalImageModel(LocalModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.pipeline_cls = kwargs.pop('pipeline_cls', None)
+        # default to DiffusionPipeline if not specified
+        if self.pipeline_cls is None:
+            if 'flux' in self.model_id.lower():
+                self.pipeline_cls = 'FluxPipeline'
+            else:
+                self.pipeline_cls = 'DiffusionPipeline'
+
+    def load_model(self):
+        # from modelscope import pipeline_cls
+        module = getattr(importlib.import_module('modelscope'), self.pipeline_cls)
+
+        logger.info(f'Loading model {self.model_id} with {self.pipeline_cls} ...')
+
+        self.model = module.from_pretrained(
+            self.model_id,
+            revision=self.model_revision,
+            torch_dtype=self.torch_dtype,
+            cache_dir=self.cache_dir,
+            **self.kwargs,
+        )
+
+        self.model.to(self.device)
+
+    def __call__(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
 
 
 def get_local_model(task_cfg: 'TaskConfig') -> Optional[LocalModel]:
@@ -64,16 +116,13 @@ def get_local_model(task_cfg: 'TaskConfig') -> Optional[LocalModel]:
     """
     if task_cfg.eval_type != EvalType.CHECKPOINT:
         return None
-    else:
-        device_map = task_cfg.model_args.get('device_map', get_device())
-        cache_dir = task_cfg.model_args.get('cache_dir', None)
-        model_precision = task_cfg.model_args.get('precision', 'torch.float16')
-        model_revision = task_cfg.model_args.get('revision', DEFAULT_MODEL_REVISION)
-
-        base_model = LocalModel(
-            model_id=task_cfg.model,
-            model_revision=model_revision,
-            device_map=device_map,
-            torch_dtype=model_precision,
-            cache_dir=cache_dir)
+    elif task_cfg.model_task == ModelTask.TEXT_GENERATION:
+        base_model = LocalChatModel(model_id=task_cfg.model, **task_cfg.model_args)
+        base_model.load_model()
+        return base_model
+    elif task_cfg.model_task == ModelTask.IMAGE_GENERATION:
+        base_model = LocalImageModel(model_id=task_cfg.model, **task_cfg.model_args)
+        base_model.load_model()
         return base_model
+    else:
+        raise ValueError(f'Unsupported model task: {task_cfg.model_task} for model checkpoint.')

evalscope/models/model.py

@@ -1,9 +1,8 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
-import random
 import time
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, List
 
 from evalscope.utils.logger import get_logger
 
@@ -95,6 +94,7 @@ class ChatBaseModel(BaseModel):
         raise NotImplementedError
 
 
+# TODO: Remove this class after refactoring all models
 class OpenAIModel(ChatBaseModel):
     """
     APIs of OpenAI models.
@@ -187,43 +187,3 @@ class OpenAIModel(ChatBaseModel):
                 time.sleep(3)
         logger.error(f'OpenAI API call failed after {self.MAX_RETRIES} retries')
         return res
-
-
-class DummyChatModel(ChatBaseModel):
-
-    MODEL_ID = 'dummy_chat_model_0801'
-    REVISION = 'v1.0.0'
-
-    def __init__(self, model_cfg: dict, **kwargs):
-        model_cfg['model_id'] = self.MODEL_ID
-        model_cfg['revision'] = self.REVISION
-        super(DummyChatModel, self).__init__(model_cfg=model_cfg)
-
-    def predict(self, inputs: dict, **kwargs) -> dict:
-
-        debug: bool = False
-        if debug:
-            messages = inputs['messages']
-            history = inputs['history']
-
-            logger.info(f'** messages: {messages}')
-            logger.info(f'** history: {history}')
-
-        choice = random.choice(['A', 'B', 'C', 'D'])
-
-        # Build response
-        res = {
-            'choices': [{
-                'index': 0,
-                'message': {
-                    'content': choice,
-                    'role': 'assistant'
-                }
-            }],
-            'created': time.time(),
-            'model': self.MODEL_ID + '-' + self.REVISION,
-            'object': 'chat.completion',
-            'usage': {}
-        }
-
-        return res

evalscope/models/register.py

@@ -1,3 +1,6 @@
+from evalscope.constants import OutputType
+from .adapters import *
+
 MODEL_ADAPTERS = {}
 
 
@@ -26,3 +29,26 @@ def get_model_adapter(name):
         raise ValueError(
             f"Model adapter '{name}' is not registered. Available model adapters: {list(MODEL_ADAPTERS.keys())}")
     return MODEL_ADAPTERS[name]
+
+
+def register_model_adapter_class(cls, name=None):
+    """
+    Register a model adapter class.
+    :param cls: The model adapter class to register
+    :param name: Optional name for the model adapter. If not provided, the class name will be used.
+    """
+    if name is None:
+        name = cls.__name__
+    if name in MODEL_ADAPTERS:
+        raise ValueError(f"Model adapter class '{name}' is already registered.")
+    MODEL_ADAPTERS[name] = cls
+
+
+# register all model adapters
+register_model_adapter_class(BaseModelAdapter, name='base')
+register_model_adapter_class(ChatGenerationModelAdapter, name=OutputType.GENERATION)
+register_model_adapter_class(ContinuationLogitsModelAdapter, name=OutputType.LOGITS)
+register_model_adapter_class(MultiChoiceModelAdapter, name=OutputType.MULTIPLE_CHOICE)
+register_model_adapter_class(CustomModelAdapter, name='custom')
+register_model_adapter_class(ServerModelAdapter, name='server')
+register_model_adapter_class(T2IModelAdapter, name=OutputType.IMAGE_GENERATION)

evalscope/perf/arguments.py

@@ -35,6 +35,7 @@ class Arguments:
     log_every_n_query: int = 10  # Log every N queries
     debug: bool = False  # Debug mode
     wandb_api_key: Optional[str] = None  # WandB API key for logging
+    swanlab_api_key: Optional[str] = None  # SwanLab API key for logging
     name: Optional[str] = None  # Name for the run
 
     # Output settings
@@ -46,6 +47,7 @@ class Arguments:
     prefix_length: int = 0  # Length of the prefix, only for random dataset
     prompt: Optional[str] = None  # The prompt text
     query_template: Optional[str] = None  # Template for the query
+    apply_chat_template: Optional[bool] = None  # Whether to apply chat template
 
     # Dataset settings
     dataset: str = 'openqa'  # Dataset type (default: 'line_by_line')
@@ -57,10 +59,10 @@ class Arguments:
     max_tokens: Optional[int] = 2048  # Maximum number of tokens in the response
     min_tokens: Optional[int] = None  # Minimum number of tokens in the response
     n_choices: Optional[int] = None  # Number of response choices
-    seed: Optional[int] = 42  # Random seed for reproducibility
+    seed: Optional[int] = 0  # Random seed for reproducibility
     stop: Optional[List[str]] = field(default_factory=list)  # Stop sequences for the response
     stop_token_ids: Optional[List[str]] = field(default_factory=list)  # Stop token IDs for the response
-    stream: Optional[bool] = False  # Whether to stream the response
+    stream: Optional[bool] = True  # Whether to stream the response
     temperature: float = 0.0  # Temperature setting for the response
     top_p: Optional[float] = None  # Top-p (nucleus) sampling setting for the response
     top_k: Optional[int] = None  # Top-k sampling setting for the response
@@ -76,12 +78,26 @@ class Arguments:
         return Arguments(**args_dict)
 
     def __post_init__(self):
+        # Set the default headers
         self.headers = self.headers or {}  # Default to empty dictionary
         if self.api_key:
             # Assuming the API key is used as a Bearer token
             self.headers['Authorization'] = f'Bearer {self.api_key}'
+
+        # Set the model ID based on the model name
         self.model_id = os.path.basename(self.model)
 
+        # Set the URL based on the dataset type
+        if self.api.startswith('local'):
+            if self.dataset.startswith('speed_benchmark'):
+                self.url = f'http://127.0.0.1:{self.port}/v1/completions'
+            else:
+                self.url = f'http://127.0.0.1:{self.port}/v1/chat/completions'
+
+        # Set the apply_chat_template flag based on the URL
+        if self.apply_chat_template is None:
+            self.apply_chat_template = self.url.strip('/').endswith('chat/completions')
+
     def __str__(self):
         return json.dumps(self.to_dict(), indent=4, default=str, ensure_ascii=False)
 
@@ -135,7 +151,8 @@ def add_argument(parser: argparse.ArgumentParser):
     parser.add_argument('--log-every-n-query', type=int, default=10, help='Logging every n query')
     parser.add_argument('--debug', action='store_true', default=False, help='Debug request send')
     parser.add_argument('--wandb-api-key', type=str, default=None, help='The wandb API key')
-    parser.add_argument('--name', type=str, help='The wandb db result name and result db name')
+    parser.add_argument('--swanlab-api-key', type=str, default=None, help='The swanlab API key')
+    parser.add_argument('--name', type=str, help='The wandb/swanlab db result name and result db name')
 
     # Prompt settings
     parser.add_argument('--max-prompt-length', type=int, default=sys.maxsize, help='Maximum input prompt length')
@@ -143,6 +160,8 @@ def add_argument(parser: argparse.ArgumentParser):
     parser.add_argument('--prefix-length', type=int, default=0, help='The prefix length')
     parser.add_argument('--prompt', type=str, required=False, default=None, help='Specified the request prompt')
     parser.add_argument('--query-template', type=str, default=None, help='Specify the query template')
+    parser.add_argument(
+        '--apply-chat-template', type=argparse.BooleanOptionalAction, default=None, help='Apply chat template to the prompt')  # noqa: E501
 
     # Output settings
     parser.add_argument('--outputs-dir', help='Outputs dir.', default='outputs')
@@ -159,10 +178,10 @@ def add_argument(parser: argparse.ArgumentParser):
     parser.add_argument(
         '--min-tokens', type=int, help='The minimum number of tokens that can be generated', default=None)
     parser.add_argument('--n-choices', type=int, help='How many completion choices to generate', default=None)
-    parser.add_argument('--seed', type=int, help='The random seed', default=42)
+    parser.add_argument('--seed', type=int, help='The random seed', default=0)
     parser.add_argument('--stop', nargs='*', help='The stop tokens', default=None)
     parser.add_argument('--stop-token-ids', nargs='*', help='Set the stop token IDs', default=None)
-    parser.add_argument('--stream', action='store_true', help='Stream output with SSE', default=False)
+    parser.add_argument('--stream', action=argparse.BooleanOptionalAction, help='Stream output with SSE', default=True)
     parser.add_argument('--temperature', type=float, help='The sample temperature', default=0.0)
     parser.add_argument('--top-p', type=float, help='Sampling top p', default=None)
     parser.add_argument('--top-k', type=int, help='Sampling top k', default=None)

evalscope/perf/benchmark.py

@@ -18,6 +18,7 @@ from evalscope.perf.utils.benchmark_util import BenchmarkData, BenchmarkMetrics
 from evalscope.perf.utils.db_util import create_result_table, get_result_db_path, insert_benchmark_data, summary_result
 from evalscope.perf.utils.handler import add_signal_handlers, exception_handler
 from evalscope.perf.utils.local_server import start_app
+from evalscope.perf.utils.log_utils import init_swanlab, init_wandb
 from evalscope.utils.logger import get_logger
 
 logger = get_logger()
@@ -56,7 +57,7 @@ async def get_requests(args: Arguments) -> AsyncGenerator[dict, None]:
 
     if args.prompt:
         prompt = load_prompt(args.prompt)
-        messages = [{'role': 'user', 'content': prompt}]
+        messages = [{'role': 'user', 'content': prompt}] if args.apply_chat_template else prompt
         generator = generate_requests_from_prompt(messages)
     elif args.dataset:
         generator = generate_requests_from_dataset()
@@ -81,6 +82,7 @@ async def send_request(
         client = AioHttpClient(args)
         async with client:
             benchmark_data = BenchmarkData(request=request)
+            benchmark_data.start_time = time.perf_counter()
             collected_messages = []
             try:
                 async for is_error, state_code, response_data in client.post(request):
@@ -106,24 +108,18 @@ async def send_request(
 
 
 @exception_handler
-async def statistic_benchmark_metric_worker(benchmark_data_queue: asyncio.Queue, args: Arguments):
+async def statistic_benchmark_metric(benchmark_data_queue: asyncio.Queue, args: Arguments):
     metrics = BenchmarkMetrics(concurrency=args.parallel)
 
     api_plugin_class = ApiRegistry(args.api)
     api_plugin = api_plugin_class(args.tokenizer_path)
 
     result_db_path = get_result_db_path(args)
-    # Initialize wandb
-    if args.wandb_api_key:
-        import datetime
-        import wandb
-        os.environ['WANDB_SILENT'] = 'true'
-        os.environ['WANDB_DIR'] = args.outputs_dir
 
-        wandb.login(key=args.wandb_api_key)
-        current_time = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
-        name = args.name if args.name else f'{args.model_id}_{current_time}'
-        wandb.init(project='perf_benchmark', name=name, config=args.to_dict())
+    if args.wandb_api_key:
+        init_wandb(args)
+    if args.swanlab_api_key:
+        init_swanlab(args)
 
     collected_benchmark_data = []
 
@@ -146,9 +142,13 @@ async def statistic_benchmark_metric_worker(benchmark_data_queue: asyncio.Queue,
             # Create a message with the updated metrics
             message = metrics.create_message()
 
-            # Log the message to wandb if the api key is provided
+            # Log the message to wandb\swanlab if the api key is provided
             if args.wandb_api_key:
+                import wandb
                 wandb.log(message)
+            if args.swanlab_api_key:
+                import swanlab
+                swanlab.log(message)
 
             # Log the message to the logger every n queries
             if int(metrics.n_total_queries) % args.log_every_n_query == 0:
@@ -169,17 +169,12 @@ async def statistic_benchmark_metric_worker(benchmark_data_queue: asyncio.Queue,
 
 
 @exception_handler
-async def start_server(args: Arguments) -> bool:
+async def connect_test(args: Arguments) -> bool:
     if args.api.startswith('local'):
         #  start local server
         server = threading.Thread(target=start_app, args=(copy.deepcopy(args), ), daemon=True)
         server.start()
 
-        if args.dataset.startswith('speed_benchmark'):
-            args.url = f'http://127.0.0.1:{args.port}/v1/completions'
-        else:
-            args.url = f'http://127.0.0.1:{args.port}/v1/chat/completions'
-
     if (not args.no_test_connection) and (not await test_connection(args)):
         raise TimeoutError('Test connection failed')
 
@@ -192,31 +187,22 @@ async def benchmark(args: Arguments) -> None:
 
     # init queue
     benchmark_data_queue = asyncio.Queue()
-
     # reset event
     data_process_completed_event.clear()
-
+    # test connection
+    await connect_test(args)
+    # start statistic benchmark metric
+    statistic_benchmark_metric_task = asyncio.create_task(statistic_benchmark_metric(benchmark_data_queue, args))
+    # start send request
     semaphore = asyncio.Semaphore(args.parallel)
+    send_request_tasks: List[asyncio.Task] = []
+    async for request in get_requests(args):
+        task = asyncio.create_task(send_request(semaphore, request, benchmark_data_queue, args))
+        send_request_tasks.append(task)
 
-    async def create_send_request_tasks():
-        tasks: List[asyncio.Task] = []
-        async for request in get_requests(args):
-            task = asyncio.create_task(send_request(semaphore, request, benchmark_data_queue, args))
-            tasks.append(task)
-        return tasks
-
-    async def run_tasks():
-        await start_server(args)
-
-        statistic_benchmark_metric_task = asyncio.create_task(
-            statistic_benchmark_metric_worker(benchmark_data_queue, args))
-        send_request_tasks = await create_send_request_tasks()
-
-        await asyncio.gather(*send_request_tasks, return_exceptions=True)
-        await benchmark_data_queue.join()
-        data_process_completed_event.set()
-
-        metrics, result_db_path = await statistic_benchmark_metric_task
-        summary_result(args, metrics, result_db_path)
+    await asyncio.gather(*send_request_tasks, return_exceptions=True)
+    await benchmark_data_queue.join()
+    data_process_completed_event.set()
 
-    await run_tasks()
+    metrics, result_db_path = await statistic_benchmark_metric_task
+    summary_result(args, metrics, result_db_path)

evalscope/perf/http_client.py

@@ -24,7 +24,6 @@ class AioHttpClient:
         self.connect_timeout = args.connect_timeout
         self.client = aiohttp.ClientSession(
             timeout=aiohttp.ClientTimeout(connect=self.connect_timeout, sock_read=self.read_timeout),
-            connector=aiohttp.TCPConnector(limit=1),
             trace_configs=[self._create_trace_config()] if args.debug else [])
 
     def _create_trace_config(self):
@@ -144,7 +143,7 @@ async def test_connection(args: Arguments) -> bool:
     async def attempt_connection():
         client = AioHttpClient(args)
         async with client:
-            if 'chat/completions' in args.url:
+            if args.apply_chat_template:
                 request = {
                     'messages': [{
                         'role': 'user',
@@ -164,7 +163,7 @@ async def test_connection(args: Arguments) -> bool:
             is_error, state_code, response_data = await asyncio.wait_for(
                 attempt_connection(), timeout=args.connect_timeout)
             if not is_error:
-                logger.info('Connection successful.')
+                logger.info('Test connection successful.')
                 return True
             logger.warning(f'Retrying...  <{state_code}> {response_data}')
         except Exception as e:

evalscope/perf/plugin/api/custom_api.py

@@ -24,7 +24,7 @@ class CustomPlugin(ApiPluginBase):
         """
         super().__init__(model_path=mode_path)
         if mode_path is not None:
-            from transformers import AutoTokenizer
+            from modelscope import AutoTokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(mode_path)
         else:
             self.tokenizer = None

evalscope/perf/plugin/api/openai_api.py

@@ -24,7 +24,7 @@ class OpenaiPlugin(ApiPluginBase):
         """
         super().__init__(model_path=mode_path)
         if mode_path is not None:
-            from transformers import AutoTokenizer
+            from modelscope import AutoTokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(mode_path)
         else:
             self.tokenizer = None
@@ -70,7 +70,7 @@ class OpenaiPlugin(ApiPluginBase):
     def __compose_query_from_parameter(self, payload: Dict, param: Arguments):
         payload['model'] = param.model
         if param.max_tokens is not None:
-            payload['max_completion_tokens'] = param.max_tokens
+            payload['max_tokens'] = param.max_tokens
         if param.min_tokens is not None:
             payload['min_tokens'] = param.min_tokens
         if param.frequency_penalty is not None:

evalscope/perf/plugin/datasets/custom.py

@@ -18,4 +18,7 @@ class CustomDatasetPlugin(DatasetPluginBase):
             prompt = item.strip()
             if len(prompt) > self.query_parameters.min_prompt_length and len(
                     prompt) < self.query_parameters.max_prompt_length:
-                yield [{'role': 'user', 'content': prompt}]
+                if self.query_parameters.apply_chat_template:
+                    yield [{'role': 'user', 'content': prompt}]
+                else:
+                    yield prompt

evalscope/perf/plugin/datasets/flickr8k.py

@@ -30,6 +30,7 @@ class FlickrDatasetPlugin(DatasetPluginBase):
 
         for item in dataset:
             pil_image = item['jpg']
+            text = item['txt']
             base64_iamge = PIL_to_base64(pil_image)
 
             yield [{
@@ -38,7 +39,7 @@ class FlickrDatasetPlugin(DatasetPluginBase):
                 'content': [
                     {
                         'type': 'text',
-                        'text': 'Describe the image'
+                        'text': text,
                     },
                     {
                         'type': 'image_url',

evalscope/perf/plugin/datasets/line_by_line.py

@@ -19,4 +19,7 @@ class LineByLineDatasetPlugin(DatasetPluginBase):
             prompt = item.strip()
             if len(prompt) > self.query_parameters.min_prompt_length and len(
                     prompt) < self.query_parameters.max_prompt_length:
-                yield [{'role': 'user', 'content': prompt}]
+                if self.query_parameters.apply_chat_template:
+                    yield [{'role': 'user', 'content': prompt}]
+                else:
+                    yield prompt

evalscope/perf/plugin/datasets/longalpaca.py

@@ -24,4 +24,7 @@ class LongAlpacaDatasetPlugin(DatasetPluginBase):
             prompt = item['instruction'].strip()
             if len(prompt) > self.query_parameters.min_prompt_length and len(
                     prompt) < self.query_parameters.max_prompt_length:
-                yield [{'role': 'user', 'content': prompt}]
+                if self.query_parameters.apply_chat_template:
+                    yield [{'role': 'user', 'content': prompt}]
+                else:
+                    yield prompt

evalscope/perf/plugin/datasets/openqa.py

@@ -29,4 +29,7 @@ class OpenqaDatasetPlugin(DatasetPluginBase):
             prompt = item['question'].strip()
             if (len(prompt) > self.query_parameters.min_prompt_length
                     and len(prompt) < self.query_parameters.max_prompt_length):
-                yield [{'role': 'user', 'content': prompt}]
+                if self.query_parameters.apply_chat_template:
+                    yield [{'role': 'user', 'content': prompt}]
+                else:
+                    yield prompt

evalscope/perf/plugin/datasets/random_dataset.py

@@ -23,8 +23,12 @@ class RandomDatasetPlugin(DatasetPluginBase):
         self.number = self.query_parameters.number or 1
 
     def build_messages(self) -> Iterator[List[Dict]]:
-        min_prompt_length = self.query_parameters.min_prompt_length - self.template_len
-        max_prompt_length = self.query_parameters.max_prompt_length - self.template_len + 1
+        if self.query_parameters.apply_chat_template:
+            min_prompt_length = self.query_parameters.min_prompt_length - self.template_len
+            max_prompt_length = self.query_parameters.max_prompt_length - self.template_len + 1
+        else:
+            min_prompt_length = self.query_parameters.min_prompt_length
+            max_prompt_length = self.query_parameters.max_prompt_length + 1
 
         assert min_prompt_length >= 0, f'min_prompt_length should be greater than or equal to the template length {self.template_len}.'  # noqa: E501
         assert max_prompt_length >= min_prompt_length, 'max_prompt_length should be greater than or equal to min_prompt_length.'  # noqa: E501
@@ -34,10 +38,13 @@ class RandomDatasetPlugin(DatasetPluginBase):
         offsets = np.random.randint(0, self.tokenizer.vocab_size, size=self.number)
 
         for i in range(self.number):
-            prompt_ids = (offsets[i] + i + np.arange(input_lens[i])) % self.tokenizer.vocab_size
-            prompt = self.tokenizer.decode(
-                self.prefix_ids + prompt_ids.tolist(), skip_special_tokens=False, clean_up_tokenization_spaces=False)
-            yield [{'role': 'user', 'content': prompt}]
+            prompt_ids = ((offsets[i] + i + np.arange(input_lens[i])) % self.tokenizer.vocab_size).tolist()
+            prompt = self.tokenizer.decode(self.prefix_ids + prompt_ids)
+
+            if self.query_parameters.apply_chat_template:
+                yield [{'role': 'user', 'content': prompt}]
+            else:
+                yield prompt
 
     def get_random_inputs(self, length: int) -> List[int]:
         if length <= 0: