[model-gateway] Add classification model support infrastructure (#16061)

Co-authored-by: Chang Su <chang.s.su@oracle.com>

sgl-model-gateway/src/core/model_card.rs

@@ -7,6 +7,8 @@
 //!
 //! Inspired by Dynamo's ModelDeploymentCard but simplified for router needs.
 
+use std::collections::HashMap;
+
 use serde::{Deserialize, Serialize};
 
 use super::{
@@ -168,6 +170,21 @@ pub struct ModelCard {
     /// User-defined metadata (for fields not covered above)
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub metadata: Option<serde_json::Value>,
+
+    // === Classification Support ===
+    /// Classification label mapping (class index -> label name).
+    /// Empty if not a classification model.
+    /// Example: {0: "negative", 1: "positive"}
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub id2label: HashMap<u32, String>,
+
+    /// Number of classification labels (0 if not a classifier).
+    #[serde(default, skip_serializing_if = "is_zero")]
+    pub num_labels: u32,
+}
+
+fn is_zero(n: &u32) -> bool {
+    *n == 0
 }
 
 fn default_model_type() -> ModelType {
@@ -193,6 +210,8 @@ impl ModelCard {
             reasoning_parser: None,
             tool_parser: None,
             metadata: None,
+            id2label: HashMap::new(),
+            num_labels: 0,
         }
     }
 
@@ -276,6 +295,19 @@ impl ModelCard {
         self
     }
 
+    /// Set the id2label mapping for classification models
+    pub fn with_id2label(mut self, id2label: HashMap<u32, String>) -> Self {
+        self.num_labels = id2label.len() as u32;
+        self.id2label = id2label;
+        self
+    }
+
+    /// Set num_labels directly (alternative to with_id2label)
+    pub fn with_num_labels(mut self, num_labels: u32) -> Self {
+        self.num_labels = num_labels;
+        self
+    }
+
     // === Query methods ===
 
     /// Check if this model matches the given ID (including aliases)
@@ -334,6 +366,20 @@ impl ModelCard {
     pub fn supports_reasoning(&self) -> bool {
         self.model_type.supports_reasoning()
     }
+
+    /// Check if this is a classification model
+    #[inline]
+    pub fn is_classifier(&self) -> bool {
+        self.num_labels > 0
+    }
+
+    /// Get label for a class index, with fallback to generic label (LABEL_N)
+    pub fn get_label(&self, class_idx: u32) -> String {
+        self.id2label
+            .get(&class_idx)
+            .cloned()
+            .unwrap_or_else(|| format!("LABEL_{}", class_idx))
+    }
 }
 
 impl Default for ModelCard {

sgl-model-gateway/src/core/steps/worker/local/create_worker.rs

@@ -175,6 +175,40 @@ fn build_model_card(
         }
     }
 
+    // Parse classification model id2label mapping
+    // The proto field is id2label_json: JSON string like {"0": "negative", "1": "positive"}
+    if let Some(id2label_json) = labels.get("id2label_json") {
+        if !id2label_json.is_empty() {
+            // Parse JSON: keys are string indices, values are label names
+            if let Ok(string_map) = serde_json::from_str::<HashMap<String, String>>(id2label_json) {
+                // Convert string keys ("0", "1") to u32 keys (0, 1)
+                let id2label: HashMap<u32, String> = string_map
+                    .into_iter()
+                    .filter_map(|(k, v)| k.parse::<u32>().ok().map(|idx| (idx, v)))
+                    .collect();
+
+                if !id2label.is_empty() {
+                    card = card.with_id2label(id2label);
+                    debug!("Parsed id2label with {} classes", card.num_labels);
+                }
+            }
+        }
+    }
+    // Fallback: if num_labels is set but id2label wasn't parsed, create default labels
+    // Match logic in serving_classify.py::_get_id2label_mapping
+    else if let Some(num_labels_str) = labels.get("num_labels") {
+        if let Ok(num_labels) = num_labels_str.parse::<u32>() {
+            if num_labels > 0 {
+                // Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...}
+                let id2label: HashMap<u32, String> = (0..num_labels)
+                    .map(|i| (i, format!("LABEL_{}", i)))
+                    .collect();
+                card = card.with_id2label(id2label);
+                debug!("Created default id2label with {} classes", num_labels);
+            }
+        }
+    }
+
     card
 }
 

sgl-model-gateway/src/core/worker.rs

@@ -242,6 +242,40 @@ pub trait Worker: Send + Sync + fmt::Debug {
         self.metadata().provider_for_model(model_id)
     }
 
+    /// Check if a model is a classifier (has id2label mapping).
+    fn is_classifier(&self, model_id: &str) -> bool {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.is_classifier())
+            .unwrap_or(false)
+    }
+
+    /// Get the id2label mapping for a classification model.
+    /// Returns None if model is not a classifier or not found.
+    fn id2label(&self, model_id: &str) -> Option<&std::collections::HashMap<u32, String>> {
+        self.metadata()
+            .find_model(model_id)
+            .filter(|m| m.is_classifier())
+            .map(|m| &m.id2label)
+    }
+
+    /// Get the number of classification labels for a model.
+    fn num_labels(&self, model_id: &str) -> u32 {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.num_labels)
+            .unwrap_or(0)
+    }
+
+    /// Get label for a class index from a classification model.
+    /// Returns generic label (LABEL_N) if model not found or index not in mapping.
+    fn get_label(&self, model_id: &str, class_idx: u32) -> String {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.get_label(class_idx))
+            .unwrap_or_else(|| format!("LABEL_{}", class_idx))
+    }
+
     /// Check if this worker supports a specific model.
     /// If models list is empty, worker accepts any model.
     fn supports_model(&self, model_id: &str) -> bool {

sgl-model-gateway/src/proto/sglang_scheduler.proto

@@ -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

sgl-model-gateway/src/routers/grpc/client.rs

@@ -105,7 +105,7 @@ impl GrpcClient {
         match self {
             Self::Sglang(client) => {
                 let info = client.get_model_info().await?;
-                Ok(ModelInfo::Sglang(info))
+                Ok(ModelInfo::Sglang(Box::new(info)))
             }
             Self::Vllm(client) => {
                 let info = client.get_model_info().await?;
@@ -151,7 +151,7 @@ impl GrpcClient {
 
 /// Unified ModelInfo wrapper
 pub enum ModelInfo {
-    Sglang(crate::grpc_client::sglang_proto::GetModelInfoResponse),
+    Sglang(Box<crate::grpc_client::sglang_proto::GetModelInfoResponse>),
     Vllm(crate::grpc_client::vllm_proto::GetModelInfoResponse),
 }