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..ac9e4a041 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -415,6 +415,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" @@ -451,6 +460,7 @@ fn get_backend_model_type( #[derive(Debug, Deserialize)] pub struct ModelConfig { pub architectures: Vec, + #[serde(default)] pub model_type: String, #[serde(alias = "n_positions")] pub max_position_embeddings: usize,