wehub-resource-sync
623 lines
21 KiB
Python
623 lines
21 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import base64
|
|
import io
|
|
import json
|
|
import numpy as np
|
|
import os
|
|
import time
|
|
import uuid
|
|
from copy import deepcopy
|
|
from dataclasses import asdict, dataclass, field, fields
|
|
from PIL import Image
|
|
from pydantic import AfterValidator, BaseModel, Field, PlainSerializer, field_validator
|
|
from typing import Annotated, Any, Dict, List, Literal, Optional, Tuple, Union
|
|
|
|
from swift.template import Messages, Tool
|
|
from swift.utils import remove_response
|
|
|
|
|
|
def serialize_ndarray(value):
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, np.ndarray):
|
|
return {
|
|
'data': base64.b64encode(value.tobytes()).decode('ascii'),
|
|
'shape': value.shape,
|
|
'dtype': str(value.dtype),
|
|
'__ndarray__': True
|
|
}
|
|
return value
|
|
|
|
|
|
def deserialize_ndarray(value):
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, dict) and value.get('__ndarray__'):
|
|
data = base64.b64decode(value['data'])
|
|
return np.frombuffer(data, dtype=value['dtype']).reshape(value['shape'])
|
|
return value
|
|
|
|
|
|
NumpyArray = Annotated[Any, PlainSerializer(serialize_ndarray, return_type=Dict), AfterValidator(deserialize_ndarray)]
|
|
|
|
|
|
@dataclass
|
|
class InferRequest:
|
|
"""
|
|
Data structure for inference requests.
|
|
|
|
Attributes:
|
|
messages (Messages):
|
|
The input conversation in messages format. Each message is a dict containing at least
|
|
a "role" field (e.g., "user", "assistant", "system") and a "content" field.
|
|
Example:
|
|
[{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image", # can also be audio/video
|
|
"image": "<url/path/base64/PIL.Image>",
|
|
},
|
|
{"type": "text", "text": "Please describe the picture."},
|
|
],
|
|
}]
|
|
The above is equivalent to:
|
|
[{"role": "user", "content": "<image>Please describe the picture."}]
|
|
with an additional argument:
|
|
images = ["<url/path/base64/PIL.Image>"]
|
|
|
|
images (List[Union[str, Image.Image]]):
|
|
Optional, a list of images associated with the request.
|
|
Each image can be a URL, local path, base64 string, or PIL.Image object.
|
|
|
|
audios (List[str]):
|
|
Optional, a list of audio resources associated with the request.
|
|
|
|
videos (List[str]):
|
|
Optional, a list of video resources associated with the request.
|
|
|
|
tools (Optional[List[Tool]]):
|
|
An optional list of tools. These should be organized in the agent_template format for
|
|
tools requested by the system, for example 'react_en'.
|
|
|
|
objects (Dict[str, Any]):
|
|
Container for additional multimodal objects, grouped by type (key).
|
|
"""
|
|
messages: Messages
|
|
|
|
images: List[Union[str, Image.Image]] = field(default_factory=list)
|
|
audios: List[str] = field(default_factory=list)
|
|
videos: List[str] = field(default_factory=list)
|
|
|
|
tools: Optional[List[Tool]] = None
|
|
objects: Dict[str, Any] = field(default_factory=dict)
|
|
chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
for key in ['images', 'audios', 'videos']:
|
|
val = getattr(self, key)
|
|
if isinstance(val, str):
|
|
setattr(self, key, [val])
|
|
assert isinstance(self.messages, list), f'messages: {self.messages}'
|
|
|
|
@staticmethod
|
|
def remove_response(messages) -> Optional[str]:
|
|
return remove_response(messages)
|
|
|
|
@staticmethod
|
|
def _to_printable(obj, key: Optional[str] = None):
|
|
if isinstance(obj, str) and key not in {'content', 'text'} and len(obj) >= 1000:
|
|
return f'<<<base64:{obj[:50]}..>>>'
|
|
elif isinstance(obj, list):
|
|
res = []
|
|
for item in obj:
|
|
res.append(InferRequest._to_printable(item))
|
|
return res
|
|
elif isinstance(obj, dict):
|
|
res = {}
|
|
for k, v in obj.items():
|
|
res[k] = InferRequest._to_printable(v, key=k)
|
|
return res
|
|
return obj
|
|
|
|
def to_printable(self):
|
|
return InferRequest._to_printable(asdict(self))
|
|
|
|
|
|
@dataclass
|
|
class RolloutInferRequest(InferRequest):
|
|
"""
|
|
An inference request class for rollout scenarios.
|
|
|
|
This class extends `InferRequest` and specifically overrides the `images` attribute
|
|
to be a list of strings for compatibility with POST requests. Each string may
|
|
represent an image URL or a Base64-encoded image.
|
|
|
|
Inherits all fields from `InferRequest`:
|
|
messages (Messages):
|
|
Input conversation messages, supporting multimodal content.
|
|
audios (List[str]):
|
|
List of audio resources associated with the request.
|
|
videos (List[str]):
|
|
List of video resources associated with the request.
|
|
tools (Optional[List[Tool]]):
|
|
List of tools, organized by the agent template (e.g. 'react_en').
|
|
objects (Dict[str, Any]):
|
|
Optional container for additional multimodal objects.
|
|
|
|
Additional / Overridden fields:
|
|
images (List[str]):
|
|
List of image resources, each as a string (URL or base64).
|
|
data_dict (Dict):
|
|
Optional dictionary for extra request data.
|
|
uuid (Optional[str]):
|
|
Optional unique identifier for this request instance.
|
|
"""
|
|
images: List[str] = field(default_factory=list)
|
|
data_dict: Dict = field(default_factory=dict)
|
|
uuid: Optional[str] = None
|
|
|
|
|
|
def random_uuid() -> str:
|
|
return str(uuid.uuid4().hex)
|
|
|
|
|
|
@dataclass
|
|
class Model:
|
|
id: str # model_type
|
|
|
|
object: str = 'model'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
owned_by: str = 'ms-swift'
|
|
|
|
|
|
@dataclass
|
|
class ModelList:
|
|
data: List[Model]
|
|
object: str = 'list'
|
|
|
|
|
|
@dataclass
|
|
class RequestConfig:
|
|
"""NOTE: The following behavior is inconsistent with the OpenAI API.
|
|
Default values for OpenAI:
|
|
temperature = 1.
|
|
top_k = -1
|
|
top_p = 1.
|
|
repetition_penalty = 1.
|
|
"""
|
|
max_tokens: Optional[int] = None # None: max_model_len - num_tokens
|
|
# None: use deploy_args
|
|
temperature: Optional[float] = None
|
|
top_k: Optional[int] = None
|
|
top_p: Optional[float] = None
|
|
repetition_penalty: Optional[float] = None
|
|
num_beams: int = 1
|
|
stop: Optional[List[str]] = field(default_factory=list)
|
|
|
|
seed: Optional[int] = None
|
|
stream: bool = False
|
|
logprobs: bool = False
|
|
top_logprobs: Optional[int] = None
|
|
prompt_logprobs: Optional[int] = None
|
|
|
|
n: int = 1
|
|
best_of: Optional[int] = None
|
|
presence_penalty: float = 0.
|
|
frequency_penalty: float = 0.
|
|
length_penalty: float = 1.
|
|
# Return token_ids additionally (non-stream)
|
|
return_details: bool = False
|
|
# vLLM structured outputs (guided decoding)
|
|
structured_outputs_regex: Optional[str] = None
|
|
|
|
def __post_init__(self):
|
|
if self.stop is None:
|
|
self.stop = []
|
|
|
|
|
|
@dataclass
|
|
class CompletionRequestMixin:
|
|
model: str
|
|
prompt: str
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingRequestMixin:
|
|
input: str
|
|
model: str
|
|
encoding_format: Literal['float', 'base64'] = 'float'
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionRequestMixin:
|
|
model: str
|
|
messages: Messages
|
|
tools: Optional[List[Tool]] = None
|
|
tool_choice: Optional[Union[str, Dict]] = None
|
|
chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
if self.tool_choice is None:
|
|
self.tool_choice = 'none' if self.tools is None else 'auto'
|
|
|
|
if self.tools:
|
|
if self.tool_choice == 'none':
|
|
self.tools = None
|
|
elif isinstance(self.tool_choice, dict):
|
|
name = self.tool_choice['function']['name']
|
|
tool = next(tool for tool in self.tools if tool['function']['name'] == name)
|
|
if tool is None:
|
|
raise ValueError(f"Tool choice '{name}' not found in tools.")
|
|
self.tools = [tool]
|
|
|
|
|
|
@dataclass
|
|
class MultiModalRequestMixin:
|
|
images: List[str] = field(default_factory=list)
|
|
audios: List[str] = field(default_factory=list)
|
|
videos: List[str] = field(default_factory=list)
|
|
objects: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
@staticmethod
|
|
def to_base64(mm_data: Union[str, Image.Image, bytes]) -> str:
|
|
if isinstance(mm_data, dict) and 'bytes' in mm_data:
|
|
mm_data = mm_data['bytes'] or mm_data['path']
|
|
if isinstance(mm_data, str) and not os.path.isfile(mm_data):
|
|
# base64 or url
|
|
return mm_data
|
|
if isinstance(mm_data, str):
|
|
# local_path
|
|
with open(mm_data, 'rb') as f:
|
|
bytes_ = f.read()
|
|
elif isinstance(mm_data, Image.Image):
|
|
bytes_io = io.BytesIO()
|
|
mm_data.save(bytes_io, format='png')
|
|
bytes_ = bytes_io.getvalue()
|
|
else:
|
|
bytes_ = mm_data
|
|
img_base64: str = base64.b64encode(bytes_).decode('utf-8')
|
|
return img_base64
|
|
|
|
def __post_init__(self):
|
|
for key in ['images', 'audios', 'videos']:
|
|
values = getattr(self, key)
|
|
if isinstance(values, str):
|
|
values = [values]
|
|
setattr(self, key, values)
|
|
for i, val in enumerate(values):
|
|
values[i] = self.to_base64(val)
|
|
|
|
|
|
@dataclass
|
|
class CompletionRequest(RequestConfig, MultiModalRequestMixin, CompletionRequestMixin):
|
|
|
|
def __post_init__(self):
|
|
RequestConfig.__post_init__(self)
|
|
MultiModalRequestMixin.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingRequest(RequestConfig, MultiModalRequestMixin, EmbeddingRequestMixin):
|
|
|
|
def __post_init__(self):
|
|
RequestConfig.__post_init__(self)
|
|
MultiModalRequestMixin.__post_init__(self)
|
|
|
|
def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
|
|
data = asdict(self)
|
|
res = []
|
|
for cls_type in [InferRequest, RequestConfig]:
|
|
parameters = set(f.name for f in fields(cls_type))
|
|
_data = {k: v for k, v in data.items() if k in parameters}
|
|
res.append(cls_type(**_data))
|
|
return tuple(res)
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionRequest(RequestConfig, MultiModalRequestMixin, ChatCompletionRequestMixin):
|
|
|
|
def __post_init__(self):
|
|
RequestConfig.__post_init__(self)
|
|
MultiModalRequestMixin.__post_init__(self)
|
|
ChatCompletionRequestMixin.__post_init__(self)
|
|
self.convert_to_base64()
|
|
|
|
def convert_to_base64(self):
|
|
for message in self.messages:
|
|
content = message['content']
|
|
if isinstance(content, str):
|
|
continue
|
|
for item in content:
|
|
key: str = item['type']
|
|
if key == 'text':
|
|
continue
|
|
|
|
key_origin = key
|
|
value = item[key]
|
|
if key.endswith('_url'):
|
|
key = key[:-len('_url')]
|
|
is_dict = False
|
|
if isinstance(value, dict):
|
|
is_dict = True
|
|
value = value['url']
|
|
if isinstance(value, str) and (value.startswith('data:') or value.startswith('http')
|
|
or len(value) > 200):
|
|
continue
|
|
|
|
# local_path / PIL.Image
|
|
if isinstance(value, str) and os.path.isfile(value):
|
|
suffix = os.path.splitext(value)[1][1:].lower()
|
|
elif isinstance(value, Image.Image):
|
|
suffix = 'jpeg'
|
|
else:
|
|
raise ValueError(f'value: {value}')
|
|
mm_data_base64 = self.to_base64(value)
|
|
new_value = f'data:{key}/{suffix};base64,{mm_data_base64}'
|
|
if is_dict:
|
|
new_value = {'url': new_value}
|
|
item[key_origin] = new_value
|
|
|
|
def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
|
|
data = asdict(self)
|
|
res = []
|
|
for cls_type in [InferRequest, RequestConfig]:
|
|
parameters = set(f.name for f in fields(cls_type))
|
|
_data = {k: v for k, v in data.items() if k in parameters}
|
|
res.append(cls_type(**_data))
|
|
return tuple(res)
|
|
|
|
@classmethod
|
|
def from_cmpl_request(cls, cmpl_request: Union[CompletionRequest, EmbeddingRequest]) -> 'ChatCompletionRequest':
|
|
cmpl_request = asdict(cmpl_request)
|
|
if 'prompt' in cmpl_request:
|
|
prompt = cmpl_request.pop('prompt')
|
|
else:
|
|
prompt = cmpl_request.pop('input')
|
|
cmpl_request['messages'] = [{'role': 'user', 'content': prompt}]
|
|
if 'encoding_format' in cmpl_request:
|
|
cmpl_request.pop('encoding_format')
|
|
return cls(**cmpl_request)
|
|
|
|
|
|
@dataclass
|
|
class UsageInfo:
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
|
|
|
|
@dataclass
|
|
class Function:
|
|
name: str
|
|
arguments: Optional[Union[str, Any]]
|
|
|
|
def __post_init__(self):
|
|
if not isinstance(self.arguments, str):
|
|
self.arguments = json.dumps(self.arguments, ensure_ascii=False)
|
|
self.name = self.name.strip()
|
|
self.arguments = self.arguments.strip()
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionMessageToolCall:
|
|
function: Function
|
|
type: str = 'function'
|
|
id: str = field(default_factory=lambda: f'toolcall-{random_uuid()}')
|
|
|
|
|
|
@dataclass
|
|
class ChatMessage:
|
|
role: Literal['system', 'user', 'assistant']
|
|
content: Union[str, List[Dict[str, Any]], int, float, List[float]]
|
|
tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
|
|
reasoning_content: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionResponseChoice:
|
|
index: int
|
|
message: ChatMessage
|
|
finish_reason: Literal['stop', 'length', None]
|
|
logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
|
|
token_ids: Optional[List[int]] = None
|
|
routed_experts: Optional[NumpyArray] = None
|
|
|
|
def to_cmpl_choice(self) -> 'CompletionResponseChoice':
|
|
self = deepcopy(self)
|
|
assert not self.message.tool_calls, f'message: {self.message}'
|
|
return CompletionResponseChoice(self.index, self.message.content, self.finish_reason, self.logprobs)
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingResponseData:
|
|
object: str = 'embedding'
|
|
index: int = 0
|
|
embedding: List[str] = field(default_factory=lambda: [])
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingResponse:
|
|
model: str
|
|
data: List[EmbeddingResponseData]
|
|
usage: UsageInfo
|
|
id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
|
|
object: str = 'list'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
|
|
|
|
@dataclass
|
|
class CompletionResponseChoice:
|
|
index: int
|
|
text: str
|
|
finish_reason: Literal['stop', 'length', None]
|
|
logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionResponse:
|
|
model: str
|
|
choices: List[ChatCompletionResponseChoice]
|
|
usage: UsageInfo
|
|
id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
|
|
object: str = 'chat.completion'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
prompt_token_ids: Optional[List[int]] = None
|
|
prompt_logprobs: Optional[List] = None
|
|
images_size: Optional[List[Tuple[int, int]]] = None
|
|
|
|
def to_cmpl_response(self) -> 'CompletionResponse':
|
|
self = deepcopy(self)
|
|
choices = [choice.to_cmpl_choice() for choice in self.choices]
|
|
id_ = f'cmpl{self.id[len("chatcmpl"):]}'
|
|
return CompletionResponse(
|
|
self.model,
|
|
choices,
|
|
self.usage,
|
|
id_,
|
|
created=self.created,
|
|
prompt_token_ids=self.prompt_token_ids,
|
|
prompt_logprobs=self.prompt_logprobs,
|
|
)
|
|
|
|
|
|
class RolloutOutput(BaseModel):
|
|
"""
|
|
Output structure for rollout.
|
|
|
|
Attributes:
|
|
response (ChatCompletionResponse):
|
|
The model's response
|
|
|
|
messages (Optional[Messages]):
|
|
(Optional) Conversation history for the final rollout; required for multi-turn scenarios.
|
|
NOTE:
|
|
- If provided, this messages sequence will overwrite the original messages.
|
|
- If not provided, 'response' will be appended as the latest turn in the original messages.
|
|
- For multi-turn training, you need to manually return the updated messages, including the full history.
|
|
- The messages should include the latest assistant response as the final message.
|
|
|
|
response_token_ids (Optional[List[List[int]]]):
|
|
(Optional) Token IDs generated at each rollout turn.
|
|
If provided, the training process will skip tokenizing the response.
|
|
|
|
response_loss_mask (Optional[List[List[int]]]):
|
|
(Optional) Loss masks corresponding to each rollout turn.
|
|
If provided, the training process will skip computing loss masks for the response (as controlled by the `loss_scale` parameter). # noqa
|
|
|
|
rollout_infos (Dict[str, Any]):
|
|
(Optional) Additional rollout information. This must be JSON-serializable.
|
|
"""
|
|
response: ChatCompletionResponse
|
|
# multi turn
|
|
messages: Optional[Messages] = None
|
|
response_token_ids: List[List[int]] = Field(default_factory=list)
|
|
response_loss_mask: List[List[int]] = Field(default_factory=list)
|
|
rollout_infos: Dict[str, Any] = Field(default_factory=dict)
|
|
# rollout logprobs for each turn (used for rollout importance sampling correction in multi-turn scenarios)
|
|
rollout_logprobs: List[List[float]] = Field(default_factory=list)
|
|
|
|
prompt_logprobs: Optional[List] = None
|
|
|
|
@field_validator('response_token_ids', 'response_loss_mask', 'rollout_logprobs', mode='before')
|
|
@classmethod
|
|
def _wrap_flat_list(cls, v):
|
|
if isinstance(v, list) and v and isinstance(v[0], (int, float)):
|
|
return [v]
|
|
return v
|
|
|
|
def model_post_init(self, __context):
|
|
# Ensure multimodal data in rollout_infos is serializable (e.g., images to base64)
|
|
super().model_post_init(__context)
|
|
self.mminfo_to_serializable()
|
|
|
|
def mminfo_to_serializable(self):
|
|
mm_keys = ['images', 'audios', 'videos']
|
|
|
|
for key, values in self.rollout_infos.items():
|
|
if key in mm_keys:
|
|
if not isinstance(values, list):
|
|
values = [values]
|
|
for i, value in enumerate(values):
|
|
values[i] = MultiModalRequestMixin.to_base64(value)
|
|
self.rollout_infos[key] = values
|
|
|
|
|
|
@dataclass
|
|
class CompletionResponse:
|
|
model: str
|
|
choices: List[CompletionResponseChoice]
|
|
usage: UsageInfo
|
|
id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
|
|
object: str = 'text_completion'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
prompt_token_ids: Optional[List[int]] = None
|
|
prompt_logprobs: Optional[List] = None
|
|
|
|
|
|
@dataclass
|
|
class DeltaMessage:
|
|
role: Literal['system', 'user', 'assistant', None] = None
|
|
content: Optional[str] = None
|
|
tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
|
|
reasoning_content: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionResponseStreamChoice:
|
|
index: int
|
|
delta: DeltaMessage
|
|
finish_reason: Literal['stop', 'length', None]
|
|
logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
|
|
|
|
def to_cmpl_choice(self) -> 'CompletionResponseStreamChoice':
|
|
self = deepcopy(self)
|
|
assert not self.delta.tool_calls
|
|
return CompletionResponseStreamChoice(self.index, self.delta.content, self.finish_reason, self.logprobs)
|
|
|
|
|
|
@dataclass
|
|
class CompletionResponseStreamChoice:
|
|
index: int
|
|
text: str
|
|
finish_reason: Literal['stop', 'length', None]
|
|
logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionStreamResponse:
|
|
model: str
|
|
choices: List[ChatCompletionResponseStreamChoice]
|
|
usage: Optional[UsageInfo] = None
|
|
id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
|
|
object: str = 'chat.completion.chunk'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
|
|
def to_cmpl_response(self) -> 'CompletionStreamResponse':
|
|
self = deepcopy(self)
|
|
choices = [choice.to_cmpl_choice() for choice in self.choices]
|
|
id_ = f'cmpl{self.id[len("chatcmpl"):]}'
|
|
return CompletionStreamResponse(self.model, choices, self.usage, id_, created=self.created)
|
|
|
|
|
|
@dataclass
|
|
class CompletionStreamResponse:
|
|
model: str
|
|
choices: List[CompletionResponseStreamChoice]
|
|
usage: Optional[UsageInfo] = None
|
|
id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
|
|
object: str = 'text_completion.chunk'
|
|
created: int = field(default_factory=lambda: int(time.time()))
|
|
|
|
|
|
class InitCommunicatorRequest(BaseModel):
|
|
host: str
|
|
port: int
|
|
world_size: int
|
|
|
|
|
|
class UpdateWeightsRequest(BaseModel):
|
|
name: str
|
|
dtype: str
|
|
shape: list[int]
|