From 28f121999a6fe56bb9502ee05bfde3683ec4eb13 Mon Sep 17 00:00:00 2001 From: Akshay Ballal Date: Sat, 26 Jul 2025 18:59:41 +0200 Subject: [PATCH 1/3] Add compute_scores method to Reranker class and update from_pretrained signature - Introduced a new `compute_scores` method in the `Reranker` class to calculate scores for given queries and documents. - Updated the `from_pretrained` method signature to include an optional `path_in_repo` parameter for model loading. - Added a new `qwen3` module for handling specific model computations. - Enhanced the `reranker` module to support dynamic model type handling during score computation. - Updated example usage in `reranker.rs` to demonstrate the new functionality. --- .../python/embed_anything/_embed_anything.pyi | 6 + python/src/models/reranker.rs | 21 +- rust/examples/reranker.rs | 1 + rust/src/reranker/mod.rs | 1 + rust/src/reranker/model.rs | 193 +++++++++++++++--- rust/src/reranker/qwen3.rs | 37 ++++ 6 files changed, 225 insertions(+), 34 deletions(-) create mode 100644 rust/src/reranker/qwen3.rs diff --git a/python/python/embed_anything/_embed_anything.pyi b/python/python/embed_anything/_embed_anything.pyi index 05b6762c..5a2c0bc0 100644 --- a/python/python/embed_anything/_embed_anything.pyi +++ b/python/python/embed_anything/_embed_anything.pyi @@ -451,6 +451,12 @@ class Reranker: Reranks the given documents for the query and returns a list of RerankerResult objects. """ + def compute_scores( + self, query: list[str], documents: list[str], batch_size: int + ) -> list[list[float]]: + """ + Computes the scores for the given query and documents. + """ class Dtype(Enum): """ Represents the data type of the model. diff --git a/python/src/models/reranker.rs b/python/src/models/reranker.rs index f1b4193e..9f85116c 100644 --- a/python/src/models/reranker.rs +++ b/python/src/models/reranker.rs @@ -105,11 +105,12 @@ impl RerankerResult { #[pymethods] impl Reranker { #[staticmethod] - #[pyo3(signature = (model_id, revision=None, dtype=None))] + #[pyo3(signature = (model_id, revision=None, dtype=None, path_in_repo=None))] pub fn from_pretrained( model_id: &str, revision: Option<&str>, dtype: Option<&Dtype>, + path_in_repo: Option<&str>, ) -> PyResult { let dtype = match dtype { Some(Dtype::F16) => embed_anything::Dtype::F16, @@ -120,7 +121,7 @@ impl Reranker { Some(Dtype::F32) => embed_anything::Dtype::F32, _ => embed_anything::Dtype::F32, }; - let model = embed_anything::reranker::model::Reranker::new(model_id, revision, dtype) + let model = embed_anything::reranker::model::Reranker::new(model_id, revision, dtype, path_in_repo) .map_err(|e| PyValueError::new_err(e.to_string()))?; Ok(Self { model }) } @@ -143,4 +144,20 @@ impl Reranker { .map(|r| RerankerResult { inner: r }) .collect::>()) } + + #[pyo3(signature = (query, documents, batch_size))] + pub fn compute_scores( + &self, + query: Vec, + documents: Vec, + batch_size: usize, + ) -> PyResult>> { + let query_refs: Vec<&str> = query.iter().map(|s| s.as_str()).collect(); + let document_refs: Vec<&str> = documents.iter().map(|s| s.as_str()).collect(); + let scores = self + .model + .compute_scores(query_refs, document_refs, batch_size) + .unwrap(); + Ok(scores) + } } diff --git a/rust/examples/reranker.rs b/rust/examples/reranker.rs index b4e51ef9..fd3e8e5e 100644 --- a/rust/examples/reranker.rs +++ b/rust/examples/reranker.rs @@ -6,6 +6,7 @@ fn main() { "jinaai/jina-reranker-v2-base-multilingual", None, Dtype::F16, + Some("onnx"), ) .unwrap(); diff --git a/rust/src/reranker/mod.rs b/rust/src/reranker/mod.rs index 65880be0..18a2a381 100644 --- a/rust/src/reranker/mod.rs +++ b/rust/src/reranker/mod.rs @@ -1 +1,2 @@ pub mod model; +pub mod qwen3; \ No newline at end of file diff --git a/rust/src/reranker/model.rs b/rust/src/reranker/model.rs index cd857b41..3b27d5eb 100644 --- a/rust/src/reranker/model.rs +++ b/rust/src/reranker/model.rs @@ -10,8 +10,8 @@ use ort::{ }; use tokenizers::{PaddingParams, Tokenizer, TruncationParams}; -use crate::embeddings::local::bert::TokenizerConfig; use crate::Dtype; +use crate::{embeddings::local::bert::TokenizerConfig, reranker::qwen3}; use serde::Serialize; #[derive(Debug, Serialize)] @@ -29,12 +29,18 @@ pub struct DocumentRank { pub struct Reranker { model: RwLock, + model_type: Option, tokenizer: Tokenizer, } impl Reranker { - pub fn new(model_id: &str, revision: Option<&str>, dtype: Dtype) -> Result { - let (_, tokenizer_filename, weights_filename, tokenizer_config_filename) = { + pub fn new( + model_id: &str, + revision: Option<&str>, + dtype: Dtype, + path_in_repo: Option<&str>, + ) -> Result { + let (config_filename, tokenizer_filename, weights_filename, tokenizer_config_filename) = { let api = Api::new().unwrap(); let api = match revision { Some(rev) => api.repo(Repo::with_revision( @@ -50,16 +56,23 @@ impl Reranker { let config = api.get("config.json")?; let tokenizer = api.get("tokenizer.json")?; let tokenizer_config = api.get("tokenizer_config.json")?; + + let mut path_in_repo = path_in_repo.unwrap_or_default().to_string(); + if !path_in_repo.is_empty() { + path_in_repo.push('/'); + } let weights = match dtype { - Dtype::Q4F16 => api.get("onnx/model_q4f16.onnx")?, - Dtype::F16 => api.get("onnx/model_fp16.onnx")?, - Dtype::INT8 => api.get("onnx/model_int8.onnx")?, - Dtype::Q4 => api.get("onnx/model_q4.onnx")?, - Dtype::UINT8 => api.get("onnx/model_uint8.onnx")?, - Dtype::BNB4 => api.get("onnx/model_bnb4.onnx")?, - Dtype::F32 => api.get("onnx/model.onnx")?, - Dtype::BF16 => api.get("onnx/model_bf16.onnx")?, - Dtype::QUANTIZED => api.get("onnx/model_quantized.onnx")?, + Dtype::Q4F16 => api.get(format!("{}model_q4f16.onnx", path_in_repo).as_str())?, + Dtype::F16 => api.get(format!("{}model_fp16.onnx", path_in_repo).as_str())?, + Dtype::INT8 => api.get(format!("{}model_int8.onnx", path_in_repo).as_str())?, + Dtype::Q4 => api.get(format!("{}model_q4.onnx", path_in_repo).as_str())?, + Dtype::UINT8 => api.get(format!("{}model_uint8.onnx", path_in_repo).as_str())?, + Dtype::BNB4 => api.get(format!("{}model_bnb4.onnx", path_in_repo).as_str())?, + Dtype::F32 => api.get(format!("{}model.onnx", path_in_repo).as_str())?, + Dtype::BF16 => api.get(format!("{}model_bf16.onnx", path_in_repo).as_str())?, + Dtype::QUANTIZED => { + api.get(format!("{}model_quantized.onnx", path_in_repo).as_str())? + } }; (config, tokenizer, weights, tokenizer_config) }; @@ -76,6 +89,13 @@ impl Reranker { (None, None) => 128, }; + let config_file = std::fs::read_to_string(config_filename)?; + let config: serde_json::Value = serde_json::from_str(&config_file)?; + let model_type = config + .get("model_type") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let mut tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?; let pp = PaddingParams { @@ -120,6 +140,7 @@ impl Reranker { Ok(Reranker { model: RwLock::new(model), + model_type, tokenizer, }) } @@ -130,37 +151,145 @@ impl Reranker { documents: Vec<&str>, batch_size: usize, ) -> Result>, E> { + + + let prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be yes or no.<|im_end|>\n<|im_start|>user\n"; + let suffix = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; + + let pairs = queries .iter() .flat_map(|query| documents.iter().map(move |doc| (*query, *doc))) .collect::>(); + + let pairs = match &self.model_type { + Some(model_type) => { + if model_type == "qwen3" { + pairs.iter().map(|(query, doc)| (format!("{}{}", prefix, query), format!("{}{}", doc, suffix))).collect::>() + } else { + pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::>() + } + } + None => { + pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::>() + } + }; + + let mut scores = Vec::with_capacity(pairs.len()); let mut model_guard = self.model.write().unwrap(); + + let false_token_id = self + .tokenizer + .token_to_id("no") + .ok_or(E::msg("no token found"))?; + let true_token_id = self + .tokenizer + .token_to_id("yes") + .ok_or(E::msg("yes token found"))?; + for pair in pairs.chunks(batch_size) { let input_ids = self.tokenize_batch_ndarray(pair)?; let attention_mask = self.get_attention_mask_ndarray(pair)?; let input_ids_tensor = ort::value::TensorRef::from_array_view(&input_ids)?; let attention_mask_tensor = ort::value::TensorRef::from_array_view(&attention_mask)?; - let outputs = model_guard.run(ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor])?; - let logits = outputs["logits".to_string()] - .try_extract_array::()? - .to_owned() - .into_dimensionality::()?; - scores.extend( - logits - .outer_iter() - .flat_map(|row| row.to_vec()) - .collect::>(), - ); + let model_inputs = model_guard + .inputs + .iter() + .map(|i| i.name.clone()) + .collect::>(); + + let seq_len = input_ids.shape()[1]; + let position_ids_vec: Vec = (0..seq_len as i64) + .collect::>() + .repeat(input_ids.shape()[0]); + let position_ids = + Array2::::from_shape_vec((input_ids.shape()[0], seq_len), position_ids_vec)?; + let position_ids_tensor = ort::value::TensorRef::from_array_view(&position_ids)?; + + let inputs = if model_inputs.contains(&"position_ids".to_string()) { + ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor, "position_ids" => position_ids_tensor] + } else { + ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor] + }; + + + match &self.model_type { + Some(model_type) => { + if model_type == "qwen3" { + let logits = qwen3::compute_scores( + &mut model_guard, + input_ids, + attention_mask, + position_ids, + false_token_id, + true_token_id, + )?; + scores.extend(logits); + } else { + let outputs = model_guard.run(inputs)?; + let logits = outputs["logits".to_string()] + .try_extract_array::()? + .to_owned() + .into_dimensionality::()?; + scores.extend( + logits + .outer_iter() + .flat_map(|row| row.to_vec()) + .collect::>(), + ); + + } + } + None => { + let outputs = model_guard.run(inputs)?; + let logits = outputs["logits".to_string()] + .try_extract_array::()? + .to_owned() + .into_dimensionality::()?; + + scores.extend( + logits + .outer_iter() + .flat_map(|row| row.to_vec()) + .collect::>(), + ); + + } + }; + } + + match &self.model_type { + Some(model_type) => { + if model_type == "qwen3" { + let scores_tensor = Tensor::from_vec( + scores.clone(), + (queries.len(), documents.len()), + &Device::Cpu, + )?; + Ok(scores_tensor.to_vec2::()?) + } else { + let scores_tensor = Tensor::from_vec( + scores.clone(), + (queries.len(), documents.len()), + &Device::Cpu, + )?; + let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?; + Ok(sigmoid_scores.to_vec2::()?) + } + } + None => { + let scores_tensor = Tensor::from_vec( + scores.clone(), + (queries.len(), documents.len()), + &Device::Cpu, + )?; + let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?; + Ok(sigmoid_scores.to_vec2::()?) + } } - let scores_tensor = Tensor::from_vec( - scores.clone(), - (queries.len(), documents.len()), - &Device::Cpu, - )?; - let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor).unwrap(); - Ok(sigmoid_scores.to_vec2::()?) + } pub fn rerank( @@ -197,7 +326,7 @@ impl Reranker { Ok(reranker_results) } - pub fn tokenize_batch_ndarray(&self, pairs: &[(&str, &str)]) -> anyhow::Result> { + pub fn tokenize_batch_ndarray(&self, pairs: &[(String, String)]) -> anyhow::Result> { let token_ids = self .tokenizer .encode_batch(pairs.to_vec(), true) @@ -222,7 +351,7 @@ impl Reranker { pub fn get_attention_mask_ndarray( &self, - pairs: &[(&str, &str)], + pairs: &[(String, String)], ) -> anyhow::Result> { let attention_mask = self .tokenizer diff --git a/rust/src/reranker/qwen3.rs b/rust/src/reranker/qwen3.rs new file mode 100644 index 00000000..6d7e06a0 --- /dev/null +++ b/rust/src/reranker/qwen3.rs @@ -0,0 +1,37 @@ +use anyhow::{Error as E, Result}; +use candle_core::{Device, IndexOp, Tensor}; +use half::f16; +use ndarray::{s, Array2}; +use ort::session::Session; + +pub fn compute_scores( + model: &mut Session, + input_ids: Array2, + attention_mask: Array2, + position_ids: Array2, + false_token_id: u32, + true_token_id: u32, +) -> Result, E> { + let input_ids_tensor = ort::value::TensorRef::from_array_view(&input_ids)?; + let attention_mask_tensor = ort::value::TensorRef::from_array_view(&attention_mask)?; + let position_ids_tensor = ort::value::TensorRef::from_array_view(&position_ids)?; + + let inputs = + ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor, "position_ids" => position_ids_tensor]; + + + let outputs = model.run(inputs)?; + let logits = outputs["logits".to_string()].try_extract_array::()?.to_owned().into_dimensionality::()?.map(|&x| x.to_f32()).into_dimensionality::()?; + + let logits = logits.slice(s![.., input_ids.shape()[1]-1, ..]).to_owned(); + let false_logits = logits.slice(s![.., false_token_id as usize]).to_owned(); + let true_logits = logits.slice(s![.., true_token_id as usize]).to_owned(); + + let false_logits_tensor = Tensor::from_vec(false_logits.to_vec(), input_ids.shape()[0], &Device::Cpu)?; + let true_logits_tensor = Tensor::from_vec(true_logits.to_vec(), input_ids.shape()[0], &Device::Cpu)?; + + let logits_tensor = Tensor::stack(&[&false_logits_tensor, &true_logits_tensor], 1)?; + let probs = candle_nn::ops::softmax(&logits_tensor, 1)?; + let scores = probs.i((.., 1))?; + Ok(scores.to_vec1::()?) +} From 2f0e169262e1a3dd5982ba99614171c23da55489 Mon Sep 17 00:00:00 2001 From: Akshay Ballal Date: Sat, 26 Jul 2025 19:05:44 +0200 Subject: [PATCH 2/3] Update reranker model to use Qwen3 and streamline score computation - Changed the model in `reranker.rs` to "zhiqing/Qwen3-Reranker-0.6B-ONNX" for improved performance. - Refactored score computation in `model.rs` to simplify handling of Qwen3 model type. - Introduced a check for model type at the beginning of the score computation process to enhance clarity and maintainability. - Updated logic to conditionally format query and document pairs based on the model type, improving code readability. --- rust/examples/reranker.rs | 6 +- rust/src/reranker/model.rs | 126 ++++++++++++------------------------- 2 files changed, 43 insertions(+), 89 deletions(-) diff --git a/rust/examples/reranker.rs b/rust/examples/reranker.rs index fd3e8e5e..c589e9af 100644 --- a/rust/examples/reranker.rs +++ b/rust/examples/reranker.rs @@ -3,10 +3,10 @@ fn main() { use embed_anything::Dtype; let reranker = embed_anything::reranker::model::Reranker::new( - "jinaai/jina-reranker-v2-base-multilingual", + "zhiqing/Qwen3-Reranker-0.6B-ONNX", + None, + Dtype::F32, None, - Dtype::F16, - Some("onnx"), ) .unwrap(); diff --git a/rust/src/reranker/model.rs b/rust/src/reranker/model.rs index 3b27d5eb..88bb9b37 100644 --- a/rust/src/reranker/model.rs +++ b/rust/src/reranker/model.rs @@ -152,34 +152,27 @@ impl Reranker { batch_size: usize, ) -> Result>, E> { + // Check model type once at the beginning + let is_qwen3 = self.model_type.as_ref().map_or(false, |t| t == "qwen3"); let prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be yes or no.<|im_end|>\n<|im_start|>user\n"; let suffix = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; - let pairs = queries .iter() .flat_map(|query| documents.iter().map(move |doc| (*query, *doc))) .collect::>(); - let pairs = match &self.model_type { - Some(model_type) => { - if model_type == "qwen3" { - pairs.iter().map(|(query, doc)| (format!("{}{}", prefix, query), format!("{}{}", doc, suffix))).collect::>() - } else { - pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::>() - } - } - None => { - pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::>() - } + let pairs = if is_qwen3 { + pairs.iter().map(|(query, doc)| (format!("{}{}", prefix, query), format!("{}{}", doc, suffix))).collect::>() + } else { + pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::>() }; let mut scores = Vec::with_capacity(pairs.len()); let mut model_guard = self.model.write().unwrap(); - let false_token_id = self .tokenizer .token_to_id("no") @@ -214,82 +207,43 @@ impl Reranker { ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor] }; - - match &self.model_type { - Some(model_type) => { - if model_type == "qwen3" { - let logits = qwen3::compute_scores( - &mut model_guard, - input_ids, - attention_mask, - position_ids, - false_token_id, - true_token_id, - )?; - scores.extend(logits); - } else { - let outputs = model_guard.run(inputs)?; - let logits = outputs["logits".to_string()] - .try_extract_array::()? - .to_owned() - .into_dimensionality::()?; - scores.extend( - logits - .outer_iter() - .flat_map(|row| row.to_vec()) - .collect::>(), - ); - - } - } - None => { - let outputs = model_guard.run(inputs)?; - let logits = outputs["logits".to_string()] - .try_extract_array::()? - .to_owned() - .into_dimensionality::()?; - - scores.extend( - logits - .outer_iter() - .flat_map(|row| row.to_vec()) - .collect::>(), - ); - - } - }; - } - - match &self.model_type { - Some(model_type) => { - if model_type == "qwen3" { - let scores_tensor = Tensor::from_vec( - scores.clone(), - (queries.len(), documents.len()), - &Device::Cpu, - )?; - Ok(scores_tensor.to_vec2::()?) - } else { - let scores_tensor = Tensor::from_vec( - scores.clone(), - (queries.len(), documents.len()), - &Device::Cpu, - )?; - let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?; - Ok(sigmoid_scores.to_vec2::()?) - } - } - None => { - let scores_tensor = Tensor::from_vec( - scores.clone(), - (queries.len(), documents.len()), - &Device::Cpu, + if is_qwen3 { + let logits = qwen3::compute_scores( + &mut model_guard, + input_ids, + attention_mask, + position_ids, + false_token_id, + true_token_id, )?; - let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?; - Ok(sigmoid_scores.to_vec2::()?) + scores.extend(logits); + } else { + let outputs = model_guard.run(inputs)?; + let logits = outputs["logits".to_string()] + .try_extract_array::()? + .to_owned() + .into_dimensionality::()?; + scores.extend( + logits + .outer_iter() + .flat_map(|row| row.to_vec()) + .collect::>(), + ); } } - + + let scores_tensor = Tensor::from_vec( + scores.clone(), + (queries.len(), documents.len()), + &Device::Cpu, + )?; + + if is_qwen3 { + Ok(scores_tensor.to_vec2::()?) + } else { + let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?; + Ok(sigmoid_scores.to_vec2::()?) + } } pub fn rerank( From 837e2b1c4f89c67b75a9262df13ba073b2ebc3b7 Mon Sep 17 00:00:00 2001 From: Akshay Ballal Date: Sun, 27 Jul 2025 20:53:46 +0200 Subject: [PATCH 3/3] Update Reranker class to include optional path_in_repo parameter and improve model type check - Added an optional `path_in_repo` parameter to the `__init__` and `from_pretrained` methods of the `Reranker` class for enhanced model loading flexibility. - Refactored the model type check in `model.rs` to use `is_some_and` for improved clarity and performance. --- .../python/embed_anything/_embed_anything.pyi | 27 +++++++++++++++++-- rust/src/reranker/model.rs | 2 +- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/python/python/embed_anything/_embed_anything.pyi b/python/python/embed_anything/_embed_anything.pyi index 5a2c0bc0..e87afd9a 100644 --- a/python/python/embed_anything/_embed_anything.pyi +++ b/python/python/embed_anything/_embed_anything.pyi @@ -431,17 +431,24 @@ class Reranker: """ def __init__( - self, model_id: str, revision: str | None = None, dtype: Dtype | None = None + self, model_id: str, revision: str | None = None, dtype: Dtype | None = None, path_in_repo: str | None = None ): """ Initializes the Reranker object. """ def from_pretrained( - model_id: str, revision: str | None = None, dtype: Dtype | None = None + model_id: str, revision: str | None = None, dtype: Dtype | None = None, path_in_repo: str | None = None ) -> Reranker: """ Loads a pre-trained Reranker model from the Hugging Face model hub. + + Args: + model_id: The ID of the model from Hugging Face. + revision: The revision of the model. + dtype: The dtype of the model. + path_in_repo: The path to the model in the repository. + """ def rerank( @@ -449,6 +456,14 @@ class Reranker: ) -> RerankerResult: """ Reranks the given documents for the query and returns a list of RerankerResult objects. + + Args: + query: The query to rerank. + documents: The list of documents to rerank. + top_k: The number of documents to return. + + Returns: + A list of RerankerResult objects. """ def compute_scores( @@ -456,6 +471,14 @@ class Reranker: ) -> list[list[float]]: """ Computes the scores for the given query and documents. + + Args: + query: The query to compute the scores for. + documents: The list of documents to compute the scores for. + batch_size: The batch size for processing the scores. + + Returns: + A list of scores for the given query and documents. """ class Dtype(Enum): """ diff --git a/rust/src/reranker/model.rs b/rust/src/reranker/model.rs index 88bb9b37..58ca6fa2 100644 --- a/rust/src/reranker/model.rs +++ b/rust/src/reranker/model.rs @@ -153,7 +153,7 @@ impl Reranker { ) -> Result>, E> { // Check model type once at the beginning - let is_qwen3 = self.model_type.as_ref().map_or(false, |t| t == "qwen3"); + let is_qwen3 = self.model_type.as_ref().is_some_and(|t| t == "qwen3"); let prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be yes or no.<|im_end|>\n<|im_start|>user\n"; let suffix = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n";