Qwen_reranker

python/python/embed_anything/_embed_anything.pyi

@@ -431,26 +431,55 @@ 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(
         self, query: list[str], documents: list[str], top_k: int
     ) -> 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.

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<Self> {
         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::<Vec<_>>())
     }
+
+    #[pyo3(signature = (query, documents, batch_size))]
+    pub fn compute_scores(
+        &self,
+        query: Vec<String>,
+        documents: Vec<String>,
+        batch_size: usize,
+    ) -> PyResult<Vec<Vec<f32>>> {
+        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)
+    }
 }

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();
 

rust/src/reranker/mod.rs

@@ -1 +1,2 @@
 pub mod model;
+pub mod qwen3;

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<Session>,
+    model_type: Option<String>,
     tokenizer: Tokenizer,
 }
 
 impl Reranker {
-    pub fn new(model_id: &str, revision: Option<&str>, dtype: Dtype) -> Result<Self, E> {
-        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<Self, E> {
+        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<Vec<Vec<f32>>, 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<think>\n\n</think>\n\n";
+
         let pairs = queries
             .iter()
             .flat_map(|query| documents.iter().map(move |doc| (*query, *doc)))
             .collect::<Vec<_>>();
+
+        let pairs = if is_qwen3 {
+            pairs.iter().map(|(query, doc)| (format!("{}{}", prefix, query), format!("{}{}", doc, suffix))).collect::<Vec<_>>()
+        } else {
+            pairs.iter().map(|(query, doc)| (query.to_string(), doc.to_string())).collect::<Vec<_>>()
+        };
+
+
         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::<f32>()?
-                .to_owned()
-                .into_dimensionality::<ndarray::Ix2>()?;
-            scores.extend(
-                logits
-                    .outer_iter()
-                    .flat_map(|row| row.to_vec())
-                    .collect::<Vec<_>>(),
-            );
+            let model_inputs = model_guard
+                .inputs
+                .iter()
+                .map(|i| i.name.clone())
+                .collect::<Vec<_>>();
+
+            let seq_len = input_ids.shape()[1];
+            let position_ids_vec: Vec<i64> = (0..seq_len as i64)
+                .collect::<Vec<i64>>()
+                .repeat(input_ids.shape()[0]);
+            let position_ids =
+                Array2::<i64>::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::<f32>()?
+                    .to_owned()
+                    .into_dimensionality::<ndarray::Ix2>()?;
+                scores.extend(
+                    logits
+                        .outer_iter()
+                        .flat_map(|row| row.to_vec())
+                        .collect::<Vec<_>>(),
+                );
+            }
         }
+
         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::<f32>()?)
+
+        if is_qwen3 {
+            Ok(scores_tensor.to_vec2::<f32>()?)
+        } else {
+            let sigmoid_scores = candle_nn::ops::sigmoid(&scores_tensor)?;
+            Ok(sigmoid_scores.to_vec2::<f32>()?)
+        }
     }
 
     pub fn rerank(
@@ -197,7 +280,7 @@ impl Reranker {
         Ok(reranker_results)
     }
 
-    pub fn tokenize_batch_ndarray(&self, pairs: &[(&str, &str)]) -> anyhow::Result<Array2<i64>> {
+    pub fn tokenize_batch_ndarray(&self, pairs: &[(String, String)]) -> anyhow::Result<Array2<i64>> {
         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<Array2<i64>> {
         let attention_mask = self
             .tokenizer

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<i64>,
+    attention_mask: Array2<i64>,
+    position_ids: Array2<i64>,
+    false_token_id: u32,
+    true_token_id: u32,
+) -> Result<Vec<f32>, 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::<f16>()?.to_owned().into_dimensionality::<ndarray::Ix3>()?.map(|&x| x.to_f32()).into_dimensionality::<ndarray::Ix3>()?;
+
+    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::<f32>()?)
+}