diff --git a/python/python/embed_anything/_embed_anything.pyi b/python/python/embed_anything/_embed_anything.pyi index 05b6762c..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,8 +456,30 @@ 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( + self, query: list[str], documents: list[str], batch_size: int + ) -> 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): """ 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..c589e9af 100644 --- a/rust/examples/reranker.rs +++ b/rust/examples/reranker.rs @@ -3,9 +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, ) .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..58ca6fa2 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,99 @@ impl Reranker { documents: Vec<&str>, batch_size: usize, ) -> Result>, E> { + + // Check model type once at the beginning + 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"; + let pairs = queries .iter() .flat_map(|query| documents.iter().map(move |doc| (*query, *doc))) .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") + .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] + }; + + if is_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::>(), + ); + } } + 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::()?) + + 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( @@ -197,7 +280,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 +305,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::()?) +}