rerank/
rerank.rs

1//! Complete example for re-ranking
2//! 
3//! Reproduces the example from <https://huggingface.co/Alibaba-NLP/gte-reranker-modernbert-base>
4//! and checks for consistency of the results.
5fn main() -> gte::util::result::Result<()> {    
6    const TOKENIZER_PATH: &str = "models/gte-modernbert-base/tokenizer.json";
7    const MODEL_PATH: &str = "models/gte-reranker-modernbert-base/onnx/model.onnx";
8    const APPLY_SIGMOID: bool = true;
9
10    let params = gte::params::Parameters::default().with_sigmoid(APPLY_SIGMOID);
11    let pipeline = gte::rerank::pipeline::RerankingPipeline::new(TOKENIZER_PATH, &params)?;
12    let model = orp::model::Model::new(MODEL_PATH, orp::params::RuntimeParameters::default())?;
13
14    let inputs = gte::rerank::input::TextInput::from_str(&[
15        ("what is the capital of China?", "Beijing"),
16        ("how to implement quick sort in python?", "Introduction of quick sort"),
17        ("how to implement quick sort in python?", "The weather is nice today"),
18    ]);
19
20    const EXPECTED_SIMILARITIES: [f32; 3] = if APPLY_SIGMOID { [0.8945664, 0.9213594, 0.15742092] } else { [2.1387, 2.4609, -1.6729] };    
21    const EPSILON: f32 = if APPLY_SIGMOID { 0.00001 } else { 0.01 };
22
23    let outputs = model.inference(inputs, &pipeline, &params)?;
24
25    println!("Scores: {:?}", outputs.scores);
26    assert!(gte::util::test::is_close_to_a(&outputs.scores.view(), &EXPECTED_SIMILARITIES, EPSILON));
27
28    Ok(())
29}