Improve reranker API

Cargo.toml

@@ -27,6 +27,6 @@ candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.9.
 tokio = { version = "1", features = ["full"] }
 async-trait = "0.1"
 pin-project-lite = "0.2"
-hf-hub = "0.4.2"
+hf-hub = { version = "0.4.3", features = ["tokio"] }
 
 [dev-dependencies]

README.md

@@ -93,6 +93,18 @@ ModernBERT powers three specialized analysis tasks with shared architecture:
 
 [→ View on HuggingFace](https://huggingface.co/MoritzLaurer/ModernBERT-base-zeroshot-v2.0)
 
+#### **Reranking Pipeline**
+*Rank documents by relevance to a query*
+
+```markdown
+ Available Sizes:
+├── 0.6B
+├── 4B
+└── 8B
+```
+
+[→ View on HuggingFace](https://huggingface.co/collections/Qwen/qwen3-rerankers)
+
 ---
 
 ***Technical Note**: All ModernBERT pipelines share the same backbone architecture, loading task-specific finetuned weights as needed.*
@@ -395,6 +407,26 @@ fn main() -> Result<()> {
 }
 ```
 
+### Reranking (Qwen3)
+
+```rust
+use transformers::pipelines::reranker_pipeline::{RerankPipelineBuilder, Qwen3RerankSize};
+use anyhow::Result;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let pipeline = RerankPipelineBuilder::qwen3(Qwen3RerankSize::Size0_6B)
+        .cpu()
+        .build()
+        .await?;
+
+    let docs = ["Rust is a systems programming language.", "Cats are cute."];
+    let ranked = pipeline.rerank("Tell me about Rust", &docs).await?;
+    println!("Best match score: {:.4}", ranked[0].score);
+    Ok(())
+}
+```
+
 ## Future Plans
 
 - Add more model families and sizes

examples/embedding.rs

@@ -2,7 +2,7 @@ use transformers::pipelines::embedding_pipeline::*;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let pipeline = EmbeddingPipelineBuilder::qwen3(Qwen3Size::Size0_6B)
+    let pipeline = EmbeddingPipelineBuilder::qwen3(Qwen3EmbeddingSize::Size0_6B)
         .build()
         .await?;
 

examples/reranker/corpus/machine_learning_doc.txt

@@ -0,0 +1,11 @@
+Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns, make predictions, and improve their performance over time through experience.
+
+There are three main types of machine learning: supervised learning, unsupervised learning, and reinforcement learning. Supervised learning uses labeled training data to learn a mapping from inputs to outputs. Common supervised learning tasks include classification (predicting categories) and regression (predicting continuous values). Popular algorithms include linear regression, decision trees, random forests, and neural networks.
+
+Unsupervised learning finds hidden patterns in data without labeled examples. Clustering algorithms like k-means group similar data points together, while dimensionality reduction techniques like PCA help visualize high-dimensional data. Association rule learning discovers relationships between different variables.
+
+Reinforcement learning trains agents to make sequential decisions by learning from rewards and penalties. The agent explores an environment, takes actions, and receives feedback to maximize cumulative reward. This approach has been successful in game playing, robotics, and autonomous systems.
+
+Deep learning, a subset of machine learning, uses artificial neural networks with multiple layers to learn complex patterns. Convolutional neural networks excel at image recognition, recurrent neural networks handle sequential data like text and speech, and transformer architectures have revolutionized natural language processing.
+
+Key concepts in machine learning include feature engineering, cross-validation, overfitting, underfitting, bias-variance tradeoff, and model evaluation metrics. Data preprocessing, feature selection, and hyperparameter tuning are crucial steps in building effective machine learning models.

examples/reranker/corpus/math_doc.txt

@@ -0,0 +1,13 @@
+Mathematics is the study of numbers, shapes, patterns, and logical structures. It provides the foundation for understanding quantitative relationships and solving problems across many fields including science, engineering, economics, and computer science.
+
+Algebra deals with symbols and the rules for manipulating those symbols to solve equations and understand relationships between variables. Linear algebra focuses on vectors, matrices, and systems of linear equations, which are fundamental to computer graphics, machine learning, and quantum mechanics.
+
+Calculus studies continuous change through derivatives and integrals. Differential calculus examines rates of change and slopes of curves, while integral calculus deals with areas under curves and accumulation of quantities. Calculus is essential for physics, engineering, and optimization problems.
+
+Geometry explores the properties of shapes, sizes, and spatial relationships. Euclidean geometry deals with flat surfaces and familiar shapes like triangles and circles. Non-Euclidean geometries describe curved spaces and are important in relativity theory and modern physics.
+
+Statistics and probability theory help us understand uncertainty and make decisions based on data. Probability measures the likelihood of events, while statistics involves collecting, analyzing, and interpreting data to draw conclusions about populations based on samples.
+
+Number theory investigates the properties of integers and has applications in cryptography and computer security. Abstract algebra studies algebraic structures like groups, rings, and fields, providing the theoretical foundation for many areas of mathematics and computer science.
+
+Discrete mathematics deals with countable, distinct objects and includes topics like graph theory, combinatorics, and logic. These areas are particularly important in computer science for algorithm design, network analysis, and formal verification of systems.

examples/reranker/corpus/misc_doc.txt

@@ -0,0 +1,13 @@
+Cooking is both an art and a science that transforms raw ingredients into delicious and nutritious meals. The fundamentals of cooking involve understanding heat transfer, chemical reactions, and flavor combinations. Different cooking methods like boiling, roasting, sautéing, and grilling each impart unique textures and flavors to food.
+
+Baking requires precise measurements and timing since it relies on chemical reactions between ingredients. Proteins coagulate, starches gelatinize, and the Maillard reaction creates complex flavors and browning. Understanding these processes helps bakers achieve consistent results and troubleshoot problems.
+
+Weather patterns are driven by the unequal heating of Earth's surface by the sun. Air masses with different temperatures and humidity levels interact to create various weather phenomena. High and low pressure systems move across the globe, bringing changes in temperature, precipitation, and wind patterns.
+
+Gardening connects us with nature and provides fresh produce, beautiful flowers, and peaceful outdoor spaces. Successful gardening requires understanding soil composition, plant nutrition, watering needs, and seasonal cycles. Composting organic waste creates nutrient-rich soil amendments while reducing household waste.
+
+Travel broadens perspectives and creates lasting memories through exposure to different cultures, cuisines, and landscapes. Planning trips involves considering factors like budget, transportation, accommodation, and local customs. Sustainable tourism practices help preserve destinations for future generations.
+
+Photography captures moments in time and allows creative expression through composition, lighting, and timing. Digital photography has made the art more accessible, but fundamental principles like the rule of thirds, depth of field, and understanding light remain important for creating compelling images.
+
+Music theory provides the framework for understanding how melodies, harmonies, and rhythms work together to create emotional and aesthetic experiences. Different musical traditions around the world have developed unique scales, instruments, and compositional techniques that reflect their cultural heritage.

examples/reranker/corpus/physics_doc.txt

@@ -0,0 +1,13 @@
+Physics is the fundamental science that seeks to understand how the universe works at all scales, from subatomic particles to galaxies. It studies matter, energy, motion, forces, space, and time, providing the foundation for all other natural sciences and many technological advances.
+
+Classical mechanics, developed by Newton, describes the motion of objects from projectiles to planets. It includes concepts like force, momentum, energy, and conservation laws. These principles govern everything from car engines to spacecraft trajectories and remain accurate for everyday-scale phenomena.
+
+Thermodynamics deals with heat, temperature, and energy transfer. The four laws of thermodynamics govern how energy flows and transforms, explaining everything from steam engines to the heat death of the universe. Concepts like entropy help us understand the direction of time and the efficiency limits of machines.
+
+Electromagnetism unifies electricity and magnetism into a single theory. Maxwell's equations describe how electric and magnetic fields interact and propagate as electromagnetic waves, including light, radio waves, and X-rays. This theory enabled the development of modern electronics and telecommunications.
+
+Quantum mechanics describes the behavior of matter and energy at atomic and subatomic scales. Unlike classical physics, quantum systems exhibit wave-particle duality, uncertainty principles, and superposition. These strange properties enable technologies like lasers, transistors, and quantum computers.
+
+Einstein's theory of relativity revolutionized our understanding of space, time, and gravity. Special relativity shows that space and time are intertwined and that nothing can travel faster than light. General relativity describes gravity as the curvature of spacetime and predicts phenomena like black holes and gravitational waves.
+
+Modern physics continues to push boundaries with particle physics exploring fundamental constituents of matter, condensed matter physics investigating complex materials, and cosmology studying the origin and evolution of the universe. These fields drive innovations in technology, medicine, and our understanding of nature.

examples/reranker/main.rs

@@ -0,0 +1,144 @@
+use anyhow::Result;
+use transformers::models::implementations::{Qwen3EmbeddingSize, Qwen3RerankSize};
+use transformers::pipelines::embedding_pipeline::*;
+use transformers::pipelines::reranker_pipeline::*;
+
+use std::fs;
+use std::path::Path;
+
+fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
+    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
+    let norm_a: f32 = a.iter().map(|v| v * v).sum::<f32>().sqrt();
+    let norm_b: f32 = b.iter().map(|v| v * v).sum::<f32>().sqrt();
+    dot / (norm_a * norm_b)
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    println!("Building pipelines...");
+
+    /* ---------- build the two pipelines ---------- */
+    let embed_pipe = EmbeddingPipelineBuilder::qwen3(Qwen3EmbeddingSize::Size0_6B)
+        .cpu() // or .cuda(0) if you have GPU
+        .build()
+        .await?;
+
+    let rerank_pipe = RerankPipelineBuilder::qwen3(Qwen3RerankSize::Size0_6B)
+        .cpu()
+        .build()
+        .await?;
+
+    println!("Pipelines built\n");
+
+    println!("Loading corpus...");
+
+    /* ---------- load corpus (directory of text files) ---------- */
+    let corpus_path = Path::new("examples").join("reranker").join("corpus");
+    let mut documents: Vec<String> = Vec::new();
+    let mut filenames: Vec<String> = Vec::new();
+
+    if corpus_path.is_dir() {
+        for entry in fs::read_dir(&corpus_path)? {
+            let entry = entry?;
+            if entry.file_type()?.is_file() {
+                let content = fs::read_to_string(entry.path())?;
+                if !content.trim().is_empty() {
+                    documents.push(content);
+                    filenames.push(entry.file_name().to_string_lossy().to_string());
+                }
+            }
+        }
+    }
+
+    // Fallback to a tiny hard-coded corpus if the directory is missing or empty
+    if documents.is_empty() {
+        documents = vec![
+            "Machine learning is a subset of artificial intelligence that focuses on algorithms."
+                .to_string(),
+            "Cooking recipes often require precise measurements and timing.".to_string(),
+            "Supervised learning uses labeled data to train models.".to_string(),
+            "The weather today is sunny and warm.".to_string(),
+            "Deep learning is a type of machine learning using neural networks.".to_string(),
+        ];
+        filenames = vec![
+            "demo_doc_1".to_string(),
+            "demo_doc_2".to_string(),
+            "demo_doc_3".to_string(),
+            "demo_doc_4".to_string(),
+            "demo_doc_5".to_string(),
+        ];
+        eprintln!("[info] corpus directory not found or empty – using built-in demo corpus");
+    }
+
+    println!("📚 Loaded {} documents from corpus\n", documents.len());
+
+    println!("Embedding corpus...");
+
+    /* ---------- embed the corpus once up front ---------- */
+    let doc_embeddings: Vec<Vec<f32>> = documents
+        .iter()
+        .map(|doc| embed_pipe.embed(doc).expect("embedding failed"))
+        .collect();
+
+    println!("Corpus embedded!");
+
+    /* ---------- incoming query ---------- */
+    let query = "How do neural networks and deep learning work?";
+    println!("🔍 Query: \"{}\"", query);
+
+    let query_emb = embed_pipe.embed(query)?;
+
+    // Helper function to truncate documents for display
+    fn truncate_doc(doc: &str, max_chars: usize) -> String {
+        if doc.len() <= max_chars {
+            doc.to_string()
+        } else {
+            format!("{}...", &doc[..max_chars])
+        }
+    }
+
+    /* ---------- stage-1: ANN recall ---------- */
+    let mut scored: Vec<(usize, f32)> = doc_embeddings
+        .iter()
+        .enumerate()
+        .map(|(i, emb)| (i, cosine_sim(&query_emb, emb)))
+        .collect();
+
+    // higher cosine = more similar
+    scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+
+    // keep the top-k (here k = 3)
+    let k = 3;
+    let candidate_idxs: Vec<usize> = scored.iter().take(k).map(|(i, _)| *i).collect();
+    let candidate_docs: Vec<&str> = candidate_idxs
+        .iter()
+        .map(|&i| documents[i].as_str())
+        .collect();
+
+    println!("=== Top-{k} after embedding recall ===");
+    for (rank, (idx, score)) in scored.iter().take(k).enumerate() {
+        println!(
+            "{:>2}. [sim {:.4}] {} | {}",
+            rank + 1,
+            score,
+            filenames[*idx],
+            truncate_doc(&documents[*idx], 80)
+        );
+    }
+
+    /* ---------- stage-2: rerank those candidates ---------- */
+    let reranked = rerank_pipe.rerank(query, &candidate_docs).await?;
+
+    println!("\n=== Final ranking from Qwen3-Reranker ===");
+    for (rank, res) in reranked.iter().enumerate() {
+        println!(
+            "{:>2}. [score {:.4}] {} | {}",
+            rank + 1,
+            res.score,
+            filenames[candidate_idxs[res.index]],
+            truncate_doc(candidate_docs[res.index], 80)
+        );
+    }
+
+    Ok(())
+}

src/lib.rs

@@ -24,5 +24,6 @@ pub use models::implementations::{
     Qwen3Model,
     Qwen3Size,
     Qwen3EmbeddingModel,
+    Qwen3EmbeddingSize,
     Qwen3RerankModel,
 };

src/models/implementations/mod.rs

@@ -7,5 +7,5 @@ pub mod qwen3_reranker;
 pub use gemma3::{Gemma3Model, Gemma3Size};
 pub use modernbert::{ModernBertModel, ModernBertSize};
 pub use qwen3::{Qwen3Model, Qwen3Size};
-pub use qwen3_embeddings::Qwen3EmbeddingModel;
+pub use qwen3_embeddings::{Qwen3EmbeddingModel, Qwen3EmbeddingSize};
 pub use qwen3_reranker::{Qwen3RerankModel, Qwen3RerankSize};