wehub-resource-sync
268 lines
9.7 KiB
Python
268 lines
9.7 KiB
Python
"""
|
|
EvalScope integration utilities for ms-swift models.
|
|
|
|
This module provides a custom ModelAPI implementation that enables batch inference
|
|
for evaluation tasks using ms-swift's TransformersEngine. It implements an asynchronous
|
|
batch processing system to improve throughput when evaluating models.
|
|
"""
|
|
|
|
from concurrent.futures import Future
|
|
from dataclasses import dataclass
|
|
from evalscope.api.messages import ChatMessage as EvalChatMessage
|
|
from evalscope.api.model import GenerateConfig, ModelAPI, ModelOutput, ModelUsage
|
|
from evalscope.api.registry import register_model_api
|
|
from evalscope.api.tool import ToolChoice, ToolInfo
|
|
from evalscope.models.utils.openai import chat_choices_from_openai
|
|
from queue import Empty, Queue
|
|
from threading import Thread
|
|
from typing import Any, List, Optional, Tuple
|
|
|
|
from swift.infer_engine import InferRequest, RequestConfig, TransformersEngine
|
|
|
|
|
|
@dataclass
|
|
class BatchInferInput:
|
|
"""
|
|
Container for batch inference input data.
|
|
|
|
Holds all necessary information for a single inference request
|
|
that will be processed as part of a batch.
|
|
"""
|
|
ms_input: InferRequest # ms-swift format request
|
|
ms_config: RequestConfig # ms-swift format configuration
|
|
batch_size: int # desired batch size for this request
|
|
engine: TransformersEngine # inference engine to use
|
|
|
|
|
|
@dataclass
|
|
class _QueueItem:
|
|
"""
|
|
Internal queue item for batch processing.
|
|
|
|
Pairs a batch input with its corresponding future for result delivery.
|
|
"""
|
|
input: BatchInferInput
|
|
future: Future[ModelOutput] # will be resolved with the inference result
|
|
|
|
|
|
# Global variables for batch processing
|
|
# These maintain the shared batch processing infrastructure across all model instances
|
|
batch_thread: Optional[Thread] = None # background thread for processing batches
|
|
batch_queue: Queue[_QueueItem] = Queue() # queue of pending inference requests
|
|
|
|
|
|
@register_model_api('swift_custom')
|
|
class EvalModel(ModelAPI):
|
|
"""
|
|
Custom ModelAPI implementation for ms-swift models with batch inference support.
|
|
|
|
This class integrates ms-swift's TransformersEngine with EvalScope's evaluation framework,
|
|
providing efficient batch processing for improved evaluation throughput.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
base_url: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
config: GenerateConfig = GenerateConfig(),
|
|
**model_args: Any,
|
|
):
|
|
"""
|
|
Initialize the EvalModel with ms-swift backend.
|
|
|
|
Args:
|
|
model_name: Name of the model for identification
|
|
base_url: Not used in this implementation (for API compatibility)
|
|
api_key: Not used in this implementation (for API compatibility)
|
|
config: Generation configuration with batch settings
|
|
**model_args: Additional arguments including 'model' and 'template'
|
|
"""
|
|
super().__init__(
|
|
model_name=model_name,
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
config=config,
|
|
)
|
|
|
|
# Extract model-specific arguments from kwargs
|
|
# This pattern allows us to collect known arguments while preserving unknown ones
|
|
def collect_model_arg(name: str) -> Optional[Any]:
|
|
value = model_args.get(name, None)
|
|
if value is not None:
|
|
model_args.pop(name)
|
|
return value
|
|
|
|
# Extract required model parameters
|
|
self.model = collect_model_arg('model') # model path or identifier
|
|
self.template = collect_model_arg('template') # conversation template
|
|
self.max_batch_size = collect_model_arg('max_batch_size') # maximum batch size
|
|
|
|
# Initialize the inference engine with batch support
|
|
self.engine = TransformersEngine(self.model, template=self.template, max_batch_size=self.max_batch_size)
|
|
|
|
def generate(
|
|
self,
|
|
input: List[EvalChatMessage],
|
|
tools: List[ToolInfo],
|
|
tool_choice: ToolChoice,
|
|
config: GenerateConfig,
|
|
) -> ModelOutput:
|
|
"""
|
|
Generate model response using batch inference.
|
|
|
|
This method queues the request for batch processing and waits for the result.
|
|
The actual inference is performed asynchronously in a background thread.
|
|
|
|
Args:
|
|
input: List of chat messages forming the conversation
|
|
tools: Available tools for function calling (if supported)
|
|
tool_choice: Tool selection strategy
|
|
config: Generation configuration
|
|
|
|
Returns:
|
|
ModelOutput containing the generated response
|
|
"""
|
|
# Ensure the background batch processing thread is running
|
|
global batch_thread
|
|
if batch_thread is None:
|
|
batch_thread = Thread(target=_process_batches, daemon=True)
|
|
batch_thread.start()
|
|
|
|
# Convert EvalScope format to ms-swift format
|
|
ms_input = convert_request(input, tools)
|
|
ms_config = convert_config(config)
|
|
|
|
# Package the request for batch processing
|
|
batch_input = BatchInferInput(
|
|
ms_input=ms_input, ms_config=ms_config, batch_size=config.batch_size, engine=self.engine)
|
|
|
|
# Create a future to receive the result asynchronously
|
|
future = Future[ModelOutput]()
|
|
|
|
# Queue the request for batch processing
|
|
batch_queue.put(_QueueItem(input=batch_input, future=future))
|
|
|
|
# Block until the result is available
|
|
return future.result()
|
|
|
|
|
|
def _process_batches() -> None:
|
|
"""
|
|
Background thread function that processes batched inference requests.
|
|
|
|
This function runs continuously, collecting requests from the queue and
|
|
processing them in batches for improved efficiency. It uses a timeout-based
|
|
approach to balance between batch size and latency.
|
|
"""
|
|
while True:
|
|
# Collect requests from the queue until timeout or batch size limit
|
|
inputs: List[Tuple[BatchInferInput, Future[ModelOutput]]] = []
|
|
|
|
while True:
|
|
try:
|
|
# Wait for new requests with a 2-second timeout
|
|
item = batch_queue.get(timeout=2)
|
|
inputs.append((item.input, item.future))
|
|
|
|
# Check if we've reached the desired batch size
|
|
if len(inputs) == item.input.batch_size:
|
|
break # Process this batch now
|
|
|
|
except Empty:
|
|
# No more requests in queue, process what we have
|
|
break
|
|
|
|
# Skip processing if no requests were collected
|
|
if len(inputs) == 0:
|
|
continue
|
|
|
|
try:
|
|
# Prepare batch inputs for ms-swift inference
|
|
ms_inputs = [item[0].ms_input for item in inputs]
|
|
ms_config = inputs[0][0].ms_config # use first config for the batch
|
|
engine = inputs[0][0].engine # use first engine for the batch
|
|
|
|
# Perform batch inference using ms-swift engine
|
|
completions = engine.infer(ms_inputs, ms_config, use_tqdm=False)
|
|
|
|
# Process results and deliver them to waiting futures
|
|
for i, (batch_input, future) in enumerate(inputs):
|
|
completion = completions[i]
|
|
|
|
# Convert ms-swift response to EvalScope format
|
|
choices = chat_choices_from_openai(completion, tools=[])
|
|
result = ModelOutput(
|
|
model=completion.model,
|
|
choices=choices,
|
|
usage=(ModelUsage(
|
|
input_tokens=completion.usage.prompt_tokens,
|
|
output_tokens=completion.usage.completion_tokens,
|
|
total_tokens=completion.usage.total_tokens,
|
|
) if completion.usage else None),
|
|
)
|
|
|
|
# Deliver the result to the waiting caller
|
|
future.set_result(result)
|
|
|
|
except Exception as ex:
|
|
# If batch processing fails, propagate the error to all waiting futures
|
|
for _, future in inputs:
|
|
future.set_exception(ex)
|
|
|
|
|
|
def convert_config(config: GenerateConfig) -> RequestConfig:
|
|
"""
|
|
Convert EvalScope GenerateConfig to ms-swift RequestConfig.
|
|
|
|
Maps configuration parameters between the two frameworks, ensuring
|
|
compatibility while maintaining the same generation behavior.
|
|
|
|
Args:
|
|
config: EvalScope generation configuration
|
|
|
|
Returns:
|
|
RequestConfig: ms-swift compatible configuration
|
|
"""
|
|
return RequestConfig(
|
|
max_tokens=config.max_tokens,
|
|
temperature=config.temperature,
|
|
top_k=config.top_k,
|
|
top_p=config.top_p,
|
|
presence_penalty=config.presence_penalty,
|
|
frequency_penalty=config.frequency_penalty,
|
|
seed=config.seed,
|
|
stream=False, # batch processing doesn't support streaming
|
|
logprobs=config.logprobs,
|
|
top_logprobs=config.top_logprobs)
|
|
|
|
|
|
def convert_request(messages: List[EvalChatMessage], tools: List[ToolInfo]) -> InferRequest:
|
|
"""
|
|
Convert EvalScope request format to ms-swift InferRequest format.
|
|
|
|
Transforms the message and tool format from EvalScope's representation
|
|
to the format expected by ms-swift's inference engine.
|
|
|
|
Args:
|
|
messages: List of chat messages in EvalScope format
|
|
tools: List of available tools in EvalScope format
|
|
|
|
Returns:
|
|
InferRequest: ms-swift compatible request object
|
|
"""
|
|
# Convert tools to ms-swift format
|
|
tools_list = []
|
|
if len(tools) > 0:
|
|
tools_list = [tool.model_dump(exclude_none=True) for tool in tools]
|
|
|
|
# Convert messages to ms-swift format
|
|
ms_messages = []
|
|
for message in messages:
|
|
ms_messages.append(message.model_dump(exclude_none=True))
|
|
|
|
return InferRequest(
|
|
messages=ms_messages,
|
|
tools=tools_list,
|
|
)
|