Skip to main content

test_reranking/
test_reranking.rs

1#[cfg(feature = "fastembed")]
2use ck_embed::create_reranker;
3
4fn main() {
5    #[cfg(not(feature = "fastembed"))]
6    {
7        println!("This example requires the 'fastembed' feature to be enabled.");
8        println!("Run with: cargo run --example test_reranking --features fastembed");
9        return;
10    }
11
12    #[cfg(feature = "fastembed")]
13    run_example();
14}
15
16#[cfg(feature = "fastembed")]
17fn run_example() {
18    println!("=== Reranking Test ===");
19
20    let mut reranker = create_reranker(None).expect("Failed to create reranker");
21
22    println!("Created reranker: {}", reranker.id());
23
24    let query = "error handling in programming";
25    let documents = vec![
26        "try catch exception handling in Java".to_string(),
27        "user interface design patterns".to_string(),
28        "error handling with Result types in Rust".to_string(),
29        "database connection management".to_string(),
30        "exception handling best practices".to_string(),
31    ];
32
33    println!("\nQuery: '{}'", query);
34    println!("Documents to rerank:");
35    for (i, doc) in documents.iter().enumerate() {
36        println!("  {}: {}", i + 1, doc);
37    }
38
39    let results = reranker
40        .rerank(query, &documents)
41        .expect("Failed to rerank documents");
42
43    println!("\nReranked results:");
44    // Sort by score (descending)
45    let mut sorted_results = results;
46    sorted_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
47
48    for (i, result) in sorted_results.iter().enumerate() {
49        println!(
50            "  {}: [Score: {:.3}] {}",
51            i + 1,
52            result.score,
53            result.document
54        );
55    }
56
57    println!("\n✅ Reranking test completed successfully!");
58}