use aprender::recommend::ContentRecommender;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
fn generate_movie_descriptions(n: usize) -> Vec<(String, String)> {
let genres = [
"action",
"comedy",
"drama",
"thriller",
"horror",
"romance",
"scifi",
"fantasy",
"mystery",
"adventure",
];
let adjectives = [
"epic",
"thrilling",
"emotional",
"intense",
"hilarious",
"dark",
"heartwarming",
"suspenseful",
"mysterious",
"explosive",
];
let nouns = [
"story",
"journey",
"adventure",
"tale",
"saga",
"quest",
"mission",
"odyssey",
"expedition",
"voyage",
];
(0..n)
.map(|i| {
let genre = genres[i % genres.len()];
let adj = adjectives[(i / 10) % adjectives.len()];
let noun = nouns[(i / 100) % nouns.len()];
let id = format!("movie_{i}");
let desc = format!("{adj} {genre} {noun} about heroes and villains");
(id, desc)
})
.collect()
}
fn bench_add_items(c: &mut Criterion) {
let mut group = c.benchmark_group("recommend_add");
for size in &[100, 1_000, 10_000] {
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| {
b.iter(|| {
let mut rec = ContentRecommender::new(16, 200, 0.95);
let items = generate_movie_descriptions(size);
for (id, desc) in items {
rec.add_item(black_box(id), black_box(desc));
}
rec
});
});
}
group.finish();
}
fn bench_recommend(c: &mut Criterion) {
let mut group = c.benchmark_group("recommend_search");
group.sample_size(50);
for size in &[100, 1_000, 10_000] {
let mut rec = ContentRecommender::new(16, 200, 0.95);
let items = generate_movie_descriptions(*size);
for (id, desc) in items {
rec.add_item(id, desc);
}
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
rec.recommend(black_box("movie_0"), black_box(10))
.expect("should succeed")
});
});
}
group.finish();
}
fn bench_recommend_latency_target(c: &mut Criterion) {
let mut rec = ContentRecommender::new(16, 200, 0.95);
let items = generate_movie_descriptions(10_000);
for (id, desc) in items {
rec.add_item(id, desc);
}
c.bench_function("recommend_10k_latency", |b| {
b.iter(|| {
rec.recommend(black_box("movie_5000"), black_box(10))
.expect("should succeed")
});
});
}
criterion_group!(
benches,
bench_add_items,
bench_recommend,
bench_recommend_latency_target
);
criterion_main!(benches);