add qwen3-reranker

src/models/common/mod.rs

@@ -5,6 +5,7 @@ pub mod generate;
 pub mod gguf;
 pub mod model_mapping;
 pub mod modules;
+pub mod reranker;
 
 /// 多模态模型的特征数据
 /// 每个模型数据不一样

src/models/common/model_mapping.rs

@@ -44,6 +44,12 @@ pub enum WhichModel {
     Qwen3Embedding4B,
     #[value(name = "Qwen/Qwen3-Embedding-8B")]
     Qwen3Embedding8B,
+    #[value(name = "Qwen/Qwen3-Reranker-0.6B")]
+    Qwen3Reranker0_6B,
+    #[value(name = "Qwen/Qwen3-Reranker-4B")]
+    Qwen3Reranker4B,
+    #[value(name = "Qwen/Qwen3-Reranker-8B")]
+    Qwen3Reranker8B,
     #[value(name = "Qwen/Qwen3-VL-2B-Instruct")]
     Qwen3VL2B,
     #[value(name = "Qwen/Qwen3-VL-4B-Instruct")]
@@ -165,6 +171,9 @@ impl WhichModel {
             | WhichModel::Qwen3Embedding4B
             | WhichModel::Qwen3Embedding8B
             | WhichModel::AllMiniLML6V2 => "embedding",
+            WhichModel::Qwen3Reranker0_6B
+            | WhichModel::Qwen3Reranker4B
+            | WhichModel::Qwen3Reranker8B => "reranker",
         }
     }
 }

src/models/common/modules.rs

@@ -1453,6 +1453,16 @@ pub fn cosine_similarity(query_vector: &Tensor, matrix: &Tensor) -> Result<Tenso
     Ok(similarity)
 }
 
+pub fn cosine_similarity_no_l2(query_vector: &Tensor, matrix: &Tensor) -> Result<Tensor> {
+    // query_vector: (n, dim)
+    // matrix: (m, dim)
+    // return (n, m)
+    let similarity = query_vector
+        .matmul(&matrix.transpose(D::Minus1, D::Minus2)?)?
+        .squeeze(D::Minus1)?;
+    Ok(similarity)
+}
+
 pub fn quick_gelu(xs: &Tensor) -> Result<Tensor> {
     let x = xs.affine(1.702, 0.0)?;
     let x = sigmoid(&x)?;

src/models/common/reranker.rs

@@ -0,0 +1,25 @@
+use anyhow::Result;
+use candle_core::Tensor;
+
+use crate::models::common::modules::{cosine_similarity, cosine_similarity_no_l2};
+pub trait TextRerank {
+    fn rerank(&mut self, query: &str, documents: &[String]) -> Result<Vec<f32>>;
+}
+
+pub enum RerankerSimilarity {
+    Cosine,
+}
+
+impl RerankerSimilarity {
+    pub fn similar(&self, query_vector: &Tensor, matrix: &Tensor, need_l2: bool) -> Result<Tensor> {
+        match self {
+            RerankerSimilarity::Cosine => {
+                if need_l2 {
+                    cosine_similarity(query_vector, matrix)
+                } else {
+                    cosine_similarity_no_l2(query_vector, matrix)
+                }
+            }
+        }
+    }
+}