mirror of https://github.com/buster-so/buster.git
56 lines
2.0 KiB
Rust
56 lines
2.0 KiB
Rust
use rerank::{Reranker, RerankResult};
|
|
use std::error::Error;
|
|
|
|
#[tokio::test]
|
|
async fn test_reranker_integration() -> Result<(), Box<dyn Error>> {
|
|
// Load environment variables from .env file
|
|
dotenv::dotenv().ok();
|
|
|
|
// Initialize the reranker
|
|
let reranker = Reranker::new()?;
|
|
|
|
// Define a sample query and documents
|
|
let query = "What is the capital of France?";
|
|
let documents = vec![
|
|
"Paris is a major European city and a global center for art, fashion, gastronomy and culture.",
|
|
"London is the capital and largest city of England and the United Kingdom.",
|
|
"The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France.",
|
|
"Berlin is the capital and largest city of Germany by both area and population.",
|
|
];
|
|
let top_n = 2;
|
|
|
|
// Perform reranking
|
|
let results: Vec<RerankResult> = reranker.rerank(query, &documents, top_n).await?;
|
|
|
|
// Assertions
|
|
assert_eq!(results.len(), top_n, "Should return top_n results");
|
|
|
|
// Check that indices are within the bounds of the original documents
|
|
for result in &results {
|
|
assert!(result.index < documents.len(), "Result index should be valid");
|
|
}
|
|
|
|
// Optional: Print results for manual verification (can be removed later)
|
|
println!("Query: {}", query);
|
|
for result in &results {
|
|
println!(
|
|
"Document Index: {}, Score: {:.4}, Document: {}",
|
|
result.index,
|
|
result.relevance_score,
|
|
documents[result.index]
|
|
);
|
|
}
|
|
|
|
// Example assertion: if we expect Paris-related documents to be ranked higher.
|
|
// This is a very basic check and might need adjustment based on actual model behavior.
|
|
if !results.is_empty() {
|
|
let first_result_doc = documents[results[0].index];
|
|
assert!(
|
|
first_result_doc.to_lowercase().contains("paris"),
|
|
"The top result for 'capital of France' should ideally mention Paris. Model output: {}",
|
|
first_result_doc
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|