39 lines
1.7 KiB
Rust
39 lines
1.7 KiB
Rust
//! Benchmarks for summarization systems
|
|
|
|
use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
|
|
use rtx_nlg::{MockModelInterface, summarization::*};
|
|
use std::sync::Arc;
|
|
|
|
fn bench_summarization(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("summarization");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let config = SummarizationConfig::default();
|
|
let summarizer = Summarizer::new(model, config);
|
|
|
|
let test_documents = [
|
|
"Short document.",
|
|
"This is a medium length document with several sentences. It contains information about various topics and should be summarized effectively.",
|
|
"This is a very long document that contains multiple paragraphs with detailed information. The document discusses artificial intelligence, machine learning algorithms, natural language processing techniques, and their applications in modern technology. It covers various aspects of these technologies including their development history, current state, and future prospects. The document also examines the challenges and opportunities in implementing these technologies across different industries and domains.",
|
|
];
|
|
|
|
for (i, doc) in test_documents.iter().enumerate() {
|
|
group.bench_with_input(
|
|
format!("doc_words_{}", doc.split_whitespace().count()),
|
|
doc,
|
|
|b, doc| {
|
|
b.iter(|| {
|
|
let result = summarizer.summarize(black_box(doc)).unwrap();
|
|
black_box(result);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(summarization_benches, bench_summarization);
|
|
criterion_main!(summarization_benches);
|