From 54ff24b2cb9e42658db4b8e6c7f0a80b76124cb3 Mon Sep 17 00:00:00 2001 From: KSW-KSM Date: Mon, 29 Dec 2025 14:12:33 +0900 Subject: [PATCH 1/3] feat: add Jina v4 candle support --- backends/candle/src/lib.rs | 159 +++++++++++++++++++- backends/candle/src/models/flash_qwen2.rs | 172 +++++++++++++++++++--- backends/candle/src/models/mod.rs | 2 +- backends/src/lib.rs | 43 +++++- router/src/lib.rs | 47 +++++- 5 files changed, 387 insertions(+), 36 deletions(-) diff --git a/backends/candle/src/lib.rs b/backends/candle/src/lib.rs index ff824f555..e817a33cc 100644 --- a/backends/candle/src/lib.rs +++ b/backends/candle/src/lib.rs @@ -31,8 +31,10 @@ use crate::models::{ use crate::models::{ FlashBertModel, FlashDistilBertModel, FlashGTEModel, FlashJinaBertModel, FlashJinaCodeBertModel, FlashMistralModel, FlashModernBertModel, FlashNomicBertModel, - FlashQwen2Model, FlashQwen3Model, + FlashQwen2Model, FlashQwen3Model, LoraWeights, }; +#[cfg(feature = "cuda")] +use std::{env, fs}; /// This enum is needed to be able to differentiate between jina models that also use /// the `bert` model type and valid Bert models. @@ -88,6 +90,106 @@ impl<'de> Deserialize<'de> for BertConfigWrapper { } } +#[derive(Debug, Clone, Deserialize)] +struct JinaV4Config { + #[serde(default)] + task_names: Vec, + text_config: Qwen2Config, +} + +fn is_jina_v4_config(value: &serde_json::Value) -> bool { + value + .get("architectures") + .and_then(|v| v.as_array()) + .map(|items| { + items + .iter() + .any(|item| item.as_str() == Some("JinaEmbeddingsV4Model")) + }) + .unwrap_or(false) +} + +#[cfg(feature = "cuda")] +fn load_jina_v4_lora( + model_path: &Path, + device: &Device, + dtype: DType, + config: &JinaV4Config, +) -> Option { + #[derive(Deserialize)] + struct AdapterConfig { + r: usize, + lora_alpha: f32, + } + + let adapter_dir = model_path.join("adapters"); + let adapter_config_path = adapter_dir.join("adapter_config.json"); + let adapter_model_path = adapter_dir.join("adapter_model.safetensors"); + + if !adapter_config_path.exists() || !adapter_model_path.exists() { + tracing::warn!("Jina v4 adapters not found; LoRA will be skipped."); + return None; + } + + let adapter_config = match fs::read_to_string(&adapter_config_path) { + Ok(content) => match serde_json::from_str::(&content) { + Ok(config) => config, + Err(err) => { + tracing::warn!("Failed to parse Jina v4 adapter_config.json: {err}"); + return None; + } + }, + Err(err) => { + tracing::warn!("Failed to read Jina v4 adapter_config.json: {err}"); + return None; + } + }; + + let mut task = env::var("JINA_V4_TASK").unwrap_or_default(); + if task.is_empty() { + task = config + .task_names + .first() + .cloned() + .unwrap_or_else(|| "retrieval".to_string()); + } else if !config.task_names.is_empty() && !config.task_names.contains(&task) { + tracing::warn!( + "JINA_V4_TASK={task} is not in config.task_names; defaulting to the first entry." + ); + task = config + .task_names + .first() + .cloned() + .unwrap_or_else(|| "retrieval".to_string()); + } + + let adapter_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[adapter_model_path], dtype, device) }; + let adapter_vb = match adapter_vb.s() { + Ok(vb) => vb, + Err(err) => { + tracing::warn!("Failed to load Jina v4 adapter weights: {err}"); + return None; + } + }; + + let lora_prefix = "base_model.model.model.language_model".to_string(); + let lora_check = format!( + "{lora_prefix}.layers.0.self_attn.q_proj.lora_A.{task}.weight" + ); + if !adapter_vb.contains_tensor(&lora_check) { + tracing::warn!("Jina v4 adapter weights missing expected keys; LoRA will be skipped."); + return None; + } + + Some(LoraWeights::new( + adapter_vb, + task, + adapter_config.r, + adapter_config.lora_alpha, + lora_prefix, + )) +} + #[derive(Deserialize)] #[serde(tag = "model_type", rename_all = "kebab-case")] enum Config { @@ -111,6 +213,7 @@ enum Config { Qwen2(Qwen2Config), #[allow(dead_code)] Qwen3(Qwen3Config), + JinaV4(JinaV4Config), Roberta(BertConfig), XlmRoberta(BertConfig), } @@ -180,9 +283,28 @@ impl CandleBackend { let config: String = std::fs::read_to_string(model_path.join("config.json")) .context("Unable to read config file") .map_err(|err| BackendError::Start(format!("{err:?}")))?; - let config: Config = serde_json::from_str(&config) + let config_value: serde_json::Value = serde_json::from_str(&config) .context("Model is not supported") .map_err(|err| BackendError::Start(format!("{err:?}")))?; + let config: Config = if is_jina_v4_config(&config_value) { + if config_value + .get("text_config") + .and_then(|text| text.get("rope_scaling")) + .is_some() + { + tracing::warn!( + "Jina v4 rope_scaling is not supported in Candle; using base rope instead." + ); + } + let jina_config: JinaV4Config = serde_json::from_value(config_value.clone()) + .context("Model is not supported") + .map_err(|err| BackendError::Start(format!("{err:?}")))?; + Config::JinaV4(jina_config) + } else { + serde_json::from_value(config_value) + .context("Model is not supported") + .map_err(|err| BackendError::Start(format!("{err:?}")))? + }; // Get candle device let device = if candle::utils::cuda_is_available() { @@ -301,6 +423,10 @@ impl CandleBackend { "Qwen2 is only supported on Cuda devices in fp16 with flash attention enabled" .to_string(), )), + (Config::JinaV4(_), Device::Cpu | Device::Metal(_)) => Err(BackendError::Start( + "Jina v4 is only supported on Cuda devices in fp16 with flash attention enabled" + .to_string(), + )), (Config::Qwen3(config), Device::Cpu | Device::Metal(_)) => { tracing::info!("Starting Qwen3 model on {:?}", device); Ok(Box::new(Qwen3Model::load(vb, &config, model_type).s()?)) @@ -488,7 +614,34 @@ impl CandleBackend { } tracing::info!("Starting FlashQwen2 model on {:?}", device); Ok(Box::new( - FlashQwen2Model::load(vb, &config, model_type).s()?, + FlashQwen2Model::load(vb, &config, model_type, None, false).s()?, + )) + } + #[cfg(feature = "cuda")] + (Config::JinaV4(config), Device::Cuda(_)) => { + if dtype != DType::F16 + || !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1")) + || &std::env::var("USE_FLASH_ATTENTION") + .unwrap_or("True".to_string()) + .to_lowercase() + != "true" + { + return Err(BackendError::Start( + "Jina v4 is only supported on Cuda devices in fp16 with flash attention v2 enabled".to_string(), + )); + } + + let lora = load_jina_v4_lora(model_path, device, dtype, &config); + tracing::info!("Starting Jina v4 model on {:?}", device); + Ok(Box::new( + FlashQwen2Model::load( + vb, + &config.text_config, + model_type, + lora, + true, + ) + .s()?, )) } #[cfg(feature = "cuda")] diff --git a/backends/candle/src/models/flash_qwen2.rs b/backends/candle/src/models/flash_qwen2.rs index f00eb3f14..1e82714ad 100644 --- a/backends/candle/src/models/flash_qwen2.rs +++ b/backends/candle/src/models/flash_qwen2.rs @@ -6,6 +6,64 @@ use candle_nn::{Embedding, Module, VarBuilder}; use candle_rotary::apply_rotary_inplace; use text_embeddings_backend_core::{Batch, ModelType, Pool}; +#[derive(Debug, Clone)] +pub struct LoraWeights { + vb: VarBuilder, + task: String, + scale: f64, + r: usize, + prefix: String, +} + +impl LoraWeights { + pub fn new(vb: VarBuilder, task: String, r: usize, lora_alpha: f32, prefix: String) -> Self { + let scale = lora_alpha as f64 / r as f64; + Self { + vb, + task, + scale, + r, + prefix, + } + } + + fn apply(&self, weight: Tensor, base_path: &str) -> Result { + let lora_a_prefix = format!("{}.{}.lora_A.{}", self.prefix, base_path, self.task); + let lora_a_name = format!("{lora_a_prefix}.weight"); + if !self.vb.contains_tensor(&lora_a_name) { + return Ok(weight); + } + + let (out_dim, in_dim) = weight.dims2()?; + let lora_a = self + .vb + .pp(&lora_a_prefix) + .get((self.r, in_dim), "weight")?; + + let lora_b_prefix = format!("{}.{}.lora_B.{}", self.prefix, base_path, self.task); + let lora_b = self + .vb + .pp(&lora_b_prefix) + .get((out_dim, self.r), "weight")?; + + let delta = (lora_b.matmul(&lora_a)? * self.scale)?; + (&weight + &delta) + } +} + +fn load_weight_with_lora( + vb: &VarBuilder, + lora: Option<&LoraWeights>, + lora_path: &str, + shape: (usize, usize), +) -> Result { + let weight = vb.get(shape, "weight")?; + match lora { + Some(lora) => lora.apply(weight, lora_path), + None => Ok(weight), + } +} + struct Qwen2Attention { qkv_linear: Linear, o_proj: Linear, @@ -21,7 +79,12 @@ struct Qwen2Attention { } impl Qwen2Attention { - pub fn load(vb: VarBuilder, config: &Qwen2Config) -> Result { + pub fn load( + vb: VarBuilder, + config: &Qwen2Config, + layer_index: usize, + lora: Option<&LoraWeights>, + ) -> Result { if config.use_sliding_window { candle::bail!("Sliding window is not supported for Qwen2",); } @@ -31,20 +94,34 @@ impl Qwen2Attention { let num_key_value_heads = config.num_key_value_heads; let hidden_size = config.hidden_size; - let query_weight = vb.pp("q_proj").get((hidden_size, hidden_size), "weight")?; + let lora_base = format!("layers.{layer_index}.self_attn"); + + let q_vb = vb.pp("q_proj"); + let query_weight = load_weight_with_lora( + &q_vb, + lora, + &format!("{lora_base}.q_proj"), + (hidden_size, hidden_size), + )?; let query_bias = vb.pp("q_proj").get(hidden_size, "bias")?; - let key_weight = vb.pp("k_proj").get( + let k_vb = vb.pp("k_proj"); + let key_weight = load_weight_with_lora( + &k_vb, + lora, + &format!("{lora_base}.k_proj"), (num_key_value_heads * attention_head_size, hidden_size), - "weight", )?; let key_bias = vb .pp("k_proj") .get(num_key_value_heads * attention_head_size, "bias")?; - let value_weight = vb.pp("v_proj").get( + let v_vb = vb.pp("v_proj"); + let value_weight = load_weight_with_lora( + &v_vb, + lora, + &format!("{lora_base}.v_proj"), (num_key_value_heads * attention_head_size, hidden_size), - "weight", )?; let value_bias = vb .pp("v_proj") @@ -54,7 +131,13 @@ impl Qwen2Attention { let qkv_bias = Tensor::cat(&[&query_bias, &key_bias, &value_bias], 0)?; let qkv_linear = Linear::new(qkv_weight, Some(qkv_bias), None); - let o_proj_weight = vb.pp("o_proj").get((hidden_size, hidden_size), "weight")?; + let o_vb = vb.pp("o_proj"); + let o_proj_weight = load_weight_with_lora( + &o_vb, + lora, + &format!("{lora_base}.o_proj"), + (hidden_size, hidden_size), + )?; let o_proj = Linear::new(o_proj_weight, None, None); @@ -134,23 +217,42 @@ struct Qwen2MLP { } impl Qwen2MLP { - pub fn load(vb: VarBuilder, config: &Qwen2Config) -> Result { + pub fn load( + vb: VarBuilder, + config: &Qwen2Config, + layer_index: usize, + lora: Option<&LoraWeights>, + ) -> Result { let intermediate_size = config.intermediate_size; - let gate_proj_weight = vb - .pp("gate_proj") - .get((intermediate_size, config.hidden_size), "weight")?; + let lora_base = format!("layers.{layer_index}.mlp"); - let up_proj_weight = vb - .pp("up_proj") - .get((intermediate_size, config.hidden_size), "weight")?; + let gate_vb = vb.pp("gate_proj"); + let gate_proj_weight = load_weight_with_lora( + &gate_vb, + lora, + &format!("{lora_base}.gate_proj"), + (intermediate_size, config.hidden_size), + )?; + + let up_vb = vb.pp("up_proj"); + let up_proj_weight = load_weight_with_lora( + &up_vb, + lora, + &format!("{lora_base}.up_proj"), + (intermediate_size, config.hidden_size), + )?; let gate_up_proj_weight = Tensor::cat(&[&gate_proj_weight, &up_proj_weight], 0)?; let gate_up_proj = Linear::new(gate_up_proj_weight, None, None); - let down_proj_weight = vb - .pp("down_proj") - .get((config.hidden_size, intermediate_size), "weight")?; + let down_vb = vb.pp("down_proj"); + let down_proj_weight = load_weight_with_lora( + &down_vb, + lora, + &format!("{lora_base}.down_proj"), + (config.hidden_size, intermediate_size), + )?; let down_proj = Linear::new(down_proj_weight, None, None); Ok(Self { @@ -185,9 +287,14 @@ struct Qwen2Layer { } impl Qwen2Layer { - pub fn load(vb: VarBuilder, config: &Qwen2Config) -> Result { - let attention = Qwen2Attention::load(vb.pp("self_attn"), config)?; - let mlp = Qwen2MLP::load(vb.pp("mlp"), config)?; + pub fn load( + vb: VarBuilder, + config: &Qwen2Config, + layer_index: usize, + lora: Option<&LoraWeights>, + ) -> Result { + let attention = Qwen2Attention::load(vb.pp("self_attn"), config, layer_index, lora)?; + let mlp = Qwen2MLP::load(vb.pp("mlp"), config, layer_index, lora)?; let input_layer_norm = RMSNorm::load( vb.pp("input_layernorm"), @@ -240,13 +347,20 @@ pub struct FlashQwen2Model { cos_cache: Tensor, sin_cache: Tensor, pool: Pool, + normalize_embeddings: bool, pub device: Device, span: tracing::Span, } impl FlashQwen2Model { - pub fn load(vb: VarBuilder, config: &Qwen2Config, model_type: ModelType) -> Result { + pub fn load( + vb: VarBuilder, + config: &Qwen2Config, + model_type: ModelType, + lora: Option, + normalize_embeddings: bool, + ) -> Result { match vb.device() { Device::Cuda(_) => {} _ => candle::bail!("FlashQwen2 requires Cuda"), @@ -279,8 +393,9 @@ impl FlashQwen2Model { config.hidden_size, ); + let lora = lora.as_ref(); let layers = (0..config.num_hidden_layers) - .map(|index| Qwen2Layer::load(vb.pp(format!("layers.{index}")), config)) + .map(|index| Qwen2Layer::load(vb.pp(format!("layers.{index}")), config, index, lora)) .collect::>>()?; let norm = RMSNorm::load(vb.pp("norm"), config.hidden_size, config.rms_norm_eps)?; @@ -305,6 +420,7 @@ impl FlashQwen2Model { cos_cache, sin_cache, pool, + normalize_embeddings, device: vb.device().clone(), span: tracing::span!(tracing::Level::TRACE, "model"), }) @@ -425,6 +541,18 @@ impl FlashQwen2Model { None }; + let pooled_embeddings = if self.normalize_embeddings { + match pooled_embeddings { + Some(embeddings) => { + let norms = embeddings.sqr()?.sum_keepdim(1)?.sqrt()?; + Some(embeddings.broadcast_div(&norms)?) + } + None => None, + } + } else { + pooled_embeddings + }; + let raw_embeddings = if has_raw_requests { if batch_size > 1 && has_pooling_requests { // Create indexing vector for the embeddings diff --git a/backends/candle/src/models/mod.rs b/backends/candle/src/models/mod.rs index 424f4e984..7cc01d5d3 100644 --- a/backends/candle/src/models/mod.rs +++ b/backends/candle/src/models/mod.rs @@ -90,7 +90,7 @@ pub use flash_modernbert::FlashModernBertModel; pub use flash_nomic::FlashNomicBertModel; #[cfg(feature = "cuda")] -pub use flash_qwen2::FlashQwen2Model; +pub use flash_qwen2::{FlashQwen2Model, LoraWeights}; #[cfg(feature = "cuda")] pub use flash_qwen3::FlashQwen3Model; diff --git a/backends/src/lib.rs b/backends/src/lib.rs index 7bf20f163..0990bd8ab 100644 --- a/backends/src/lib.rs +++ b/backends/src/lib.rs @@ -4,7 +4,7 @@ use hf_hub::api::tokio::{ApiError, ApiRepo}; use rand::Rng; use std::cmp::{max, min}; use std::env; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::process::Command; use std::sync::Arc; use std::thread::JoinHandle; @@ -30,6 +30,27 @@ use text_embeddings_backend_ort::OrtBackend; #[cfg(feature = "python")] use text_embeddings_backend_python::PythonBackend; +#[cfg(feature = "candle")] +#[derive(Debug, Deserialize)] +struct MinimalConfig { + #[serde(default)] + architectures: Vec, +} + +#[cfg(feature = "candle")] +fn is_jina_v4_model(model_path: &Path) -> bool { + let cfg_path = model_path.join("config.json"); + if let Ok(cfg) = std::fs::read_to_string(cfg_path) { + if let Ok(cfg) = serde_json::from_str::(&cfg) { + return cfg + .architectures + .iter() + .any(|arch| arch == "JinaEmbeddingsV4Model"); + } + } + false +} + fn powers_of_two(max_value: usize) -> Vec { let mut result = Vec::new(); let mut power: usize = 1; @@ -365,8 +386,12 @@ async fn init_backend( ) -> Result, BackendError> { let mut backend_start_failed = false; let api_repo = api_repo.map(Arc::new); + #[cfg(feature = "candle")] + let jina_v4_model = is_jina_v4_model(&model_path); + #[cfg(not(feature = "candle"))] + let jina_v4_model = false; - if cfg!(feature = "ort") { + if cfg!(feature = "ort") && !jina_v4_model { #[cfg(feature = "ort")] { if let Some(api_repo) = api_repo.as_ref() { @@ -420,6 +445,12 @@ async fn init_backend( } tracing::info!("Model weights downloaded in {:?}", start.elapsed()); + + if jina_v4_model { + if let Err(err) = download_jina_v4_adapters(api_repo.clone()).await { + tracing::warn!("Failed to download Jina v4 adapters: {err}"); + } + } } } @@ -633,6 +664,14 @@ async fn download_safetensors(api: Arc) -> Result, ApiErro Ok(safetensors_files) } +async fn download_jina_v4_adapters(api: Arc) -> Result<(), ApiError> { + tracing::info!("Downloading `adapters/adapter_config.json`"); + let _ = api.get("adapters/adapter_config.json").await?; + tracing::info!("Downloading `adapters/adapter_model.safetensors`"); + let _ = api.get("adapters/adapter_model.safetensors").await?; + Ok(()) +} + #[cfg(feature = "ort")] async fn download_onnx(api: &ApiRepo) -> Result, ApiError> { let mut model_files: Vec = Vec::new(); diff --git a/router/src/lib.rs b/router/src/lib.rs index 1ad777ac0..c61088c33 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -114,12 +114,14 @@ pub async fn run( text_embeddings_backend::ModelType::Classifier => { let id2label = config .id2label + .clone() .context("`config.json` does not contain `id2label`")?; let n_classes = id2label.len(); let classifier_model = ClassifierModel { id2label, label2id: config .label2id + .clone() .context("`config.json` does not contain `label2id`")?, }; if n_classes > 1 { @@ -144,7 +146,7 @@ pub async fn run( // Old Qwen2 repos updates the post processor manually instead of into the tokenizer.json. // Newer ones (https://huggingface.co/jinaai/jina-code-embeddings-0.5b/tree/main) have it in the tokenizer.json. This is to support both cases. // https://huggingface.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct/blob/main/tokenization_qwen.py#L246 - if config.model_type == "qwen2" + if config.model_type() == Some("qwen2") && config .auto_map .as_ref() @@ -172,10 +174,10 @@ pub async fn run( } // Position IDs offset. Used for Roberta and camembert. - let position_offset = if &config.model_type == "xlm-roberta" - || &config.model_type == "camembert" - || &config.model_type == "roberta" - { + let position_offset = if matches!( + config.model_type(), + Some("xlm-roberta") | Some("camembert") | Some("roberta") + ) { config.pad_token_id + 1 } else { 0 @@ -263,7 +265,7 @@ pub async fn run( // NOTE: `gemma3_text` won't support Float16 but only Float32, given that with `candle-cuda` // feature, the default `Dtype::Float16` this overrides that to prevent issues when running a // `gemma3_text` model without specifying a `--dtype` - let dtype = if dtype.is_none() && config.model_type == "gemma3_text" { + let dtype = if dtype.is_none() && config.model_type() == Some("gemma3_text") { DType::Float32 } else { dtype.unwrap_or_default() @@ -415,6 +417,15 @@ fn get_backend_model_type( } } + if config + .architectures + .iter() + .any(|arch| arch == "JinaEmbeddingsV4Model") + { + let pool = pooling.unwrap_or(text_embeddings_backend::Pool::Mean); + return Ok(text_embeddings_backend::ModelType::Embedding(pool)); + } + if Some(text_embeddings_backend::Pool::Splade) == pooling { return Err(anyhow!( "Splade pooling is not supported: model is not a ForMaskedLM model" @@ -435,7 +446,10 @@ fn get_backend_model_type( Pool::try_from(config)? } Err(err) => { - if !config.model_type.to_lowercase().contains("bert") { + let is_bert = config + .model_type() + .is_some_and(|t| t.to_lowercase().contains("bert")); + if !is_bert { return Err(err).context("The `--pooling` arg is not set and we could not find a pooling configuration (`1_Pooling/config.json`) for this model."); } tracing::warn!("The `--pooling` arg is not set and we could not find a pooling configuration (`1_Pooling/config.json`) for this model but the model is a BERT variant. Defaulting to `CLS` pooling."); @@ -451,16 +465,33 @@ fn get_backend_model_type( #[derive(Debug, Deserialize)] pub struct ModelConfig { pub architectures: Vec, - pub model_type: String, + pub model_type: Option, #[serde(alias = "n_positions")] pub max_position_embeddings: usize, #[serde(default)] pub pad_token_id: usize, + #[serde(default)] + pub text_config: Option, pub id2label: Option>, pub label2id: Option>, pub auto_map: Option>, } +impl ModelConfig { + fn model_type(&self) -> Option<&str> { + self.model_type + .as_deref() + .or_else(|| self.text_config.as_ref()?.model_type.as_deref()) + } +} + +#[derive(Debug, Deserialize)] +pub struct TextConfig { + pub model_type: Option, + #[serde(default)] + pub sliding_window: Option, +} + #[derive(Debug, Clone, PartialEq, Deserialize)] pub struct PoolConfig { pooling_mode_cls_token: bool, From 79428a476e27b6ee51bbd4652a81d2c16e69cd04 Mon Sep 17 00:00:00 2001 From: KSW-KSM Date: Mon, 29 Dec 2025 14:27:40 +0900 Subject: [PATCH 2/3] chore: make router jina v4 handling conservative --- router/src/lib.rs | 37 +++++++++---------------------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/router/src/lib.rs b/router/src/lib.rs index c61088c33..2bc9ae28f 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -146,7 +146,7 @@ pub async fn run( // Old Qwen2 repos updates the post processor manually instead of into the tokenizer.json. // Newer ones (https://huggingface.co/jinaai/jina-code-embeddings-0.5b/tree/main) have it in the tokenizer.json. This is to support both cases. // https://huggingface.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct/blob/main/tokenization_qwen.py#L246 - if config.model_type() == Some("qwen2") + if config.model_type == "qwen2" && config .auto_map .as_ref() @@ -174,10 +174,10 @@ pub async fn run( } // Position IDs offset. Used for Roberta and camembert. - let position_offset = if matches!( - config.model_type(), - Some("xlm-roberta") | Some("camembert") | Some("roberta") - ) { + let position_offset = if &config.model_type == "xlm-roberta" + || &config.model_type == "camembert" + || &config.model_type == "roberta" + { config.pad_token_id + 1 } else { 0 @@ -265,7 +265,7 @@ pub async fn run( // NOTE: `gemma3_text` won't support Float16 but only Float32, given that with `candle-cuda` // feature, the default `Dtype::Float16` this overrides that to prevent issues when running a // `gemma3_text` model without specifying a `--dtype` - let dtype = if dtype.is_none() && config.model_type() == Some("gemma3_text") { + let dtype = if dtype.is_none() && config.model_type == "gemma3_text" { DType::Float32 } else { dtype.unwrap_or_default() @@ -446,10 +446,7 @@ fn get_backend_model_type( Pool::try_from(config)? } Err(err) => { - let is_bert = config - .model_type() - .is_some_and(|t| t.to_lowercase().contains("bert")); - if !is_bert { + if !config.model_type.to_lowercase().contains("bert") { return Err(err).context("The `--pooling` arg is not set and we could not find a pooling configuration (`1_Pooling/config.json`) for this model."); } tracing::warn!("The `--pooling` arg is not set and we could not find a pooling configuration (`1_Pooling/config.json`) for this model but the model is a BERT variant. Defaulting to `CLS` pooling."); @@ -465,33 +462,17 @@ fn get_backend_model_type( #[derive(Debug, Deserialize)] pub struct ModelConfig { pub architectures: Vec, - pub model_type: Option, + #[serde(default)] + pub model_type: String, #[serde(alias = "n_positions")] pub max_position_embeddings: usize, #[serde(default)] pub pad_token_id: usize, - #[serde(default)] - pub text_config: Option, pub id2label: Option>, pub label2id: Option>, pub auto_map: Option>, } -impl ModelConfig { - fn model_type(&self) -> Option<&str> { - self.model_type - .as_deref() - .or_else(|| self.text_config.as_ref()?.model_type.as_deref()) - } -} - -#[derive(Debug, Deserialize)] -pub struct TextConfig { - pub model_type: Option, - #[serde(default)] - pub sliding_window: Option, -} - #[derive(Debug, Clone, PartialEq, Deserialize)] pub struct PoolConfig { pooling_mode_cls_token: bool, From ce51bc3979bc153be60118454d12d999a453ad0c Mon Sep 17 00:00:00 2001 From: KSW-KSM Date: Mon, 29 Dec 2025 14:30:32 +0900 Subject: [PATCH 3/3] chore: restore classifier config ownership --- router/src/lib.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/router/src/lib.rs b/router/src/lib.rs index 2bc9ae28f..ac9e4a041 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -114,14 +114,12 @@ pub async fn run( text_embeddings_backend::ModelType::Classifier => { let id2label = config .id2label - .clone() .context("`config.json` does not contain `id2label`")?; let n_classes = id2label.len(); let classifier_model = ClassifierModel { id2label, label2id: config .label2id - .clone() .context("`config.json` does not contain `label2id`")?, }; if n_classes > 1 {