[model-gateway] multimodality initialization (#13350)
This commit is contained in:
@@ -4,7 +4,10 @@ version = "0.2.3"
|
||||
edition = "2021"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
default = ["grpc-client"]
|
||||
grpc-client = []
|
||||
grpc-server = []
|
||||
|
||||
vendored-openssl = ["openssl/vendored"]
|
||||
|
||||
[lints.rust]
|
||||
@@ -36,6 +39,7 @@ serde_json = { version = "1.0", default-features = false, features = [
|
||||
"std",
|
||||
"preserve_order",
|
||||
] }
|
||||
serde_bytes = "0.11"
|
||||
bytes = "1.8.0"
|
||||
rand = "0.9.2"
|
||||
reqwest = { version = "0.12.8", features = ["stream", "blocking", "json", "rustls-tls"], default-features = false }
|
||||
@@ -86,6 +90,9 @@ oracle = { version = "0.6.3", features = ["chrono"] }
|
||||
subtle = "2.6"
|
||||
rustpython-parser = "0.4.0"
|
||||
num-traits = "0.2"
|
||||
image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] }
|
||||
ndarray = "0.16"
|
||||
base64 = "0.22"
|
||||
openai-harmony = { git = "https://github.com/openai/harmony", tag = "v0.0.4" }
|
||||
openmetrics-parser = "0.4.4"
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ pub mod grpc_client;
|
||||
pub mod mcp;
|
||||
pub mod metrics;
|
||||
pub mod middleware;
|
||||
pub mod multimodal;
|
||||
pub mod policies;
|
||||
pub mod protocols;
|
||||
pub mod reasoning_parser;
|
||||
|
||||
47
sgl-router/src/multimodal/error.rs
Normal file
47
sgl-router/src/multimodal/error.rs
Normal file
@@ -0,0 +1,47 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use super::types::Modality;
|
||||
|
||||
pub type MultiModalResult<T> = Result<T, MultiModalError>;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum MediaConnectorError {
|
||||
#[error("unsupported media scheme: {0}")]
|
||||
UnsupportedScheme(String),
|
||||
#[error("invalid media URL: {0}")]
|
||||
InvalidUrl(String),
|
||||
#[error("media domain '{0}' is not in the allow list")]
|
||||
DisallowedDomain(String),
|
||||
#[error("local media path is not allowed: {0}")]
|
||||
DisallowedLocalPath(String),
|
||||
#[error("HTTP error while fetching media: {0}")]
|
||||
Http(#[from] reqwest::Error),
|
||||
#[error("I/O error while reading media: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("base64 decode error: {0}")]
|
||||
Base64Decode(#[from] base64::DecodeError),
|
||||
#[error("data URL parse error: {0}")]
|
||||
DataUrl(String),
|
||||
#[error("media decode task failed: {0}")]
|
||||
Blocking(#[from] tokio::task::JoinError),
|
||||
#[error("image decode error: {0}")]
|
||||
Image(#[from] image::ImageError),
|
||||
#[error("media fetch timed out after {0:?}")]
|
||||
Timeout(Duration),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum MultiModalError {
|
||||
#[error(transparent)]
|
||||
Media(#[from] MediaConnectorError),
|
||||
#[error("unsupported content part: {0}")]
|
||||
UnsupportedContent(&'static str),
|
||||
#[error("too many {modality:?} items provided. limit={limit}")]
|
||||
ModalityLimit { modality: Modality, limit: usize },
|
||||
#[error("tracker task join error: {0}")]
|
||||
Join(#[from] tokio::task::JoinError),
|
||||
#[error("tracker validation error: {0}")]
|
||||
Validation(String),
|
||||
}
|
||||
206
sgl-router/src/multimodal/media.rs
Normal file
206
sgl-router/src/multimodal/media.rs
Normal file
@@ -0,0 +1,206 @@
|
||||
use std::{collections::HashSet, path::PathBuf, sync::Arc, time::Duration};
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
|
||||
use bytes::Bytes;
|
||||
use image::DynamicImage;
|
||||
use reqwest::Client;
|
||||
use tokio::{fs, task};
|
||||
use url::Url;
|
||||
|
||||
use super::{
|
||||
error::MediaConnectorError,
|
||||
types::{ImageDetail, ImageFrame, ImageSource},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MediaConnectorConfig {
|
||||
pub allowed_domains: Option<Vec<String>>,
|
||||
pub allowed_local_media_path: Option<PathBuf>,
|
||||
pub fetch_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for MediaConnectorConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
allowed_domains: None,
|
||||
allowed_local_media_path: None,
|
||||
fetch_timeout: Duration::from_secs(10),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct ImageFetchConfig {
|
||||
pub detail: ImageDetail,
|
||||
}
|
||||
|
||||
impl Default for ImageFetchConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
detail: ImageDetail::Auto,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum MediaSource {
|
||||
Url(String),
|
||||
DataUrl(String),
|
||||
InlineBytes(Vec<u8>),
|
||||
File(PathBuf),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MediaConnector {
|
||||
client: Client,
|
||||
allowed_domains: Option<HashSet<String>>,
|
||||
allowed_local_media_path: Option<PathBuf>,
|
||||
fetch_timeout: Duration,
|
||||
}
|
||||
|
||||
impl MediaConnector {
|
||||
pub fn new(client: Client, config: MediaConnectorConfig) -> Result<Self, MediaConnectorError> {
|
||||
let allowed_domains = config.allowed_domains.map(|domains| {
|
||||
domains
|
||||
.into_iter()
|
||||
.map(|d| d.to_ascii_lowercase())
|
||||
.collect::<HashSet<_>>()
|
||||
});
|
||||
|
||||
let allowed_local_media_path = if let Some(path) = config.allowed_local_media_path {
|
||||
Some(std::fs::canonicalize(path)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
client,
|
||||
allowed_domains,
|
||||
allowed_local_media_path,
|
||||
fetch_timeout: config.fetch_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn fetch_image(
|
||||
&self,
|
||||
source: MediaSource,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, MediaConnectorError> {
|
||||
match source {
|
||||
MediaSource::Url(url) => self.fetch_http_image(url, cfg).await,
|
||||
MediaSource::DataUrl(data_url) => self.fetch_data_url(data_url, cfg).await,
|
||||
MediaSource::InlineBytes(bytes) => {
|
||||
self.decode_image(bytes.into(), cfg.detail, ImageSource::InlineBytes)
|
||||
.await
|
||||
}
|
||||
MediaSource::File(path) => self.fetch_file(path, cfg).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_http_image(
|
||||
&self,
|
||||
url: String,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, MediaConnectorError> {
|
||||
let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?;
|
||||
self.ensure_domain_allowed(&parsed)?;
|
||||
|
||||
let mut req = self.client.get(parsed.as_str());
|
||||
if self.fetch_timeout > Duration::ZERO {
|
||||
req = req.timeout(self.fetch_timeout);
|
||||
}
|
||||
|
||||
let resp = req.send().await.map_err(|err| {
|
||||
if err.is_timeout() {
|
||||
MediaConnectorError::Timeout(self.fetch_timeout)
|
||||
} else {
|
||||
MediaConnectorError::Http(err)
|
||||
}
|
||||
})?;
|
||||
|
||||
let resp = resp.error_for_status()?;
|
||||
let bytes = resp.bytes().await?;
|
||||
self.decode_image(
|
||||
bytes,
|
||||
cfg.detail,
|
||||
ImageSource::Url {
|
||||
url: parsed.to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn fetch_data_url(
|
||||
&self,
|
||||
data_url: String,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, MediaConnectorError> {
|
||||
let (metadata, data) = data_url
|
||||
.split_once(',')
|
||||
.ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?;
|
||||
|
||||
if !metadata.ends_with(";base64") {
|
||||
return Err(MediaConnectorError::DataUrl(
|
||||
"only base64 encoded data URLs are supported".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let data = data.trim();
|
||||
let decoded = BASE64_STANDARD.decode(data)?;
|
||||
self.decode_image(decoded.into(), cfg.detail, ImageSource::DataUrl)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn fetch_file(
|
||||
&self,
|
||||
path: PathBuf,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, MediaConnectorError> {
|
||||
let allowed_root = self
|
||||
.allowed_local_media_path
|
||||
.as_ref()
|
||||
.ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?;
|
||||
|
||||
let canonical = fs::canonicalize(&path).await?;
|
||||
if !canonical.starts_with(allowed_root) {
|
||||
return Err(MediaConnectorError::DisallowedLocalPath(
|
||||
path.display().to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let bytes = fs::read(&canonical).await?;
|
||||
self.decode_image(
|
||||
bytes.into(),
|
||||
cfg.detail,
|
||||
ImageSource::File { path: canonical },
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn ensure_domain_allowed(&self, url: &Url) -> Result<(), MediaConnectorError> {
|
||||
if let Some(allowed) = &self.allowed_domains {
|
||||
let host = url
|
||||
.host_str()
|
||||
.map(|h| h.to_ascii_lowercase())
|
||||
.ok_or_else(|| MediaConnectorError::InvalidUrl(url.to_string()))?;
|
||||
if !allowed.contains(&host) {
|
||||
return Err(MediaConnectorError::DisallowedDomain(host));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn decode_image(
|
||||
&self,
|
||||
bytes: Bytes,
|
||||
detail: ImageDetail,
|
||||
source: ImageSource,
|
||||
) -> Result<Arc<ImageFrame>, MediaConnectorError> {
|
||||
let raw: Arc<Vec<u8>> = Arc::new(bytes.to_vec());
|
||||
let raw_clone = raw.clone();
|
||||
let image: DynamicImage =
|
||||
task::spawn_blocking(move || image::load_from_memory(&raw_clone)).await??;
|
||||
|
||||
Ok(Arc::new(ImageFrame::new(image, raw, detail, source)))
|
||||
}
|
||||
}
|
||||
11
sgl-router/src/multimodal/mod.rs
Normal file
11
sgl-router/src/multimodal/mod.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
pub mod error;
|
||||
pub mod media;
|
||||
pub mod registry;
|
||||
pub mod tracker;
|
||||
pub mod types;
|
||||
|
||||
pub use error::{MediaConnectorError, MultiModalError, MultiModalResult};
|
||||
pub use media::{ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource};
|
||||
pub use registry::{ModelProcessorSpec, ModelRegistry};
|
||||
pub use tracker::{AsyncMultiModalTracker, TrackerConfig, TrackerOutput};
|
||||
pub use types::*;
|
||||
461
sgl-router/src/multimodal/registry.rs
Normal file
461
sgl-router/src/multimodal/registry.rs
Normal file
@@ -0,0 +1,461 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_json::{json, Value};
|
||||
use thiserror::Error;
|
||||
|
||||
use super::types::{ImageSize, Modality, PromptReplacement, TokenId};
|
||||
use crate::tokenizer::traits::Tokenizer as TokenizerTrait;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ModelRegistryError {
|
||||
#[error("unsupported model: {0}")]
|
||||
UnsupportedModel(String),
|
||||
#[error("token '{token}' not found in tokenizer vocabulary")]
|
||||
TokenNotFound { token: String },
|
||||
#[error("missing config field '{field}'")]
|
||||
MissingConfigField { field: String },
|
||||
}
|
||||
|
||||
pub type RegistryResult<T> = Result<T, ModelRegistryError>;
|
||||
|
||||
/// Metadata about the current model used to derive tokenizer/config dependent fields.
|
||||
pub struct ModelMetadata<'a> {
|
||||
pub model_id: &'a str,
|
||||
pub tokenizer: &'a dyn TokenizerTrait,
|
||||
pub config: &'a Value,
|
||||
}
|
||||
|
||||
impl<'a> ModelMetadata<'a> {
|
||||
pub fn token_id(&self, token: &str) -> RegistryResult<TokenId> {
|
||||
self.tokenizer
|
||||
.token_to_id(token)
|
||||
.map(|id| id as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::TokenNotFound {
|
||||
token: token.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config_u32(&self, path: &[&str]) -> Option<u32> {
|
||||
Self::find_value(self.config, path).and_then(|value| value.as_u64().map(|v| v as u32))
|
||||
}
|
||||
|
||||
fn find_value<'v>(value: &'v Value, path: &[&str]) -> Option<&'v Value> {
|
||||
let mut current = value;
|
||||
for key in path {
|
||||
current = current.get(*key)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ModelProcessorSpec: Send + Sync {
|
||||
fn name(&self) -> &'static str;
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool;
|
||||
fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult<String>;
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId>;
|
||||
fn modality_limits(&self, metadata: &ModelMetadata)
|
||||
-> RegistryResult<HashMap<Modality, usize>>;
|
||||
fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult<Value>;
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>>;
|
||||
}
|
||||
|
||||
pub struct ModelRegistry {
|
||||
specs: Vec<LazySpec>,
|
||||
}
|
||||
|
||||
impl ModelRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
specs: vec![
|
||||
LazySpec::new("llava", || Box::new(LlavaSpec)),
|
||||
LazySpec::new("qwen_vl", || Box::new(QwenVLVisionSpec)),
|
||||
LazySpec::new("phi3_v", || Box::new(Phi3VisionSpec)),
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lookup<'a>(&'a self, metadata: &ModelMetadata) -> Option<&'a dyn ModelProcessorSpec> {
|
||||
for spec in &self.specs {
|
||||
let spec_ref = spec.get();
|
||||
if spec_ref.matches(metadata) {
|
||||
return Some(spec_ref);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ModelRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
struct LazySpec {
|
||||
inner: Lazy<Box<dyn ModelProcessorSpec>>,
|
||||
}
|
||||
|
||||
impl LazySpec {
|
||||
fn new(_id: &'static str, factory: fn() -> Box<dyn ModelProcessorSpec>) -> Self {
|
||||
Self {
|
||||
inner: Lazy::new(factory),
|
||||
}
|
||||
}
|
||||
|
||||
fn get(&self) -> &dyn ModelProcessorSpec {
|
||||
self.inner.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
struct LlavaSpec;
|
||||
|
||||
impl LlavaSpec {
|
||||
fn patch_size(metadata: &ModelMetadata) -> u32 {
|
||||
metadata
|
||||
.config_u32(&["vision_config", "patch_size"])
|
||||
.unwrap_or(14)
|
||||
}
|
||||
|
||||
fn tokens_per_image(metadata: &ModelMetadata, size: &ImageSize) -> usize {
|
||||
let patch = Self::patch_size(metadata);
|
||||
let cols = size.width.div_ceil(patch) as usize;
|
||||
let rows = size.height.div_ceil(patch) as usize;
|
||||
cols * rows
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelProcessorSpec for LlavaSpec {
|
||||
fn name(&self) -> &'static str {
|
||||
"llava"
|
||||
}
|
||||
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool {
|
||||
metadata.model_id.to_ascii_lowercase().contains("llava")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_string())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
if let Some(value) = metadata.config_u32(&["image_token_index"]) {
|
||||
return Ok(value as TokenId);
|
||||
}
|
||||
metadata.token_id("<image>")
|
||||
}
|
||||
|
||||
fn modality_limits(
|
||||
&self,
|
||||
_metadata: &ModelMetadata,
|
||||
) -> RegistryResult<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 4)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let token_id = self.placeholder_token_id(metadata)?;
|
||||
let token = self.placeholder_token(metadata)?;
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|size| {
|
||||
let count = Self::tokens_per_image(metadata, size);
|
||||
PromptReplacement::repeated(Modality::Image, &token, token_id, count)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
struct QwenVLVisionSpec;
|
||||
|
||||
impl QwenVLVisionSpec {
|
||||
fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["vision_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "vision_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn start_token_id(metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["vision_start_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "vision_start_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn patch_grid(metadata: &ModelMetadata, size: &ImageSize) -> (usize, usize) {
|
||||
let patch = metadata
|
||||
.config_u32(&["vision_config", "patch_size"])
|
||||
.unwrap_or(14);
|
||||
let cols = size.width.div_ceil(patch) as usize;
|
||||
let rows = size.height.div_ceil(patch) as usize;
|
||||
(rows, cols)
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelProcessorSpec for QwenVLVisionSpec {
|
||||
fn name(&self) -> &'static str {
|
||||
"qwen_vl"
|
||||
}
|
||||
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool {
|
||||
metadata.model_id.to_ascii_lowercase().contains("qwen")
|
||||
&& metadata.model_id.to_ascii_lowercase().contains("vl")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_string())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["image_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "image_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn modality_limits(
|
||||
&self,
|
||||
_metadata: &ModelMetadata,
|
||||
) -> RegistryResult<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 10)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let start_token_id = Self::start_token_id(metadata)?;
|
||||
let pad_token_id = Self::pad_token_id(metadata)?;
|
||||
let placeholder_token = self.placeholder_token(metadata)?;
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|size| {
|
||||
let (rows, cols) = Self::patch_grid(metadata, size);
|
||||
let pad_len = rows * cols;
|
||||
let mut tokens = Vec::with_capacity(pad_len + 1);
|
||||
tokens.push(start_token_id);
|
||||
tokens.extend(std::iter::repeat_n(pad_token_id, pad_len));
|
||||
PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
struct Phi3VisionSpec;
|
||||
|
||||
impl Phi3VisionSpec {
|
||||
fn tokens_per_image(metadata: &ModelMetadata) -> usize {
|
||||
metadata
|
||||
.config_u32(&["img_processor", "num_img_tokens"])
|
||||
.unwrap_or(256) as usize
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelProcessorSpec for Phi3VisionSpec {
|
||||
fn name(&self) -> &'static str {
|
||||
"phi3_v"
|
||||
}
|
||||
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool {
|
||||
let id = metadata.model_id.to_ascii_lowercase();
|
||||
id.contains("phi") && id.contains("vision")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_owned())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata.token_id("<image>")
|
||||
}
|
||||
|
||||
fn modality_limits(
|
||||
&self,
|
||||
_metadata: &ModelMetadata,
|
||||
) -> RegistryResult<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 4)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let token_id = self.placeholder_token_id(metadata)?;
|
||||
let token = self.placeholder_token(metadata)?;
|
||||
let count = Self::tokens_per_image(metadata);
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|_| PromptReplacement::repeated(Modality::Image, &token, token_id, count))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::tokenizer::traits::{
|
||||
Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait,
|
||||
};
|
||||
|
||||
struct TestTokenizer {
|
||||
vocab: HashMap<String, u32>,
|
||||
}
|
||||
|
||||
impl TestTokenizer {
|
||||
fn new(pairs: &[(&str, u32)]) -> Self {
|
||||
let vocab = pairs
|
||||
.iter()
|
||||
.map(|(token, id)| ((*token).to_string(), *id))
|
||||
.collect();
|
||||
Self { vocab }
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for TestTokenizer {
|
||||
fn encode(&self, _input: &str) -> anyhow::Result<Encoding> {
|
||||
Ok(Encoding::Sp(Vec::new()))
|
||||
}
|
||||
|
||||
fn encode_batch(&self, inputs: &[&str]) -> anyhow::Result<Vec<Encoding>> {
|
||||
inputs.iter().map(|_| self.encode("")).collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for TestTokenizer {
|
||||
fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerTrait for TestTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.vocab.len()
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
static TOKENS: Lazy<SpecialTokens> = Lazy::new(|| SpecialTokens {
|
||||
bos_token: None,
|
||||
eos_token: None,
|
||||
unk_token: None,
|
||||
sep_token: None,
|
||||
pad_token: None,
|
||||
cls_token: None,
|
||||
mask_token: None,
|
||||
additional_special_tokens: vec![],
|
||||
});
|
||||
&TOKENS
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
self.vocab.get(token).copied()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llava_prompt_replacement_uses_config_ids() {
|
||||
let tokenizer = TestTokenizer::new(&[("<image>", 32000)]);
|
||||
let config = json!({
|
||||
"model_type": "llava",
|
||||
"image_token_index": 32000,
|
||||
"vision_config": {"patch_size": 14}
|
||||
});
|
||||
let metadata = ModelMetadata {
|
||||
model_id: "llava-v1.5",
|
||||
tokenizer: &tokenizer,
|
||||
config: &config,
|
||||
};
|
||||
let registry = ModelRegistry::new();
|
||||
let spec = registry.lookup(&metadata).expect("llava spec");
|
||||
let replacements = spec
|
||||
.prompt_replacements(&metadata, &[ImageSize::new(336, 336)])
|
||||
.unwrap();
|
||||
assert_eq!(replacements[0].tokens.len(), 576);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen_vision_uses_config_token_ids() {
|
||||
let tokenizer = TestTokenizer::new(&[("<image>", 999)]);
|
||||
let config = json!({
|
||||
"model_type": "qwen2_vl",
|
||||
"vision_start_token_id": 151652,
|
||||
"vision_token_id": 151654,
|
||||
"image_token_id": 151655,
|
||||
"vision_config": {"patch_size": 14}
|
||||
});
|
||||
let metadata = ModelMetadata {
|
||||
model_id: "Qwen2-VL-7B",
|
||||
tokenizer: &tokenizer,
|
||||
config: &config,
|
||||
};
|
||||
let registry = ModelRegistry::new();
|
||||
let spec = registry.lookup(&metadata).expect("qwen spec");
|
||||
let replacements = spec
|
||||
.prompt_replacements(&metadata, &[ImageSize::new(448, 448)])
|
||||
.unwrap();
|
||||
// 448/14 = 32 patches => 1024 pad tokens + 1 start token
|
||||
assert_eq!(replacements[0].tokens.len(), 1025);
|
||||
assert_eq!(replacements[0].tokens[0], 151652);
|
||||
assert_eq!(replacements[0].tokens[1], 151654);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn phi3_uses_num_img_tokens() {
|
||||
let tokenizer = TestTokenizer::new(&[("<image>", 555)]);
|
||||
let config = json!({
|
||||
"model_type": "phi3_v",
|
||||
"img_processor": {"num_img_tokens": 144}
|
||||
});
|
||||
let metadata = ModelMetadata {
|
||||
model_id: "Phi-3-vision",
|
||||
tokenizer: &tokenizer,
|
||||
config: &config,
|
||||
};
|
||||
let registry = ModelRegistry::new();
|
||||
let spec = registry.lookup(&metadata).expect("phi3 spec");
|
||||
let replacements = spec
|
||||
.prompt_replacements(&metadata, &[ImageSize::new(336, 336)])
|
||||
.unwrap();
|
||||
assert_eq!(replacements[0].tokens.len(), 144);
|
||||
assert_eq!(replacements[0].tokens[0], 555);
|
||||
}
|
||||
}
|
||||
176
sgl-router/src/multimodal/tracker.rs
Normal file
176
sgl-router/src/multimodal/tracker.rs
Normal file
@@ -0,0 +1,176 @@
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use super::{
|
||||
error::{MultiModalError, MultiModalResult},
|
||||
media::{ImageFetchConfig, MediaConnector, MediaSource},
|
||||
types::{
|
||||
ChatContentPart, ConversationSegment, ImageDetail, Modality, MultiModalData,
|
||||
MultiModalUUIDs, PlaceholderHandle, PlaceholderMap, TrackedMedia, DEFAULT_PLACEHOLDERS,
|
||||
},
|
||||
};
|
||||
|
||||
type PendingTask = JoinHandle<MultiModalResult<TrackedMedia>>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrackerConfig {
|
||||
pub placeholder_tokens: HashMap<Modality, String>,
|
||||
pub modality_limits: HashMap<Modality, usize>,
|
||||
}
|
||||
|
||||
impl Default for TrackerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
placeholder_tokens: DEFAULT_PLACEHOLDERS
|
||||
.iter()
|
||||
.map(|(k, v)| (*k, (*v).to_string()))
|
||||
.collect(),
|
||||
modality_limits: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TrackerOutput {
|
||||
pub conversation: Vec<ConversationSegment>,
|
||||
pub data: MultiModalData,
|
||||
pub uuids: MultiModalUUIDs,
|
||||
pub placeholders: PlaceholderMap,
|
||||
}
|
||||
|
||||
pub struct AsyncMultiModalTracker {
|
||||
media_connector: Arc<MediaConnector>,
|
||||
config: TrackerConfig,
|
||||
pending: HashMap<Modality, Vec<PendingTask>>,
|
||||
placeholders: PlaceholderMap,
|
||||
conversation: Vec<ConversationSegment>,
|
||||
uuids: MultiModalUUIDs,
|
||||
counts: HashMap<Modality, usize>,
|
||||
}
|
||||
|
||||
impl AsyncMultiModalTracker {
|
||||
pub fn new(media_connector: Arc<MediaConnector>, config: TrackerConfig) -> Self {
|
||||
Self {
|
||||
media_connector,
|
||||
config,
|
||||
pending: HashMap::new(),
|
||||
placeholders: PlaceholderMap::new(),
|
||||
conversation: Vec::new(),
|
||||
uuids: HashMap::new(),
|
||||
counts: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_part(&mut self, part: ChatContentPart) -> MultiModalResult<()> {
|
||||
match part {
|
||||
ChatContentPart::Text { text } => {
|
||||
if !text.is_empty() {
|
||||
self.conversation.push(ConversationSegment::text(text));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
ChatContentPart::ImageUrl { url, detail, uuid } => {
|
||||
self.enqueue_image(MediaSource::Url(url), detail.unwrap_or_default(), uuid)
|
||||
}
|
||||
ChatContentPart::ImageData {
|
||||
data,
|
||||
mime_type: _,
|
||||
uuid,
|
||||
detail,
|
||||
} => self.enqueue_image(
|
||||
MediaSource::InlineBytes(data),
|
||||
detail.unwrap_or_default(),
|
||||
uuid,
|
||||
),
|
||||
ChatContentPart::ImageEmbeds { .. } => {
|
||||
Err(MultiModalError::UnsupportedContent("image_embeds"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn finalize(mut self) -> MultiModalResult<TrackerOutput> {
|
||||
let mut data = MultiModalData::new();
|
||||
for (modality, tasks) in self.pending.drain() {
|
||||
let mut items = Vec::with_capacity(tasks.len());
|
||||
for task in tasks {
|
||||
let media = task.await??;
|
||||
items.push(media);
|
||||
}
|
||||
data.insert(modality, items);
|
||||
}
|
||||
|
||||
Ok(TrackerOutput {
|
||||
conversation: self.conversation,
|
||||
data,
|
||||
uuids: self.uuids,
|
||||
placeholders: self.placeholders,
|
||||
})
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, modality: Modality) -> &str {
|
||||
if let Some(token) = self.config.placeholder_tokens.get(&modality) {
|
||||
token.as_str()
|
||||
} else {
|
||||
DEFAULT_PLACEHOLDERS
|
||||
.get(&modality)
|
||||
.copied()
|
||||
.unwrap_or("<mm>")
|
||||
}
|
||||
}
|
||||
|
||||
fn next_index(&mut self, modality: Modality) -> MultiModalResult<usize> {
|
||||
let count = self.counts.entry(modality).or_insert(0);
|
||||
let next = *count;
|
||||
let limit = self.config.modality_limits.get(&modality).copied();
|
||||
if let Some(limit) = limit {
|
||||
if next >= limit {
|
||||
return Err(MultiModalError::ModalityLimit { modality, limit });
|
||||
}
|
||||
}
|
||||
*count += 1;
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
fn record_uuid(&mut self, modality: Modality, uuid: Option<String>) {
|
||||
self.uuids.entry(modality).or_default().push(uuid);
|
||||
}
|
||||
|
||||
fn add_placeholder(&mut self, token: String, handle: PlaceholderHandle) {
|
||||
self.placeholders.entry(token).or_default().push(handle);
|
||||
}
|
||||
|
||||
fn enqueue_image(
|
||||
&mut self,
|
||||
source: MediaSource,
|
||||
detail: ImageDetail,
|
||||
uuid: Option<String>,
|
||||
) -> MultiModalResult<()> {
|
||||
let modality = Modality::Image;
|
||||
let idx = self.next_index(modality)?;
|
||||
let token = self.placeholder_token(modality).to_string();
|
||||
let text_position = self.conversation.len();
|
||||
self.conversation
|
||||
.push(ConversationSegment::placeholder(token.clone()));
|
||||
self.add_placeholder(
|
||||
token.clone(),
|
||||
PlaceholderHandle {
|
||||
modality,
|
||||
item_index: idx,
|
||||
text_position,
|
||||
},
|
||||
);
|
||||
self.record_uuid(modality, uuid);
|
||||
|
||||
let connector = Arc::clone(&self.media_connector);
|
||||
let handle = tokio::spawn(async move {
|
||||
let frame = connector
|
||||
.fetch_image(source, ImageFetchConfig { detail })
|
||||
.await?;
|
||||
Ok(TrackedMedia::Image(frame))
|
||||
});
|
||||
|
||||
self.pending.entry(modality).or_default().push(handle);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
290
sgl-router/src/multimodal/types.rs
Normal file
290
sgl-router/src/multimodal/types.rs
Normal file
@@ -0,0 +1,290 @@
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap},
|
||||
fmt,
|
||||
path::PathBuf,
|
||||
sync::{Arc, LazyLock},
|
||||
};
|
||||
|
||||
use image::DynamicImage;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Supported multimodal modalities.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Modality {
|
||||
Image,
|
||||
ImageEmbeds,
|
||||
Audio,
|
||||
Video,
|
||||
}
|
||||
|
||||
impl fmt::Display for Modality {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Modality::Image => write!(f, "image"),
|
||||
Modality::ImageEmbeds => write!(f, "image_embeds"),
|
||||
Modality::Audio => write!(f, "audio"),
|
||||
Modality::Video => write!(f, "video"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple helper used for default placeholder tokens per modality.
|
||||
pub static DEFAULT_PLACEHOLDERS: LazyLock<HashMap<Modality, &'static str>> = LazyLock::new(|| {
|
||||
HashMap::from([
|
||||
(Modality::Image, "<image>"),
|
||||
(Modality::ImageEmbeds, "<image_embeds>"),
|
||||
(Modality::Audio, "<audio>"),
|
||||
(Modality::Video, "<video>"),
|
||||
])
|
||||
});
|
||||
|
||||
/// Detail level passed by OpenAI style APIs.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ImageDetail {
|
||||
#[default]
|
||||
Auto,
|
||||
Low,
|
||||
High,
|
||||
}
|
||||
|
||||
/// A normalized content part understood by the tracker.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ChatContentPart {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
ImageUrl {
|
||||
url: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
detail: Option<ImageDetail>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
},
|
||||
ImageData {
|
||||
data: Vec<u8>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
mime_type: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
detail: Option<ImageDetail>,
|
||||
},
|
||||
ImageEmbeds {
|
||||
payload: Value,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Represents one segment in the conversation after placeholders are inserted.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ConversationSegment {
|
||||
Text(String),
|
||||
Placeholder { token: String },
|
||||
}
|
||||
|
||||
impl ConversationSegment {
|
||||
pub fn text<S: Into<String>>(text: S) -> Self {
|
||||
ConversationSegment::Text(text.into())
|
||||
}
|
||||
|
||||
pub fn placeholder<S: Into<String>>(token: S) -> Self {
|
||||
ConversationSegment::Placeholder {
|
||||
token: token.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metadata describing where placeholder tokens were inserted.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PlaceholderHandle {
|
||||
pub modality: Modality,
|
||||
pub item_index: usize,
|
||||
pub text_position: usize,
|
||||
}
|
||||
|
||||
pub type PlaceholderMap = HashMap<String, Vec<PlaceholderHandle>>;
|
||||
|
||||
/// Image source metadata (useful for hashing & tracing).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum ImageSource {
|
||||
Url { url: String },
|
||||
DataUrl,
|
||||
InlineBytes,
|
||||
File { path: PathBuf },
|
||||
}
|
||||
|
||||
/// Concrete image payload captured by the media connector.
|
||||
#[derive(Debug)]
|
||||
pub struct ImageFrame {
|
||||
image: DynamicImage,
|
||||
raw_bytes: Arc<Vec<u8>>,
|
||||
pub detail: ImageDetail,
|
||||
pub source: ImageSource,
|
||||
}
|
||||
|
||||
impl ImageFrame {
|
||||
pub fn new(
|
||||
image: DynamicImage,
|
||||
raw_bytes: Arc<Vec<u8>>,
|
||||
detail: ImageDetail,
|
||||
source: ImageSource,
|
||||
) -> Self {
|
||||
Self {
|
||||
image,
|
||||
raw_bytes,
|
||||
detail,
|
||||
source,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn data(&self) -> &DynamicImage {
|
||||
&self.image
|
||||
}
|
||||
|
||||
pub fn raw_bytes(&self) -> &[u8] {
|
||||
self.raw_bytes.as_slice()
|
||||
}
|
||||
|
||||
pub fn source(&self) -> &ImageSource {
|
||||
&self.source
|
||||
}
|
||||
|
||||
pub fn size(&self) -> ImageSize {
|
||||
ImageSize::new(self.image.width(), self.image.height())
|
||||
}
|
||||
}
|
||||
|
||||
/// Container for all supported multimodal media objects.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TrackedMedia {
|
||||
Image(Arc<ImageFrame>),
|
||||
/// Placeholder variants for future modalities.
|
||||
Audio,
|
||||
Video,
|
||||
Embeddings,
|
||||
}
|
||||
|
||||
pub type MultiModalData = HashMap<Modality, Vec<TrackedMedia>>;
|
||||
pub type MultiModalUUIDs = HashMap<Modality, Vec<Option<String>>>;
|
||||
|
||||
pub type TokenId = i32;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ImageSize {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
}
|
||||
|
||||
impl ImageSize {
|
||||
pub fn new(width: u32, height: u32) -> Self {
|
||||
Self { width, height }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct PlaceholderRange {
|
||||
pub offset: usize,
|
||||
pub length: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MultiModalTensor {
|
||||
pub shape: Vec<usize>,
|
||||
pub dtype: String,
|
||||
#[serde(with = "serde_bytes")]
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum MultiModalValue {
|
||||
Tensor(MultiModalTensor),
|
||||
Json(Value),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct MultiModalInputs {
|
||||
pub prompt_token_ids: Vec<u32>,
|
||||
#[serde(default)]
|
||||
pub mm_kwargs: BTreeMap<String, Vec<MultiModalValue>>,
|
||||
#[serde(default)]
|
||||
pub mm_hashes: BTreeMap<String, Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub mm_placeholders: BTreeMap<String, Vec<PlaceholderRange>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_salt: Option<String>,
|
||||
}
|
||||
|
||||
impl MultiModalInputs {
|
||||
pub fn new(prompt_token_ids: Vec<u32>) -> Self {
|
||||
Self {
|
||||
prompt_token_ids,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptReplacement {
|
||||
pub modality: Modality,
|
||||
pub placeholder_token: String,
|
||||
pub tokens: Vec<TokenId>,
|
||||
}
|
||||
|
||||
impl PromptReplacement {
|
||||
pub fn repeated(
|
||||
modality: Modality,
|
||||
placeholder_token: &str,
|
||||
token_id: TokenId,
|
||||
count: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
modality,
|
||||
placeholder_token: placeholder_token.to_string(),
|
||||
tokens: vec![token_id; count],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sequence(modality: Modality, placeholder_token: &str, sequence: Vec<TokenId>) -> Self {
|
||||
Self {
|
||||
modality,
|
||||
placeholder_token: placeholder_token.to_string(),
|
||||
tokens: sequence,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn multimodal_inputs_defaults() {
|
||||
let inputs = MultiModalInputs::new(vec![1, 2, 3]);
|
||||
assert_eq!(inputs.prompt_token_ids, vec![1, 2, 3]);
|
||||
assert!(inputs.mm_kwargs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn placeholder_range_serializes() {
|
||||
let range = PlaceholderRange {
|
||||
offset: 10,
|
||||
length: 4,
|
||||
};
|
||||
let json = serde_json::to_string(&range).unwrap();
|
||||
assert!(json.contains("offset"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_replacement_builders() {
|
||||
let rep = PromptReplacement::repeated(Modality::Image, "<image>", 100, 3);
|
||||
assert_eq!(rep.tokens, vec![100, 100, 100]);
|
||||
}
|
||||
}
|
||||
151
sgl-router/tests/multimodal_tracker_test.rs
Normal file
151
sgl-router/tests/multimodal_tracker_test.rs
Normal file
@@ -0,0 +1,151 @@
|
||||
use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
|
||||
use reqwest::Client;
|
||||
use sgl_model_gateway::multimodal::{
|
||||
AsyncMultiModalTracker, ChatContentPart, ConversationSegment, ImageFetchConfig, ImageSource,
|
||||
MediaConnector, MediaConnectorConfig, MediaSource, Modality, TrackerConfig,
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
const TINY_PNG_BASE64: &str =
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=";
|
||||
|
||||
fn tiny_png_bytes() -> Vec<u8> {
|
||||
BASE64_STANDARD
|
||||
.decode(TINY_PNG_BASE64)
|
||||
.expect("decode tiny png fixture")
|
||||
}
|
||||
|
||||
fn test_connector(allowed_path: Option<PathBuf>) -> MediaConnector {
|
||||
let client = Client::builder()
|
||||
.timeout(Duration::from_secs(5))
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client");
|
||||
MediaConnector::new(
|
||||
client,
|
||||
MediaConnectorConfig {
|
||||
allowed_domains: None,
|
||||
allowed_local_media_path: allowed_path,
|
||||
fetch_timeout: Duration::from_secs(5),
|
||||
},
|
||||
)
|
||||
.expect("media connector")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_image_from_inline_bytes() {
|
||||
let connector = test_connector(None);
|
||||
let bytes = tiny_png_bytes();
|
||||
let frame = connector
|
||||
.fetch_image(
|
||||
MediaSource::InlineBytes(bytes.clone()),
|
||||
ImageFetchConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("inline image");
|
||||
assert_eq!(frame.data().width(), 1);
|
||||
assert_eq!(frame.data().height(), 1);
|
||||
assert_eq!(frame.raw_bytes(), bytes.as_slice());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_image_from_data_url() {
|
||||
let connector = test_connector(None);
|
||||
let bytes = tiny_png_bytes();
|
||||
let data_url = format!(
|
||||
"data:image/png;base64,{}",
|
||||
BASE64_STANDARD.encode(bytes.clone())
|
||||
);
|
||||
|
||||
let frame = connector
|
||||
.fetch_image(MediaSource::DataUrl(data_url), ImageFetchConfig::default())
|
||||
.await
|
||||
.expect("data url");
|
||||
assert_eq!(frame.data().width(), 1);
|
||||
assert_eq!(frame.raw_bytes(), bytes.as_slice());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_image_from_file() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let allowed_root = std::fs::canonicalize(tmp.path()).expect("canonical tmp path");
|
||||
let file_path = allowed_root.join("tiny.png");
|
||||
std::fs::write(&file_path, tiny_png_bytes()).expect("write png");
|
||||
|
||||
let connector = test_connector(Some(allowed_root));
|
||||
let frame = connector
|
||||
.fetch_image(
|
||||
MediaSource::File(file_path.clone()),
|
||||
ImageFetchConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("file png");
|
||||
assert_eq!(frame.data().width(), 1);
|
||||
let expected = std::fs::canonicalize(&file_path).expect("canonical path");
|
||||
match frame.source() {
|
||||
ImageSource::File { path } => assert_eq!(path, &expected),
|
||||
other => panic!("expected file source, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tracker_collects_conversation_and_placeholders() {
|
||||
let connector = Arc::new(test_connector(None));
|
||||
let mut tracker = AsyncMultiModalTracker::new(
|
||||
connector,
|
||||
TrackerConfig {
|
||||
placeholder_tokens: Default::default(),
|
||||
modality_limits: HashMap::from([(Modality::Image, 2)]),
|
||||
},
|
||||
);
|
||||
|
||||
tracker
|
||||
.push_part(ChatContentPart::Text {
|
||||
text: "before".into(),
|
||||
})
|
||||
.expect("text part");
|
||||
tracker
|
||||
.push_part(ChatContentPart::ImageData {
|
||||
data: tiny_png_bytes(),
|
||||
mime_type: Some("image/png".into()),
|
||||
uuid: Some("img-1".into()),
|
||||
detail: None,
|
||||
})
|
||||
.expect("image part");
|
||||
tracker
|
||||
.push_part(ChatContentPart::Text {
|
||||
text: "after".into(),
|
||||
})
|
||||
.expect("text part");
|
||||
|
||||
let output = tracker.finalize().await.expect("tracker finalize");
|
||||
assert_eq!(output.conversation.len(), 3);
|
||||
assert!(matches!(
|
||||
&output.conversation[0],
|
||||
ConversationSegment::Text(text) if text == "before"
|
||||
));
|
||||
assert!(matches!(
|
||||
&output.conversation[1],
|
||||
ConversationSegment::Placeholder { token } if token == "<image>"
|
||||
));
|
||||
assert!(matches!(
|
||||
&output.conversation[2],
|
||||
ConversationSegment::Text(text) if text == "after"
|
||||
));
|
||||
|
||||
let images = output.data.get(&Modality::Image).expect("image entry");
|
||||
assert_eq!(images.len(), 1);
|
||||
|
||||
let uuids = output.uuids.get(&Modality::Image).expect("uuid entry");
|
||||
assert_eq!(uuids, &vec![Some("img-1".into())]);
|
||||
|
||||
let placeholders = output
|
||||
.placeholders
|
||||
.get("<image>")
|
||||
.expect("placeholder entry");
|
||||
assert_eq!(placeholders.len(), 1);
|
||||
assert_eq!(placeholders[0].text_position, 1);
|
||||
assert_eq!(placeholders[0].item_index, 0);
|
||||
}
|
||||
Reference in New Issue
Block a user