[model-gateway] Add classification model support infrastructure (#16061)
Co-authored-by: Chang Su <chang.s.su@oracle.com>
This commit is contained in:
@@ -5,6 +5,7 @@ Uses GrpcRequestManager for orchestration without tokenization.
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
@@ -334,6 +335,9 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
|
||||
pad_token_id=self.model_info["pad_token_id"],
|
||||
bos_token_id=self.model_info["bos_token_id"],
|
||||
max_req_input_len=self.model_info["max_req_input_len"],
|
||||
# Classification model support
|
||||
id2label_json=self.model_info.get("id2label_json") or "",
|
||||
num_labels=self.model_info.get("num_labels") or 0,
|
||||
)
|
||||
|
||||
async def GetServerInfo(
|
||||
@@ -743,6 +747,22 @@ async def serve_grpc(
|
||||
|
||||
# Update model info from scheduler info and model config
|
||||
if model_info is None:
|
||||
# Extract classification labels from HuggingFace config (if available)
|
||||
# Match logic in serving_classify.py::_get_id2label_mapping
|
||||
hf_config = model_config.hf_config
|
||||
id2label = getattr(hf_config, "id2label", None)
|
||||
num_labels = getattr(hf_config, "num_labels", 0) or 0
|
||||
|
||||
# If no id2label but num_labels exists, create default mapping
|
||||
if not id2label and num_labels:
|
||||
id2label = {i: f"LABEL_{i}" for i in range(num_labels)}
|
||||
elif id2label and not num_labels:
|
||||
num_labels = len(id2label)
|
||||
|
||||
# Convert to JSON string for proto transport
|
||||
# id2label is a dict like {0: "negative", 1: "positive"}
|
||||
id2label_json = json.dumps(id2label) if id2label else ""
|
||||
|
||||
model_info = {
|
||||
"model_name": server_args.model_path,
|
||||
"max_context_length": scheduler_info.get(
|
||||
@@ -750,12 +770,15 @@ async def serve_grpc(
|
||||
),
|
||||
"vocab_size": scheduler_info.get("vocab_size", 128256),
|
||||
"supports_vision": scheduler_info.get("supports_vision", False),
|
||||
"model_type": getattr(model_config.hf_config, "model_type", None),
|
||||
"architectures": getattr(model_config.hf_config, "architectures", None),
|
||||
"model_type": getattr(hf_config, "model_type", None),
|
||||
"architectures": getattr(hf_config, "architectures", None),
|
||||
"max_req_input_len": scheduler_info.get("max_req_input_len", 8192),
|
||||
"eos_token_ids": scheduler_info.get("eos_token_ids", []),
|
||||
"pad_token_id": scheduler_info.get("pad_token_id", 0),
|
||||
"bos_token_id": scheduler_info.get("bos_token_id", 1),
|
||||
# Classification model support
|
||||
"id2label_json": id2label_json,
|
||||
"num_labels": num_labels or 0,
|
||||
}
|
||||
|
||||
# Create request manager with the correct port args
|
||||
|
||||
@@ -428,6 +428,12 @@ message GetModelInfoResponse {
|
||||
int32 bos_token_id = 13;
|
||||
int32 max_req_input_len = 14;
|
||||
repeated string architectures = 15;
|
||||
|
||||
// Classification model support (from HuggingFace config.json)
|
||||
// id2label maps class indices to label names, e.g., {"0": "negative", "1": "positive"}
|
||||
string id2label_json = 16;
|
||||
// Number of classification labels (0 if not a classifier)
|
||||
int32 num_labels = 17;
|
||||
}
|
||||
|
||||
// Get server information
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -432,7 +432,7 @@ class GetModelInfoRequest(_message.Message):
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
class GetModelInfoResponse(_message.Message):
|
||||
__slots__ = ("model_path", "tokenizer_path", "is_generation", "preferred_sampling_params", "weight_version", "served_model_name", "max_context_length", "vocab_size", "supports_vision", "model_type", "eos_token_ids", "pad_token_id", "bos_token_id", "max_req_input_len", "architectures")
|
||||
__slots__ = ("model_path", "tokenizer_path", "is_generation", "preferred_sampling_params", "weight_version", "served_model_name", "max_context_length", "vocab_size", "supports_vision", "model_type", "eos_token_ids", "pad_token_id", "bos_token_id", "max_req_input_len", "architectures", "id2label_json", "num_labels")
|
||||
MODEL_PATH_FIELD_NUMBER: _ClassVar[int]
|
||||
TOKENIZER_PATH_FIELD_NUMBER: _ClassVar[int]
|
||||
IS_GENERATION_FIELD_NUMBER: _ClassVar[int]
|
||||
@@ -448,6 +448,8 @@ class GetModelInfoResponse(_message.Message):
|
||||
BOS_TOKEN_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
MAX_REQ_INPUT_LEN_FIELD_NUMBER: _ClassVar[int]
|
||||
ARCHITECTURES_FIELD_NUMBER: _ClassVar[int]
|
||||
ID2LABEL_JSON_FIELD_NUMBER: _ClassVar[int]
|
||||
NUM_LABELS_FIELD_NUMBER: _ClassVar[int]
|
||||
model_path: str
|
||||
tokenizer_path: str
|
||||
is_generation: bool
|
||||
@@ -463,7 +465,9 @@ class GetModelInfoResponse(_message.Message):
|
||||
bos_token_id: int
|
||||
max_req_input_len: int
|
||||
architectures: _containers.RepeatedScalarFieldContainer[str]
|
||||
def __init__(self, model_path: _Optional[str] = ..., tokenizer_path: _Optional[str] = ..., is_generation: bool = ..., preferred_sampling_params: _Optional[str] = ..., weight_version: _Optional[str] = ..., served_model_name: _Optional[str] = ..., max_context_length: _Optional[int] = ..., vocab_size: _Optional[int] = ..., supports_vision: bool = ..., model_type: _Optional[str] = ..., eos_token_ids: _Optional[_Iterable[int]] = ..., pad_token_id: _Optional[int] = ..., bos_token_id: _Optional[int] = ..., max_req_input_len: _Optional[int] = ..., architectures: _Optional[_Iterable[str]] = ...) -> None: ...
|
||||
id2label_json: str
|
||||
num_labels: int
|
||||
def __init__(self, model_path: _Optional[str] = ..., tokenizer_path: _Optional[str] = ..., is_generation: bool = ..., preferred_sampling_params: _Optional[str] = ..., weight_version: _Optional[str] = ..., served_model_name: _Optional[str] = ..., max_context_length: _Optional[int] = ..., vocab_size: _Optional[int] = ..., supports_vision: bool = ..., model_type: _Optional[str] = ..., eos_token_ids: _Optional[_Iterable[int]] = ..., pad_token_id: _Optional[int] = ..., bos_token_id: _Optional[int] = ..., max_req_input_len: _Optional[int] = ..., architectures: _Optional[_Iterable[str]] = ..., id2label_json: _Optional[str] = ..., num_labels: _Optional[int] = ...) -> None: ...
|
||||
|
||||
class GetServerInfoRequest(_message.Message):
|
||||
__slots__ = ()
|
||||
|
||||
Reference in New Issue
Block a user