Squashed 'vendor/ruvector/' content from commit b64c2172
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
This commit is contained in:
759
examples/vibecast-7sense/benches/api_benchmark.rs
Normal file
759
examples/vibecast-7sense/benches/api_benchmark.rs
Normal file
@@ -0,0 +1,759 @@
|
||||
//! API Benchmark Suite for 7sense
|
||||
//!
|
||||
//! Performance targets from ADR-004:
|
||||
//! - Query latency: <100ms total (end-to-end)
|
||||
//! - Neighbor search: <50ms p99
|
||||
//! - Evidence pack generation: <200ms
|
||||
|
||||
use criterion::{
|
||||
black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
mod utils;
|
||||
use utils::*;
|
||||
|
||||
// ============================================================================
|
||||
// Simulated API Types
|
||||
// ============================================================================
|
||||
|
||||
/// Neighbor search request
|
||||
#[derive(Clone, Debug)]
|
||||
struct NeighborSearchRequest {
|
||||
embedding: Vec<f32>,
|
||||
k: usize,
|
||||
filter: Option<SearchFilter>,
|
||||
include_metadata: bool,
|
||||
}
|
||||
|
||||
/// Search filter for neighbor queries
|
||||
#[derive(Clone, Debug)]
|
||||
struct SearchFilter {
|
||||
species: Option<Vec<String>>,
|
||||
location: Option<BoundingBox>,
|
||||
time_range: Option<TimeRange>,
|
||||
min_confidence: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct BoundingBox {
|
||||
min_lat: f32,
|
||||
max_lat: f32,
|
||||
min_lon: f32,
|
||||
max_lon: f32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct TimeRange {
|
||||
start: i64,
|
||||
end: i64,
|
||||
}
|
||||
|
||||
/// Neighbor search response
|
||||
#[derive(Clone, Debug)]
|
||||
struct NeighborSearchResponse {
|
||||
results: Vec<SearchResult>,
|
||||
total_time_ms: u64,
|
||||
cache_hit: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct SearchResult {
|
||||
id: String,
|
||||
distance: f32,
|
||||
metadata: Option<EmbeddingMetadata>,
|
||||
}
|
||||
|
||||
/// Embedding metadata
|
||||
#[derive(Clone, Debug)]
|
||||
struct EmbeddingMetadata {
|
||||
recording_id: String,
|
||||
species: Option<String>,
|
||||
call_type: Option<String>,
|
||||
location: Option<Location>,
|
||||
timestamp: i64,
|
||||
confidence: f32,
|
||||
audio_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Location {
|
||||
lat: f32,
|
||||
lon: f32,
|
||||
site_name: Option<String>,
|
||||
}
|
||||
|
||||
/// Evidence pack for interpretability
|
||||
#[derive(Clone, Debug)]
|
||||
struct EvidencePack {
|
||||
query_embedding: Vec<f32>,
|
||||
neighbors: Vec<NeighborEvidence>,
|
||||
cluster_info: ClusterInfo,
|
||||
spectrogram_url: Option<String>,
|
||||
attention_map: Option<Vec<Vec<f32>>>,
|
||||
confidence_breakdown: ConfidenceBreakdown,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct NeighborEvidence {
|
||||
result: SearchResult,
|
||||
similarity_score: f32,
|
||||
contributing_features: Vec<FeatureContribution>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct FeatureContribution {
|
||||
feature_name: String,
|
||||
contribution: f32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct ClusterInfo {
|
||||
cluster_id: i32,
|
||||
cluster_size: usize,
|
||||
centroid_distance: f32,
|
||||
typical_species: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct ConfidenceBreakdown {
|
||||
neighbor_agreement: f32,
|
||||
cluster_membership: f32,
|
||||
embedding_quality: f32,
|
||||
overall: f32,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Simulated API Service
|
||||
// ============================================================================
|
||||
|
||||
/// Simulated API service for benchmarking
|
||||
struct ApiService {
|
||||
index: SimpleHnswIndex,
|
||||
metadata_store: HashMap<usize, EmbeddingMetadata>,
|
||||
cluster_centroids: Vec<Vec<f32>>,
|
||||
cluster_assignments: Vec<i32>,
|
||||
}
|
||||
|
||||
impl ApiService {
|
||||
fn new(index: SimpleHnswIndex, num_clusters: usize) -> Self {
|
||||
let size = index.len();
|
||||
|
||||
// Generate fake metadata
|
||||
let mut metadata_store = HashMap::new();
|
||||
let species = ["Robin", "Sparrow", "Blackbird", "Thrush", "Finch"];
|
||||
let call_types = ["song", "call", "alarm", "contact"];
|
||||
|
||||
for i in 0..size {
|
||||
metadata_store.insert(
|
||||
i,
|
||||
EmbeddingMetadata {
|
||||
recording_id: format!("rec_{}", i),
|
||||
species: Some(species[i % species.len()].to_string()),
|
||||
call_type: Some(call_types[i % call_types.len()].to_string()),
|
||||
location: Some(Location {
|
||||
lat: 51.5 + (i as f32 * 0.001),
|
||||
lon: -0.1 + (i as f32 * 0.001),
|
||||
site_name: Some(format!("Site {}", i % 10)),
|
||||
}),
|
||||
timestamp: 1700000000 + (i as i64 * 300),
|
||||
confidence: 0.7 + (i as f32 % 30) / 100.0,
|
||||
audio_url: Some(format!("https://audio.example.com/{}.wav", i)),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Generate cluster centroids and assignments
|
||||
let cluster_centroids = generate_random_vectors(num_clusters, PERCH_EMBEDDING_DIM);
|
||||
let cluster_assignments: Vec<i32> = (0..size).map(|i| (i % num_clusters) as i32).collect();
|
||||
|
||||
Self {
|
||||
index,
|
||||
metadata_store,
|
||||
cluster_centroids,
|
||||
cluster_assignments,
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute neighbor search
|
||||
fn neighbor_search(&self, request: &NeighborSearchRequest) -> NeighborSearchResponse {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Perform HNSW search
|
||||
let raw_results = self.index.search(&request.embedding, request.k * 2);
|
||||
|
||||
// Apply filters
|
||||
let filtered_results: Vec<_> = raw_results
|
||||
.into_iter()
|
||||
.filter(|(idx, _)| self.apply_filter(*idx, &request.filter))
|
||||
.take(request.k)
|
||||
.collect();
|
||||
|
||||
// Build response with optional metadata
|
||||
let results: Vec<SearchResult> = filtered_results
|
||||
.into_iter()
|
||||
.map(|(idx, distance)| SearchResult {
|
||||
id: format!("emb_{}", idx),
|
||||
distance,
|
||||
metadata: if request.include_metadata {
|
||||
self.metadata_store.get(&idx).cloned()
|
||||
} else {
|
||||
None
|
||||
},
|
||||
})
|
||||
.collect();
|
||||
|
||||
NeighborSearchResponse {
|
||||
results,
|
||||
total_time_ms: start.elapsed().as_millis() as u64,
|
||||
cache_hit: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_filter(&self, idx: usize, filter: &Option<SearchFilter>) -> bool {
|
||||
match filter {
|
||||
None => true,
|
||||
Some(f) => {
|
||||
if let Some(metadata) = self.metadata_store.get(&idx) {
|
||||
// Species filter
|
||||
if let Some(species_list) = &f.species {
|
||||
if let Some(species) = &metadata.species {
|
||||
if !species_list.contains(species) {
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Confidence filter
|
||||
if let Some(min_conf) = f.min_confidence {
|
||||
if metadata.confidence < min_conf {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Time range filter
|
||||
if let Some(time_range) = &f.time_range {
|
||||
if metadata.timestamp < time_range.start
|
||||
|| metadata.timestamp > time_range.end
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Location filter
|
||||
if let Some(bbox) = &f.location {
|
||||
if let Some(loc) = &metadata.location {
|
||||
if loc.lat < bbox.min_lat
|
||||
|| loc.lat > bbox.max_lat
|
||||
|| loc.lon < bbox.min_lon
|
||||
|| loc.lon > bbox.max_lon
|
||||
{
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate evidence pack for interpretability
|
||||
fn generate_evidence_pack(&self, embedding: &[f32], k: usize) -> EvidencePack {
|
||||
// Get neighbors
|
||||
let raw_results = self.index.search(embedding, k);
|
||||
|
||||
let neighbors: Vec<NeighborEvidence> = raw_results
|
||||
.iter()
|
||||
.map(|(idx, distance)| {
|
||||
let metadata = self.metadata_store.get(idx).cloned();
|
||||
let similarity = 1.0 / (1.0 + distance);
|
||||
|
||||
// Generate feature contributions (mock)
|
||||
let contributions: Vec<FeatureContribution> = (0..5)
|
||||
.map(|i| FeatureContribution {
|
||||
feature_name: format!("feature_{}", i),
|
||||
contribution: similarity * (1.0 - i as f32 * 0.1),
|
||||
})
|
||||
.collect();
|
||||
|
||||
NeighborEvidence {
|
||||
result: SearchResult {
|
||||
id: format!("emb_{}", idx),
|
||||
distance: *distance,
|
||||
metadata,
|
||||
},
|
||||
similarity_score: similarity,
|
||||
contributing_features: contributions,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Compute cluster info
|
||||
let cluster_info = self.compute_cluster_info(embedding);
|
||||
|
||||
// Compute confidence breakdown
|
||||
let confidence_breakdown = self.compute_confidence(embedding, &neighbors);
|
||||
|
||||
EvidencePack {
|
||||
query_embedding: embedding.to_vec(),
|
||||
neighbors,
|
||||
cluster_info,
|
||||
spectrogram_url: Some("https://spectrograms.example.com/query.png".to_string()),
|
||||
attention_map: Some(self.generate_attention_map()),
|
||||
confidence_breakdown,
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_cluster_info(&self, embedding: &[f32]) -> ClusterInfo {
|
||||
// Find nearest cluster
|
||||
let mut best_cluster = 0;
|
||||
let mut best_distance = f32::MAX;
|
||||
|
||||
for (i, centroid) in self.cluster_centroids.iter().enumerate() {
|
||||
let dist = l2_distance(embedding, centroid);
|
||||
if dist < best_distance {
|
||||
best_distance = dist;
|
||||
best_cluster = i;
|
||||
}
|
||||
}
|
||||
|
||||
// Count cluster members
|
||||
let cluster_size = self
|
||||
.cluster_assignments
|
||||
.iter()
|
||||
.filter(|&&c| c == best_cluster as i32)
|
||||
.count();
|
||||
|
||||
ClusterInfo {
|
||||
cluster_id: best_cluster as i32,
|
||||
cluster_size,
|
||||
centroid_distance: best_distance,
|
||||
typical_species: vec!["Robin".to_string(), "Sparrow".to_string()],
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_confidence(&self, _embedding: &[f32], neighbors: &[NeighborEvidence]) -> ConfidenceBreakdown {
|
||||
// Compute neighbor agreement
|
||||
let neighbor_agreement = if !neighbors.is_empty() {
|
||||
let avg_sim: f32 = neighbors.iter().map(|n| n.similarity_score).sum::<f32>()
|
||||
/ neighbors.len() as f32;
|
||||
avg_sim
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
ConfidenceBreakdown {
|
||||
neighbor_agreement,
|
||||
cluster_membership: 0.85,
|
||||
embedding_quality: 0.92,
|
||||
overall: (neighbor_agreement + 0.85 + 0.92) / 3.0,
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_attention_map(&self) -> Vec<Vec<f32>> {
|
||||
// Generate a small mock attention map
|
||||
(0..32)
|
||||
.map(|i| (0..128).map(|j| ((i * j) % 100) as f32 / 100.0).collect())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Neighbor Search Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark neighbor search endpoint
|
||||
fn benchmark_neighbor_search_endpoint(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("neighbor_search_endpoint");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(15));
|
||||
|
||||
// Build test service
|
||||
let index = setup_test_index(50_000);
|
||||
let service = ApiService::new(index, 50);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
|
||||
// Basic search without metadata
|
||||
group.bench_function("basic_k10", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 10,
|
||||
filter: None,
|
||||
include_metadata: false,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// Search with metadata
|
||||
group.bench_function("with_metadata_k10", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 10,
|
||||
filter: None,
|
||||
include_metadata: true,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// Search with filters
|
||||
group.bench_function("filtered_k10", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 10,
|
||||
filter: Some(SearchFilter {
|
||||
species: Some(vec!["Robin".to_string(), "Sparrow".to_string()]),
|
||||
location: None,
|
||||
time_range: None,
|
||||
min_confidence: Some(0.8),
|
||||
}),
|
||||
include_metadata: true,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// Different k values
|
||||
for &k in &[10, 50, 100] {
|
||||
group.bench_with_input(BenchmarkId::new("k", k), &k, |b, &k| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k,
|
||||
filter: None,
|
||||
include_metadata: true,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark search throughput under concurrent load
|
||||
fn benchmark_search_throughput(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("search_throughput");
|
||||
group.sample_size(20);
|
||||
group.measurement_time(Duration::from_secs(20));
|
||||
|
||||
let index = setup_test_index(50_000);
|
||||
let service = ApiService::new(index, 50);
|
||||
|
||||
// Batch of queries
|
||||
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
let requests: Vec<NeighborSearchRequest> = queries
|
||||
.into_iter()
|
||||
.map(|embedding| NeighborSearchRequest {
|
||||
embedding,
|
||||
k: 10,
|
||||
filter: None,
|
||||
include_metadata: true,
|
||||
})
|
||||
.collect();
|
||||
|
||||
group.throughput(Throughput::Elements(requests.len() as u64));
|
||||
group.bench_function("batch_100_queries", |b| {
|
||||
b.iter(|| {
|
||||
for request in &requests {
|
||||
black_box(service.neighbor_search(request));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Evidence Pack Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark evidence pack generation
|
||||
fn benchmark_evidence_pack_generation(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("evidence_pack_generation");
|
||||
group.sample_size(30);
|
||||
group.measurement_time(Duration::from_secs(15));
|
||||
|
||||
let index = setup_test_index(50_000);
|
||||
let service = ApiService::new(index, 50);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
|
||||
// Basic evidence pack
|
||||
group.bench_function("basic", |b| {
|
||||
b.iter(|| black_box(service.generate_evidence_pack(&query, 10)));
|
||||
});
|
||||
|
||||
// Different neighbor counts
|
||||
for &k in &[5, 10, 20, 50] {
|
||||
group.bench_with_input(BenchmarkId::new("neighbors", k), &k, |b, &k| {
|
||||
b.iter(|| black_box(service.generate_evidence_pack(&query, k)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Filter Performance Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark filter application performance
|
||||
fn benchmark_filter_performance(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("filter_performance");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
let index = setup_test_index(50_000);
|
||||
let service = ApiService::new(index, 50);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
|
||||
// No filter
|
||||
group.bench_function("no_filter", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 100,
|
||||
filter: None,
|
||||
include_metadata: false,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// Species filter only
|
||||
group.bench_function("species_filter", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 100,
|
||||
filter: Some(SearchFilter {
|
||||
species: Some(vec!["Robin".to_string()]),
|
||||
location: None,
|
||||
time_range: None,
|
||||
min_confidence: None,
|
||||
}),
|
||||
include_metadata: false,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// Confidence filter only
|
||||
group.bench_function("confidence_filter", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 100,
|
||||
filter: Some(SearchFilter {
|
||||
species: None,
|
||||
location: None,
|
||||
time_range: None,
|
||||
min_confidence: Some(0.9),
|
||||
}),
|
||||
include_metadata: false,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
// All filters combined
|
||||
group.bench_function("all_filters", |b| {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 100,
|
||||
filter: Some(SearchFilter {
|
||||
species: Some(vec!["Robin".to_string(), "Sparrow".to_string()]),
|
||||
location: Some(BoundingBox {
|
||||
min_lat: 51.0,
|
||||
max_lat: 52.0,
|
||||
min_lon: -1.0,
|
||||
max_lon: 1.0,
|
||||
}),
|
||||
time_range: Some(TimeRange {
|
||||
start: 1700000000,
|
||||
end: 1710000000,
|
||||
}),
|
||||
min_confidence: Some(0.8),
|
||||
}),
|
||||
include_metadata: false,
|
||||
};
|
||||
b.iter(|| black_box(service.neighbor_search(&request)));
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Latency Analysis
|
||||
// ============================================================================
|
||||
|
||||
/// Analyze end-to-end latency against targets
|
||||
fn analyze_api_latency() {
|
||||
use std::time::Instant;
|
||||
|
||||
println!("\n=== API Latency Analysis ===\n");
|
||||
|
||||
// Build service
|
||||
let index = setup_test_index(100_000);
|
||||
let service = ApiService::new(index, 50);
|
||||
|
||||
let queries = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
|
||||
// Neighbor search latency
|
||||
let mut search_latencies = Vec::new();
|
||||
for query in &queries {
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query.clone(),
|
||||
k: 10,
|
||||
filter: None,
|
||||
include_metadata: true,
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
let _ = service.neighbor_search(&request);
|
||||
search_latencies.push(start.elapsed());
|
||||
}
|
||||
|
||||
let search_stats = PerformanceStats::from_latencies(search_latencies);
|
||||
println!("Neighbor Search (k=10, with metadata):");
|
||||
println!("{}", search_stats.report());
|
||||
println!(
|
||||
" p99 target: {}ms ({})",
|
||||
targets::QUERY_LATENCY_P99_MS,
|
||||
if search_stats.p99 <= Duration::from_millis(targets::QUERY_LATENCY_P99_MS) {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
println!(
|
||||
" Total target: {}ms ({})",
|
||||
targets::TOTAL_QUERY_LATENCY_MS,
|
||||
if search_stats.p99 <= Duration::from_millis(targets::TOTAL_QUERY_LATENCY_MS) {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
println!();
|
||||
|
||||
// Evidence pack latency
|
||||
let mut evidence_latencies = Vec::new();
|
||||
for query in queries.iter().take(100) {
|
||||
let start = Instant::now();
|
||||
let _ = service.generate_evidence_pack(query, 10);
|
||||
evidence_latencies.push(start.elapsed());
|
||||
}
|
||||
|
||||
let evidence_stats = PerformanceStats::from_latencies(evidence_latencies);
|
||||
println!("Evidence Pack Generation (10 neighbors):");
|
||||
println!("{}", evidence_stats.report());
|
||||
println!(
|
||||
" p99 target: 200ms ({})",
|
||||
if evidence_stats.p99 <= Duration::from_millis(200) {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Criterion Groups
|
||||
// ============================================================================
|
||||
|
||||
criterion_group!(
|
||||
name = search_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_neighbor_search_endpoint, benchmark_search_throughput
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = evidence_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_evidence_pack_generation
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = filter_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_filter_performance
|
||||
);
|
||||
|
||||
criterion_main!(search_benches, evidence_benches, filter_benches);
|
||||
|
||||
// ============================================================================
|
||||
// Tests
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_neighbor_search_basic() {
|
||||
let index = setup_test_index(1000);
|
||||
let service = ApiService::new(index, 10);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query,
|
||||
k: 10,
|
||||
filter: None,
|
||||
include_metadata: true,
|
||||
};
|
||||
|
||||
let response = service.neighbor_search(&request);
|
||||
assert_eq!(response.results.len(), 10);
|
||||
assert!(response.results.iter().all(|r| r.metadata.is_some()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_neighbor_search_with_filter() {
|
||||
let index = setup_test_index(1000);
|
||||
let service = ApiService::new(index, 10);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
let request = NeighborSearchRequest {
|
||||
embedding: query,
|
||||
k: 10,
|
||||
filter: Some(SearchFilter {
|
||||
species: Some(vec!["Robin".to_string()]),
|
||||
location: None,
|
||||
time_range: None,
|
||||
min_confidence: Some(0.7),
|
||||
}),
|
||||
include_metadata: true,
|
||||
};
|
||||
|
||||
let response = service.neighbor_search(&request);
|
||||
// All results should match filter
|
||||
for result in &response.results {
|
||||
if let Some(metadata) = &result.metadata {
|
||||
assert_eq!(metadata.species, Some("Robin".to_string()));
|
||||
assert!(metadata.confidence >= 0.7);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evidence_pack_generation() {
|
||||
let index = setup_test_index(1000);
|
||||
let service = ApiService::new(index, 10);
|
||||
|
||||
let query = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
let evidence = service.generate_evidence_pack(&query, 10);
|
||||
|
||||
assert_eq!(evidence.neighbors.len(), 10);
|
||||
assert!(evidence.confidence_breakdown.overall > 0.0);
|
||||
assert!(evidence.cluster_info.cluster_size > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore] // Run with: cargo test --release -- --ignored --nocapture
|
||||
fn run_api_latency_analysis() {
|
||||
analyze_api_latency();
|
||||
}
|
||||
}
|
||||
684
examples/vibecast-7sense/benches/clustering_benchmark.rs
Normal file
684
examples/vibecast-7sense/benches/clustering_benchmark.rs
Normal file
@@ -0,0 +1,684 @@
|
||||
//! Clustering Benchmark Suite for 7sense
|
||||
//!
|
||||
//! Benchmarks for clustering algorithms used in bird call analysis:
|
||||
//! - HDBSCAN for species/call-type clustering
|
||||
//! - Cluster assignment for new embeddings
|
||||
//! - Motif detection in audio sequences
|
||||
//! - Centroid computation and updates
|
||||
|
||||
use criterion::{
|
||||
black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
|
||||
};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::time::Duration;
|
||||
|
||||
mod utils;
|
||||
use utils::*;
|
||||
|
||||
/// Number of clusters for benchmark
|
||||
const NUM_CLUSTERS: usize = 50;
|
||||
|
||||
// ============================================================================
|
||||
// Simplified HDBSCAN Implementation for Benchmarking
|
||||
// ============================================================================
|
||||
|
||||
/// Simplified HDBSCAN-like clustering for benchmarking
|
||||
/// In production, this would use the actual HDBSCAN algorithm
|
||||
struct SimpleHdbscan {
|
||||
min_cluster_size: usize,
|
||||
min_samples: usize,
|
||||
epsilon: f32,
|
||||
}
|
||||
|
||||
impl SimpleHdbscan {
|
||||
fn new(min_cluster_size: usize, min_samples: usize, epsilon: f32) -> Self {
|
||||
Self {
|
||||
min_cluster_size,
|
||||
min_samples,
|
||||
epsilon,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fit the clustering model on embeddings
|
||||
/// Returns cluster labels (-1 for noise)
|
||||
fn fit(&self, embeddings: &[Vec<f32>]) -> Vec<i32> {
|
||||
let n = embeddings.len();
|
||||
let mut labels = vec![-1i32; n];
|
||||
let mut cluster_id = 0;
|
||||
let mut visited = vec![false; n];
|
||||
|
||||
for i in 0..n {
|
||||
if visited[i] {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Find neighbors within epsilon
|
||||
let neighbors = self.region_query(embeddings, i);
|
||||
|
||||
if neighbors.len() >= self.min_samples {
|
||||
// Expand cluster
|
||||
let cluster_members = self.expand_cluster(embeddings, i, &neighbors, &mut visited);
|
||||
|
||||
if cluster_members.len() >= self.min_cluster_size {
|
||||
for &member in &cluster_members {
|
||||
labels[member] = cluster_id;
|
||||
}
|
||||
cluster_id += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
labels
|
||||
}
|
||||
|
||||
fn region_query(&self, embeddings: &[Vec<f32>], point_idx: usize) -> Vec<usize> {
|
||||
let point = &embeddings[point_idx];
|
||||
embeddings
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, other)| l2_distance(point, other) <= self.epsilon)
|
||||
.map(|(idx, _)| idx)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn expand_cluster(
|
||||
&self,
|
||||
embeddings: &[Vec<f32>],
|
||||
seed: usize,
|
||||
initial_neighbors: &[usize],
|
||||
visited: &mut [bool],
|
||||
) -> Vec<usize> {
|
||||
let mut cluster = vec![seed];
|
||||
visited[seed] = true;
|
||||
|
||||
let mut to_process: Vec<usize> = initial_neighbors.to_vec();
|
||||
|
||||
while let Some(point_idx) = to_process.pop() {
|
||||
if visited[point_idx] {
|
||||
continue;
|
||||
}
|
||||
visited[point_idx] = true;
|
||||
cluster.push(point_idx);
|
||||
|
||||
let neighbors = self.region_query(embeddings, point_idx);
|
||||
if neighbors.len() >= self.min_samples {
|
||||
to_process.extend(neighbors.iter().filter(|&&n| !visited[n]));
|
||||
}
|
||||
}
|
||||
|
||||
cluster
|
||||
}
|
||||
}
|
||||
|
||||
/// Cluster assignment for new embeddings
|
||||
struct ClusterAssigner {
|
||||
centroids: Vec<Vec<f32>>,
|
||||
cluster_ids: Vec<usize>,
|
||||
}
|
||||
|
||||
impl ClusterAssigner {
|
||||
fn new(centroids: Vec<Vec<f32>>) -> Self {
|
||||
let cluster_ids = (0..centroids.len()).collect();
|
||||
Self {
|
||||
centroids,
|
||||
cluster_ids,
|
||||
}
|
||||
}
|
||||
|
||||
/// Assign a single embedding to the nearest cluster
|
||||
fn assign(&self, embedding: &[f32]) -> (usize, f32) {
|
||||
let mut best_cluster = 0;
|
||||
let mut best_distance = f32::MAX;
|
||||
|
||||
for (i, centroid) in self.centroids.iter().enumerate() {
|
||||
let dist = l2_distance(embedding, centroid);
|
||||
if dist < best_distance {
|
||||
best_distance = dist;
|
||||
best_cluster = self.cluster_ids[i];
|
||||
}
|
||||
}
|
||||
|
||||
(best_cluster, best_distance)
|
||||
}
|
||||
|
||||
/// Batch assign embeddings
|
||||
fn batch_assign(&self, embeddings: &[Vec<f32>]) -> Vec<(usize, f32)> {
|
||||
embeddings.iter().map(|e| self.assign(e)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute cluster centroids from labeled embeddings
|
||||
fn compute_centroids(embeddings: &[Vec<f32>], labels: &[i32]) -> HashMap<i32, Vec<f32>> {
|
||||
let dims = embeddings[0].len();
|
||||
let mut sums: HashMap<i32, Vec<f32>> = HashMap::new();
|
||||
let mut counts: HashMap<i32, usize> = HashMap::new();
|
||||
|
||||
for (embedding, &label) in embeddings.iter().zip(labels.iter()) {
|
||||
if label >= 0 {
|
||||
let sum = sums.entry(label).or_insert_with(|| vec![0.0; dims]);
|
||||
for (s, &e) in sum.iter_mut().zip(embedding.iter()) {
|
||||
*s += e;
|
||||
}
|
||||
*counts.entry(label).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let mut centroids = HashMap::new();
|
||||
for (label, sum) in sums {
|
||||
let count = counts[&label] as f32;
|
||||
let centroid: Vec<f32> = sum.iter().map(|&s| s / count).collect();
|
||||
centroids.insert(label, centroid);
|
||||
}
|
||||
|
||||
centroids
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Motif Detection
|
||||
// ============================================================================
|
||||
|
||||
/// Simplified motif detector for recurring audio patterns
|
||||
struct MotifDetector {
|
||||
min_length: usize,
|
||||
max_gap: usize,
|
||||
similarity_threshold: f32,
|
||||
}
|
||||
|
||||
impl MotifDetector {
|
||||
fn new(min_length: usize, max_gap: usize, similarity_threshold: f32) -> Self {
|
||||
Self {
|
||||
min_length,
|
||||
max_gap,
|
||||
similarity_threshold,
|
||||
}
|
||||
}
|
||||
|
||||
/// Detect motifs in a sequence of embeddings
|
||||
fn detect_motifs(&self, embeddings: &[Vec<f32>]) -> Vec<Motif> {
|
||||
let mut motifs = Vec::new();
|
||||
let n = embeddings.len();
|
||||
|
||||
// Simplified matrix profile approach
|
||||
for i in 0..n.saturating_sub(self.min_length) {
|
||||
for j in (i + self.min_length)..n.saturating_sub(self.min_length) {
|
||||
// Check if subsequences are similar
|
||||
let sim = self.subsequence_similarity(embeddings, i, j, self.min_length);
|
||||
|
||||
if sim >= self.similarity_threshold {
|
||||
motifs.push(Motif {
|
||||
start_a: i,
|
||||
start_b: j,
|
||||
length: self.min_length,
|
||||
similarity: sim,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
motifs
|
||||
}
|
||||
|
||||
fn subsequence_similarity(
|
||||
&self,
|
||||
embeddings: &[Vec<f32>],
|
||||
start_a: usize,
|
||||
start_b: usize,
|
||||
length: usize,
|
||||
) -> f32 {
|
||||
let mut total_sim = 0.0;
|
||||
|
||||
for i in 0..length {
|
||||
let sim = cosine_similarity(&embeddings[start_a + i], &embeddings[start_b + i]);
|
||||
total_sim += sim;
|
||||
}
|
||||
|
||||
total_sim / length as f32
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct Motif {
|
||||
start_a: usize,
|
||||
start_b: usize,
|
||||
length: usize,
|
||||
similarity: f32,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HDBSCAN Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark HDBSCAN clustering
|
||||
fn benchmark_hdbscan(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hdbscan");
|
||||
group.sample_size(20);
|
||||
group.measurement_time(Duration::from_secs(30));
|
||||
|
||||
for &size in &[500, 1000, 2000] {
|
||||
// Generate clustered data for more realistic benchmark
|
||||
let embeddings = generate_clustered_vectors(size, PERCH_EMBEDDING_DIM, NUM_CLUSTERS, 0.1);
|
||||
|
||||
let hdbscan = SimpleHdbscan::new(5, 3, 0.5);
|
||||
|
||||
group.throughput(Throughput::Elements(size as u64));
|
||||
group.bench_with_input(BenchmarkId::new("fit", size), &size, |b, _| {
|
||||
b.iter(|| black_box(hdbscan.fit(&embeddings)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark HDBSCAN with different parameters
|
||||
fn benchmark_hdbscan_params(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hdbscan_params");
|
||||
group.sample_size(10);
|
||||
group.measurement_time(Duration::from_secs(20));
|
||||
|
||||
let size = 1000;
|
||||
let embeddings = generate_clustered_vectors(size, PERCH_EMBEDDING_DIM, NUM_CLUSTERS, 0.1);
|
||||
|
||||
for min_cluster_size in [5, 10, 20] {
|
||||
let hdbscan = SimpleHdbscan::new(min_cluster_size, 3, 0.5);
|
||||
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("min_cluster_size", min_cluster_size),
|
||||
&min_cluster_size,
|
||||
|b, _| {
|
||||
b.iter(|| black_box(hdbscan.fit(&embeddings)));
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
for epsilon in [0.3, 0.5, 0.7] {
|
||||
let hdbscan = SimpleHdbscan::new(5, 3, epsilon);
|
||||
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("epsilon", format!("{:.1}", epsilon)),
|
||||
&epsilon,
|
||||
|b, _| {
|
||||
b.iter(|| black_box(hdbscan.fit(&embeddings)));
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Cluster Assignment Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark cluster assignment for new embeddings
|
||||
fn benchmark_cluster_assignment(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("cluster_assignment");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
// Generate centroids
|
||||
let centroids = generate_random_vectors(NUM_CLUSTERS, PERCH_EMBEDDING_DIM);
|
||||
let assigner = ClusterAssigner::new(centroids);
|
||||
|
||||
// Benchmark single assignment
|
||||
let single_embedding = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
group.bench_function("single", |b| {
|
||||
b.iter(|| black_box(assigner.assign(&single_embedding)));
|
||||
});
|
||||
|
||||
// Benchmark batch assignment
|
||||
for &batch_size in &[100, 1000, 10000] {
|
||||
let embeddings = generate_random_vectors(batch_size, PERCH_EMBEDDING_DIM);
|
||||
|
||||
group.throughput(Throughput::Elements(batch_size as u64));
|
||||
group.bench_with_input(BenchmarkId::new("batch", batch_size), &batch_size, |b, _| {
|
||||
b.iter(|| black_box(assigner.batch_assign(&embeddings)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark cluster assignment with different numbers of clusters
|
||||
fn benchmark_cluster_assignment_scalability(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("cluster_assignment_scalability");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
let embeddings = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
|
||||
for num_clusters in [10, 50, 100, 200, 500] {
|
||||
let centroids = generate_random_vectors(num_clusters, PERCH_EMBEDDING_DIM);
|
||||
let assigner = ClusterAssigner::new(centroids);
|
||||
|
||||
group.throughput(Throughput::Elements(embeddings.len() as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("num_clusters", num_clusters),
|
||||
&num_clusters,
|
||||
|b, _| {
|
||||
b.iter(|| black_box(assigner.batch_assign(&embeddings)));
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Centroid Computation Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark centroid computation
|
||||
fn benchmark_centroid_computation(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("centroid_computation");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
for &size in &[1000, 5000, 10000] {
|
||||
let embeddings = generate_clustered_vectors(size, PERCH_EMBEDDING_DIM, NUM_CLUSTERS, 0.1);
|
||||
|
||||
// Create synthetic labels
|
||||
let labels: Vec<i32> = (0..size).map(|i| (i % NUM_CLUSTERS) as i32).collect();
|
||||
|
||||
group.throughput(Throughput::Elements(size as u64));
|
||||
group.bench_with_input(BenchmarkId::new("size", size), &size, |b, _| {
|
||||
b.iter(|| black_box(compute_centroids(&embeddings, &labels)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark incremental centroid update
|
||||
fn benchmark_centroid_update(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("centroid_update");
|
||||
group.sample_size(100);
|
||||
|
||||
// Pre-compute initial centroid
|
||||
let cluster_size = 1000;
|
||||
let cluster_embeddings = generate_random_vectors(cluster_size, PERCH_EMBEDDING_DIM);
|
||||
let initial_centroid: Vec<f32> = (0..PERCH_EMBEDDING_DIM)
|
||||
.map(|d| {
|
||||
cluster_embeddings.iter().map(|e| e[d]).sum::<f32>() / cluster_size as f32
|
||||
})
|
||||
.collect();
|
||||
|
||||
// New embedding to add
|
||||
let new_embedding = generate_random_vectors(1, PERCH_EMBEDDING_DIM).remove(0);
|
||||
|
||||
group.bench_function("incremental_update", |b| {
|
||||
b.iter(|| {
|
||||
// Incremental centroid update formula:
|
||||
// new_centroid = old_centroid + (new_point - old_centroid) / (n + 1)
|
||||
let n = cluster_size as f32;
|
||||
let updated: Vec<f32> = initial_centroid
|
||||
.iter()
|
||||
.zip(new_embedding.iter())
|
||||
.map(|(&c, &e)| c + (e - c) / (n + 1.0))
|
||||
.collect();
|
||||
black_box(updated)
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Motif Detection Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark motif detection
|
||||
fn benchmark_motif_detection(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("motif_detection");
|
||||
group.sample_size(20);
|
||||
group.measurement_time(Duration::from_secs(20));
|
||||
|
||||
let detector = MotifDetector::new(3, 10, 0.8);
|
||||
|
||||
for &seq_length in &[50, 100, 200] {
|
||||
// Generate sequence with some repeated patterns
|
||||
let mut embeddings = generate_clustered_vectors(seq_length, PERCH_EMBEDDING_DIM, 10, 0.05);
|
||||
|
||||
group.throughput(Throughput::Elements(seq_length as u64));
|
||||
group.bench_with_input(BenchmarkId::new("seq_length", seq_length), &seq_length, |b, _| {
|
||||
b.iter(|| black_box(detector.detect_motifs(&embeddings)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Silhouette Score Computation
|
||||
// ============================================================================
|
||||
|
||||
/// Compute silhouette score for cluster quality assessment
|
||||
fn compute_silhouette_score(embeddings: &[Vec<f32>], labels: &[i32]) -> f32 {
|
||||
let n = embeddings.len();
|
||||
if n < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let unique_labels: HashSet<i32> = labels.iter().filter(|&&l| l >= 0).copied().collect();
|
||||
if unique_labels.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut total_score = 0.0;
|
||||
let mut count = 0;
|
||||
|
||||
for i in 0..n {
|
||||
let label_i = labels[i];
|
||||
if label_i < 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute a(i): mean intra-cluster distance
|
||||
let mut intra_sum = 0.0;
|
||||
let mut intra_count = 0;
|
||||
for j in 0..n {
|
||||
if i != j && labels[j] == label_i {
|
||||
intra_sum += l2_distance(&embeddings[i], &embeddings[j]);
|
||||
intra_count += 1;
|
||||
}
|
||||
}
|
||||
let a_i = if intra_count > 0 {
|
||||
intra_sum / intra_count as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Compute b(i): min mean inter-cluster distance
|
||||
let mut b_i = f32::MAX;
|
||||
for &other_label in &unique_labels {
|
||||
if other_label != label_i {
|
||||
let mut inter_sum = 0.0;
|
||||
let mut inter_count = 0;
|
||||
for j in 0..n {
|
||||
if labels[j] == other_label {
|
||||
inter_sum += l2_distance(&embeddings[i], &embeddings[j]);
|
||||
inter_count += 1;
|
||||
}
|
||||
}
|
||||
if inter_count > 0 {
|
||||
let mean_inter = inter_sum / inter_count as f32;
|
||||
b_i = b_i.min(mean_inter);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Silhouette coefficient for point i
|
||||
if b_i.is_finite() {
|
||||
let s_i = (b_i - a_i) / a_i.max(b_i);
|
||||
total_score += s_i;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
total_score / count as f32
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Benchmark silhouette score computation
|
||||
fn benchmark_silhouette_score(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("silhouette_score");
|
||||
group.sample_size(10);
|
||||
group.measurement_time(Duration::from_secs(30));
|
||||
|
||||
for &size in &[100, 500] {
|
||||
let embeddings = generate_clustered_vectors(size, PERCH_EMBEDDING_DIM, 10, 0.1);
|
||||
let labels: Vec<i32> = (0..size).map(|i| (i % 10) as i32).collect();
|
||||
|
||||
group.bench_with_input(BenchmarkId::new("size", size), &size, |b, _| {
|
||||
b.iter(|| black_box(compute_silhouette_score(&embeddings, &labels)));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Criterion Groups
|
||||
// ============================================================================
|
||||
|
||||
criterion_group!(
|
||||
name = hdbscan_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_hdbscan, benchmark_hdbscan_params
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = assignment_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_cluster_assignment, benchmark_cluster_assignment_scalability
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = centroid_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_centroid_computation, benchmark_centroid_update
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = motif_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_motif_detection
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = quality_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_silhouette_score
|
||||
);
|
||||
|
||||
criterion_main!(
|
||||
hdbscan_benches,
|
||||
assignment_benches,
|
||||
centroid_benches,
|
||||
motif_benches,
|
||||
quality_benches
|
||||
);
|
||||
|
||||
// ============================================================================
|
||||
// Tests
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hdbscan_clustering() {
|
||||
let embeddings = generate_clustered_vectors(100, 128, 5, 0.05);
|
||||
let hdbscan = SimpleHdbscan::new(5, 3, 0.5);
|
||||
|
||||
let labels = hdbscan.fit(&embeddings);
|
||||
assert_eq!(labels.len(), 100);
|
||||
|
||||
// Should have some non-noise labels
|
||||
let non_noise: Vec<_> = labels.iter().filter(|&&l| l >= 0).collect();
|
||||
assert!(!non_noise.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cluster_assignment() {
|
||||
let centroids = generate_random_vectors(10, 128);
|
||||
let assigner = ClusterAssigner::new(centroids.clone());
|
||||
|
||||
// Assign a centroid to itself should return that cluster
|
||||
let (cluster, dist) = assigner.assign(¢roids[5]);
|
||||
assert_eq!(cluster, 5);
|
||||
assert!(dist < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_centroid_computation() {
|
||||
let embeddings = vec![
|
||||
vec![1.0, 0.0],
|
||||
vec![0.0, 1.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![-1.0, -1.0],
|
||||
];
|
||||
let labels = vec![0, 0, 0, 1];
|
||||
|
||||
let centroids = compute_centroids(&embeddings, &labels);
|
||||
|
||||
assert_eq!(centroids.len(), 2);
|
||||
|
||||
// Cluster 0 centroid should be (2/3, 2/3)
|
||||
let c0 = ¢roids[&0];
|
||||
assert!((c0[0] - 2.0 / 3.0).abs() < 1e-5);
|
||||
assert!((c0[1] - 2.0 / 3.0).abs() < 1e-5);
|
||||
|
||||
// Cluster 1 centroid should be (-1, -1)
|
||||
let c1 = ¢roids[&1];
|
||||
assert!((c1[0] - (-1.0)).abs() < 1e-5);
|
||||
assert!((c1[1] - (-1.0)).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_motif_detection() {
|
||||
// Create sequence with a repeated pattern
|
||||
let mut embeddings = Vec::new();
|
||||
let pattern: Vec<Vec<f32>> = (0..3)
|
||||
.map(|i| {
|
||||
let mut v = vec![0.0f32; 128];
|
||||
v[i] = 1.0;
|
||||
v
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Insert pattern twice with gap
|
||||
embeddings.extend(pattern.clone());
|
||||
embeddings.extend(generate_random_vectors(5, 128));
|
||||
embeddings.extend(pattern);
|
||||
|
||||
let detector = MotifDetector::new(3, 10, 0.9);
|
||||
let motifs = detector.detect_motifs(&embeddings);
|
||||
|
||||
// Should detect at least one motif
|
||||
// Note: Due to noise, this may not always work perfectly
|
||||
println!("Found {} motifs", motifs.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_silhouette_score() {
|
||||
// Perfect clustering: two well-separated clusters
|
||||
let embeddings = vec![
|
||||
vec![0.0, 0.0],
|
||||
vec![0.1, 0.0],
|
||||
vec![0.0, 0.1],
|
||||
vec![10.0, 10.0],
|
||||
vec![10.1, 10.0],
|
||||
vec![10.0, 10.1],
|
||||
];
|
||||
let labels = vec![0, 0, 0, 1, 1, 1];
|
||||
|
||||
let score = compute_silhouette_score(&embeddings, &labels);
|
||||
|
||||
// Score should be close to 1 for well-separated clusters
|
||||
assert!(score > 0.5, "Silhouette score {} too low", score);
|
||||
}
|
||||
}
|
||||
563
examples/vibecast-7sense/benches/embedding_benchmark.rs
Normal file
563
examples/vibecast-7sense/benches/embedding_benchmark.rs
Normal file
@@ -0,0 +1,563 @@
|
||||
//! Embedding Benchmark Suite for 7sense
|
||||
//!
|
||||
//! Performance targets from ADR-004:
|
||||
//! - Embedding inference: >100 segments/second
|
||||
//! - Mel spectrogram compute: <20ms per segment
|
||||
//! - Embedding normalization: <5ms per segment
|
||||
//! - Batch ingestion: 1M vectors/minute
|
||||
|
||||
use criterion::{
|
||||
black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
mod utils;
|
||||
use utils::*;
|
||||
|
||||
/// Audio segment parameters (5 seconds at 32kHz)
|
||||
const AUDIO_SAMPLE_RATE: usize = 32_000;
|
||||
const SEGMENT_DURATION_SECS: f32 = 5.0;
|
||||
const SEGMENT_SAMPLES: usize = (AUDIO_SAMPLE_RATE as f32 * SEGMENT_DURATION_SECS) as usize;
|
||||
|
||||
/// Mel spectrogram parameters
|
||||
const N_MELS: usize = 128;
|
||||
const N_FFT: usize = 2048;
|
||||
const HOP_LENGTH: usize = 512;
|
||||
const MEL_FRAMES: usize = (SEGMENT_SAMPLES / HOP_LENGTH) + 1;
|
||||
|
||||
// ============================================================================
|
||||
// Simulated Audio Processing
|
||||
// ============================================================================
|
||||
|
||||
/// Generate synthetic audio samples for benchmarking
|
||||
fn generate_audio_segment() -> Vec<f32> {
|
||||
let mut samples = Vec::with_capacity(SEGMENT_SAMPLES);
|
||||
let mut seed = 12345u64;
|
||||
|
||||
for i in 0..SEGMENT_SAMPLES {
|
||||
// Simple synthetic audio with multiple frequencies
|
||||
let t = i as f32 / AUDIO_SAMPLE_RATE as f32;
|
||||
let sample = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * 0.3
|
||||
+ (2.0 * std::f32::consts::PI * 880.0 * t).sin() * 0.2
|
||||
+ (2.0 * std::f32::consts::PI * 1320.0 * t).sin() * 0.1;
|
||||
|
||||
// Add some noise
|
||||
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
let noise = ((seed >> 33) as f32 / u32::MAX as f32) * 0.1 - 0.05;
|
||||
|
||||
samples.push(sample + noise);
|
||||
}
|
||||
|
||||
samples
|
||||
}
|
||||
|
||||
/// Simulated mel spectrogram computation
|
||||
/// In production, this would use actual FFT and mel filterbank
|
||||
fn compute_mel_spectrogram(audio: &[f32]) -> Vec<Vec<f32>> {
|
||||
let num_frames = (audio.len() / HOP_LENGTH) + 1;
|
||||
let mut spectrogram = Vec::with_capacity(num_frames);
|
||||
|
||||
for frame_idx in 0..num_frames {
|
||||
let start = frame_idx * HOP_LENGTH;
|
||||
let end = (start + N_FFT).min(audio.len());
|
||||
|
||||
// Simulated FFT and mel filterbank
|
||||
let mut mel_frame = vec![0.0f32; N_MELS];
|
||||
for (i, &sample) in audio[start..end].iter().enumerate() {
|
||||
let bin = i % N_MELS;
|
||||
mel_frame[bin] += sample.abs();
|
||||
}
|
||||
|
||||
// Apply log scaling
|
||||
for val in mel_frame.iter_mut() {
|
||||
*val = (*val + 1e-10).ln();
|
||||
}
|
||||
|
||||
spectrogram.push(mel_frame);
|
||||
}
|
||||
|
||||
spectrogram
|
||||
}
|
||||
|
||||
/// Simulated embedding inference (mock ONNX model)
|
||||
/// In production, this would use the actual Perch 2.0 model
|
||||
fn compute_embedding(spectrogram: &[Vec<f32>]) -> Vec<f32> {
|
||||
let mut embedding = vec![0.0f32; PERCH_EMBEDDING_DIM];
|
||||
|
||||
// Simulated neural network computation
|
||||
for (i, frame) in spectrogram.iter().enumerate() {
|
||||
for (j, &mel) in frame.iter().enumerate() {
|
||||
let embed_idx = (i * N_MELS + j) % PERCH_EMBEDDING_DIM;
|
||||
embedding[embed_idx] += mel * 0.01;
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize to unit length
|
||||
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in embedding.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
|
||||
embedding
|
||||
}
|
||||
|
||||
/// L2 normalize an embedding vector
|
||||
fn normalize_embedding(embedding: &mut [f32]) {
|
||||
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in embedding.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Spectrogram Generation Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark mel spectrogram generation
|
||||
fn benchmark_spectrogram_generation(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("spectrogram_generation");
|
||||
group.sample_size(100);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
let audio = generate_audio_segment();
|
||||
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("single_segment", |b| {
|
||||
b.iter(|| black_box(compute_mel_spectrogram(&audio)));
|
||||
});
|
||||
|
||||
// Batch spectrogram computation
|
||||
let batch_sizes = [10, 50, 100];
|
||||
for &batch_size in &batch_sizes {
|
||||
let audio_batch: Vec<Vec<f32>> = (0..batch_size).map(|_| generate_audio_segment()).collect();
|
||||
|
||||
group.throughput(Throughput::Elements(batch_size as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("batch", batch_size),
|
||||
&batch_size,
|
||||
|b, _| {
|
||||
b.iter(|| {
|
||||
for audio in &audio_batch {
|
||||
black_box(compute_mel_spectrogram(audio));
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Embedding Inference Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark embedding inference (mock ONNX)
|
||||
fn benchmark_embedding_inference(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("embedding_inference");
|
||||
group.sample_size(100);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
// Pre-compute spectrogram
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("single_inference", |b| {
|
||||
b.iter(|| black_box(compute_embedding(&spectrogram)));
|
||||
});
|
||||
|
||||
// Batch inference
|
||||
let batch_sizes = [10, 32, 64, 128];
|
||||
for &batch_size in &batch_sizes {
|
||||
let spectrograms: Vec<Vec<Vec<f32>>> = (0..batch_size)
|
||||
.map(|_| {
|
||||
let audio = generate_audio_segment();
|
||||
compute_mel_spectrogram(&audio)
|
||||
})
|
||||
.collect();
|
||||
|
||||
group.throughput(Throughput::Elements(batch_size as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("batch", batch_size),
|
||||
&batch_size,
|
||||
|b, _| {
|
||||
b.iter(|| {
|
||||
for spec in &spectrograms {
|
||||
black_box(compute_embedding(spec));
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark full pipeline: audio -> spectrogram -> embedding
|
||||
fn benchmark_full_pipeline(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("full_embedding_pipeline");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(15));
|
||||
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("single_segment", |b| {
|
||||
b.iter(|| {
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
let embedding = compute_embedding(&spectrogram);
|
||||
black_box(embedding)
|
||||
});
|
||||
});
|
||||
|
||||
// Batch pipeline
|
||||
for &batch_size in &[10, 50, 100] {
|
||||
group.throughput(Throughput::Elements(batch_size as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("batch", batch_size),
|
||||
&batch_size,
|
||||
|b, &size| {
|
||||
b.iter(|| {
|
||||
for _ in 0..size {
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
let embedding = compute_embedding(&spectrogram);
|
||||
black_box(embedding);
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Normalization Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark embedding normalization
|
||||
fn benchmark_normalization(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("normalization");
|
||||
group.sample_size(100);
|
||||
group.measurement_time(Duration::from_secs(5));
|
||||
|
||||
// Generate random unnormalized embeddings
|
||||
let embeddings: Vec<Vec<f32>> = (0..1000)
|
||||
.map(|i| {
|
||||
let mut vec = vec![0.0f32; PERCH_EMBEDDING_DIM];
|
||||
let mut seed = (i as u64).wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
for v in vec.iter_mut() {
|
||||
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
*v = ((seed >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0;
|
||||
}
|
||||
vec
|
||||
})
|
||||
.collect();
|
||||
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("single", |b| {
|
||||
let mut embedding = embeddings[0].clone();
|
||||
b.iter(|| {
|
||||
normalize_embedding(&mut embedding);
|
||||
black_box(&embedding);
|
||||
});
|
||||
});
|
||||
|
||||
group.throughput(Throughput::Elements(1000));
|
||||
group.bench_function("batch_1000", |b| {
|
||||
let mut batch = embeddings.clone();
|
||||
b.iter(|| {
|
||||
for embedding in batch.iter_mut() {
|
||||
normalize_embedding(embedding);
|
||||
}
|
||||
black_box(&batch);
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Quantization Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark scalar quantization (float32 -> int8)
|
||||
fn benchmark_quantization(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("quantization");
|
||||
group.sample_size(100);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
// Generate embeddings and calibrate quantizer
|
||||
let embeddings = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
let mut quantizer = ScalarQuantizer::new(PERCH_EMBEDDING_DIM);
|
||||
quantizer.calibrate(&embeddings);
|
||||
|
||||
// Benchmark quantization
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("quantize_single", |b| {
|
||||
let embedding = &embeddings[0];
|
||||
b.iter(|| black_box(quantizer.quantize(embedding)));
|
||||
});
|
||||
|
||||
// Batch quantization
|
||||
group.throughput(Throughput::Elements(1000));
|
||||
group.bench_function("quantize_batch_1000", |b| {
|
||||
b.iter(|| {
|
||||
for embedding in &embeddings {
|
||||
black_box(quantizer.quantize(embedding));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// Benchmark dequantization
|
||||
let quantized: Vec<Vec<u8>> = embeddings.iter().map(|e| quantizer.quantize(e)).collect();
|
||||
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("dequantize_single", |b| {
|
||||
let q = &quantized[0];
|
||||
b.iter(|| black_box(quantizer.dequantize(q)));
|
||||
});
|
||||
|
||||
group.throughput(Throughput::Elements(1000));
|
||||
group.bench_function("dequantize_batch_1000", |b| {
|
||||
b.iter(|| {
|
||||
for q in &quantized {
|
||||
black_box(quantizer.dequantize(q));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark quantization error measurement
|
||||
fn benchmark_quantization_error(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("quantization_error");
|
||||
group.sample_size(50);
|
||||
|
||||
let embeddings = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
let mut quantizer = ScalarQuantizer::new(PERCH_EMBEDDING_DIM);
|
||||
quantizer.calibrate(&embeddings);
|
||||
|
||||
group.bench_function("measure_error", |b| {
|
||||
b.iter(|| {
|
||||
let mut total_error = 0.0f32;
|
||||
let mut max_error = 0.0f32;
|
||||
|
||||
for embedding in &embeddings {
|
||||
let quantized = quantizer.quantize(embedding);
|
||||
let dequantized = quantizer.dequantize(&quantized);
|
||||
|
||||
let error: f32 = embedding
|
||||
.iter()
|
||||
.zip(dequantized.iter())
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.sum();
|
||||
|
||||
total_error += error;
|
||||
max_error = max_error.max(error);
|
||||
}
|
||||
|
||||
black_box((total_error / embeddings.len() as f32, max_error))
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Throughput Analysis
|
||||
// ============================================================================
|
||||
|
||||
/// Analyze embedding throughput against targets
|
||||
fn analyze_embedding_throughput() {
|
||||
use std::time::Instant;
|
||||
|
||||
println!("\n=== Embedding Throughput Analysis ===\n");
|
||||
|
||||
// Target: 100 segments/second
|
||||
let target_segments_per_sec = targets::EMBEDDING_SEGMENTS_PER_SECOND;
|
||||
let num_segments = 100;
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
for _ in 0..num_segments {
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
let _embedding = compute_embedding(&spectrogram);
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
let throughput = num_segments as f64 / elapsed.as_secs_f64();
|
||||
|
||||
println!("Processed {} segments in {:?}", num_segments, elapsed);
|
||||
println!("Throughput: {:.1} segments/sec", throughput);
|
||||
println!(
|
||||
"Target: {} segments/sec ({})",
|
||||
target_segments_per_sec,
|
||||
if throughput >= target_segments_per_sec as f64 {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Half-Precision (float16) Simulation
|
||||
// ============================================================================
|
||||
|
||||
/// Simulate float16 quantization for warm tier storage
|
||||
fn simulate_float16(embedding: &[f32]) -> Vec<u16> {
|
||||
embedding
|
||||
.iter()
|
||||
.map(|&v| half::f16::from_f32(v).to_bits())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Benchmark float16 conversion
|
||||
fn benchmark_float16_conversion(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("float16_conversion");
|
||||
group.sample_size(100);
|
||||
|
||||
let embeddings = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
|
||||
group.throughput(Throughput::Elements(1000));
|
||||
group.bench_function("to_float16", |b| {
|
||||
b.iter(|| {
|
||||
for embedding in &embeddings {
|
||||
black_box(simulate_float16(embedding));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// Benchmark float16 -> float32 conversion
|
||||
let float16_embeddings: Vec<Vec<u16>> = embeddings.iter().map(|e| simulate_float16(e)).collect();
|
||||
|
||||
group.bench_function("from_float16", |b| {
|
||||
b.iter(|| {
|
||||
for embedding in &float16_embeddings {
|
||||
let restored: Vec<f32> = embedding
|
||||
.iter()
|
||||
.map(|&bits| half::f16::from_bits(bits).to_f32())
|
||||
.collect();
|
||||
black_box(restored);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Criterion Groups
|
||||
// ============================================================================
|
||||
|
||||
criterion_group!(
|
||||
name = spectrogram_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_spectrogram_generation
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = inference_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_embedding_inference, benchmark_full_pipeline
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = normalization_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_normalization
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = quantization_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_quantization, benchmark_quantization_error, benchmark_float16_conversion
|
||||
);
|
||||
|
||||
criterion_main!(
|
||||
spectrogram_benches,
|
||||
inference_benches,
|
||||
normalization_benches,
|
||||
quantization_benches
|
||||
);
|
||||
|
||||
// ============================================================================
|
||||
// Tests
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_audio_generation() {
|
||||
let audio = generate_audio_segment();
|
||||
assert_eq!(audio.len(), SEGMENT_SAMPLES);
|
||||
|
||||
// Check samples are in reasonable range
|
||||
for &sample in &audio {
|
||||
assert!(sample.abs() < 2.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectrogram_generation() {
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
|
||||
assert!(!spectrogram.is_empty());
|
||||
assert_eq!(spectrogram[0].len(), N_MELS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_embedding_computation() {
|
||||
let audio = generate_audio_segment();
|
||||
let spectrogram = compute_mel_spectrogram(&audio);
|
||||
let embedding = compute_embedding(&spectrogram);
|
||||
|
||||
assert_eq!(embedding.len(), PERCH_EMBEDDING_DIM);
|
||||
|
||||
// Check normalization
|
||||
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
assert!((norm - 1.0).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quantization_roundtrip() {
|
||||
let embeddings = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
let mut quantizer = ScalarQuantizer::new(PERCH_EMBEDDING_DIM);
|
||||
quantizer.calibrate(&embeddings);
|
||||
|
||||
for embedding in &embeddings {
|
||||
let quantized = quantizer.quantize(embedding);
|
||||
let dequantized = quantizer.dequantize(&quantized);
|
||||
|
||||
// Check dimensions preserved
|
||||
assert_eq!(dequantized.len(), embedding.len());
|
||||
|
||||
// Check error is bounded
|
||||
let max_error: f32 = embedding
|
||||
.iter()
|
||||
.zip(dequantized.iter())
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0.0, f32::max);
|
||||
|
||||
// Max error per dimension should be small
|
||||
assert!(max_error < 0.1, "Max error {} too large", max_error);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore] // Run with: cargo test --release -- --ignored --nocapture
|
||||
fn run_throughput_analysis() {
|
||||
analyze_embedding_throughput();
|
||||
}
|
||||
}
|
||||
505
examples/vibecast-7sense/benches/hnsw_benchmark.rs
Normal file
505
examples/vibecast-7sense/benches/hnsw_benchmark.rs
Normal file
@@ -0,0 +1,505 @@
|
||||
//! HNSW Benchmark Suite for 7sense
|
||||
//!
|
||||
//! Performance targets from ADR-004:
|
||||
//! - HNSW Search: 150x speedup vs brute force
|
||||
//! - Query Latency p99: < 50ms
|
||||
//! - Recall@10: >= 0.95
|
||||
//! - Recall@100: >= 0.98
|
||||
//! - Insert Throughput: >= 10,000 vectors/s
|
||||
//! - Build Time: < 30 min for 1M vectors
|
||||
|
||||
use criterion::{
|
||||
black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
|
||||
};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
mod utils;
|
||||
use utils::*;
|
||||
|
||||
/// Index sizes to benchmark
|
||||
const SMALL_INDEX: usize = 10_000;
|
||||
const MEDIUM_INDEX: usize = 100_000;
|
||||
const LARGE_INDEX: usize = 500_000;
|
||||
|
||||
/// K values for search benchmarks
|
||||
const K_VALUES: &[usize] = &[10, 50, 100];
|
||||
|
||||
// ============================================================================
|
||||
// HNSW Search Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark HNSW search performance with different index sizes and k values
|
||||
fn benchmark_hnsw_search(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_search");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
// Generate query vectors once
|
||||
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
|
||||
for &size in &[SMALL_INDEX, MEDIUM_INDEX] {
|
||||
// Build index
|
||||
println!("Building index with {} vectors...", size);
|
||||
let index = setup_test_index(size);
|
||||
|
||||
for &k in K_VALUES {
|
||||
group.throughput(Throughput::Elements(queries.len() as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new(format!("size_{}_k_{}", size, k), k),
|
||||
&k,
|
||||
|b, &k| {
|
||||
b.iter(|| {
|
||||
for query in &queries {
|
||||
black_box(index.search(query, k));
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark HNSW search with different ef_search values
|
||||
fn benchmark_hnsw_search_ef(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_search_ef");
|
||||
group.sample_size(30);
|
||||
group.measurement_time(Duration::from_secs(8));
|
||||
|
||||
let size = MEDIUM_INDEX;
|
||||
let mut index = setup_test_index(size);
|
||||
let queries = generate_random_vectors(50, PERCH_EMBEDDING_DIM);
|
||||
let k = 10;
|
||||
|
||||
for ef in [64, 128, 256, 512] {
|
||||
index.set_ef_search(ef);
|
||||
|
||||
group.throughput(Throughput::Elements(queries.len() as u64));
|
||||
group.bench_with_input(BenchmarkId::new("ef", ef), &ef, |b, _| {
|
||||
b.iter(|| {
|
||||
for query in &queries {
|
||||
black_box(index.search(query, k));
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark HNSW vs brute force to calculate speedup ratio
|
||||
fn benchmark_hnsw_vs_brute_force(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_vs_brute_force");
|
||||
group.sample_size(20);
|
||||
group.measurement_time(Duration::from_secs(15));
|
||||
|
||||
// Use smaller index for brute force comparison
|
||||
let size = 10_000;
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
|
||||
let queries = generate_random_vectors(20, PERCH_EMBEDDING_DIM);
|
||||
let k = 10;
|
||||
|
||||
// Benchmark brute force
|
||||
group.bench_function("brute_force", |b| {
|
||||
b.iter(|| {
|
||||
for query in &queries {
|
||||
black_box(brute_force_knn(query, &vectors, k));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// Benchmark HNSW
|
||||
group.bench_function("hnsw", |b| {
|
||||
b.iter(|| {
|
||||
for query in &queries {
|
||||
black_box(index.search(query, k));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HNSW Insert Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark single vector insertion
|
||||
fn benchmark_hnsw_insert_single(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_insert_single");
|
||||
group.sample_size(50);
|
||||
group.measurement_time(Duration::from_secs(10));
|
||||
|
||||
// Benchmark insertion into indices of different sizes
|
||||
for &initial_size in &[1000, 10_000, 50_000] {
|
||||
let vectors_to_insert = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("initial_size", initial_size),
|
||||
&initial_size,
|
||||
|b, &size| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
// Setup: create index with initial vectors
|
||||
setup_test_index(size)
|
||||
},
|
||||
|mut index| {
|
||||
// Insert new vectors
|
||||
for vec in &vectors_to_insert {
|
||||
black_box(index.add(vec.clone()));
|
||||
}
|
||||
},
|
||||
criterion::BatchSize::SmallInput,
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark batch vector insertion
|
||||
fn benchmark_hnsw_insert_batch(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_insert_batch");
|
||||
group.sample_size(20);
|
||||
group.measurement_time(Duration::from_secs(15));
|
||||
|
||||
for &batch_size in &[100, 1000, 5000] {
|
||||
let vectors = generate_random_vectors(batch_size, PERCH_EMBEDDING_DIM);
|
||||
|
||||
group.throughput(Throughput::Elements(batch_size as u64));
|
||||
group.bench_with_input(
|
||||
BenchmarkId::new("batch_size", batch_size),
|
||||
&batch_size,
|
||||
|b, _| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
// Setup: create empty index
|
||||
SimpleHnswIndex::new_default()
|
||||
},
|
||||
|mut index| {
|
||||
// Insert batch
|
||||
black_box(index.batch_add(vectors.clone()));
|
||||
},
|
||||
criterion::BatchSize::SmallInput,
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HNSW Build Benchmarks
|
||||
// ============================================================================
|
||||
|
||||
/// Benchmark index construction time
|
||||
fn benchmark_hnsw_build(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_build");
|
||||
group.sample_size(10);
|
||||
group.measurement_time(Duration::from_secs(30));
|
||||
|
||||
for &size in &[1000, 5000, 10_000] {
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
|
||||
group.throughput(Throughput::Elements(size as u64));
|
||||
group.bench_with_input(BenchmarkId::new("vectors", size), &size, |b, _| {
|
||||
b.iter(|| {
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
black_box(index)
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
/// Benchmark index construction with different M parameters
|
||||
fn benchmark_hnsw_build_m_param(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hnsw_build_m_param");
|
||||
group.sample_size(10);
|
||||
group.measurement_time(Duration::from_secs(20));
|
||||
|
||||
let size = 5000;
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
|
||||
for m in [16, 24, 32, 48] {
|
||||
group.bench_with_input(BenchmarkId::new("M", m), &m, |b, &m| {
|
||||
b.iter(|| {
|
||||
let mut index =
|
||||
SimpleHnswIndex::new(PERCH_EMBEDDING_DIM, m, DEFAULT_EF_CONSTRUCTION, DEFAULT_EF_SEARCH);
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
black_box(index)
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Recall Measurement
|
||||
// ============================================================================
|
||||
|
||||
/// Measure and report recall metrics (not a benchmark, but a validation)
|
||||
fn measure_recall(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("recall_measurement");
|
||||
group.sample_size(10);
|
||||
|
||||
let size = 10_000;
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
|
||||
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
|
||||
// This benchmark measures time to compute recall (including brute force)
|
||||
group.bench_function("recall_computation", |b| {
|
||||
b.iter(|| {
|
||||
let mut total_recall_10 = 0.0;
|
||||
let mut total_recall_100 = 0.0;
|
||||
|
||||
for query in &queries {
|
||||
let hnsw_results = index.search(query, 100);
|
||||
let ground_truth = brute_force_knn(query, &vectors, 100);
|
||||
|
||||
total_recall_10 += measure_recall_at_k(&hnsw_results, &ground_truth, 10);
|
||||
total_recall_100 += measure_recall_at_k(&hnsw_results, &ground_truth, 100);
|
||||
}
|
||||
|
||||
let avg_recall_10 = total_recall_10 / queries.len() as f32;
|
||||
let avg_recall_100 = total_recall_100 / queries.len() as f32;
|
||||
|
||||
black_box((avg_recall_10, avg_recall_100))
|
||||
});
|
||||
});
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Speedup Ratio Calculation
|
||||
// ============================================================================
|
||||
|
||||
/// Calculate and report the speedup ratio of HNSW vs brute force
|
||||
/// This is run as a single iteration with detailed output
|
||||
fn calculate_speedup_ratio() {
|
||||
println!("\n=== HNSW vs Brute Force Speedup Analysis ===\n");
|
||||
|
||||
for &size in &[1_000, 5_000, 10_000, 50_000] {
|
||||
println!("Index size: {} vectors", size);
|
||||
println!("Dimension: {}", PERCH_EMBEDDING_DIM);
|
||||
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
|
||||
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
||||
let k = 10;
|
||||
|
||||
// Time brute force
|
||||
let bf_start = Instant::now();
|
||||
for query in &queries {
|
||||
let _ = brute_force_knn(query, &vectors, k);
|
||||
}
|
||||
let bf_time = bf_start.elapsed();
|
||||
|
||||
// Time HNSW
|
||||
let hnsw_start = Instant::now();
|
||||
for query in &queries {
|
||||
let _ = index.search(query, k);
|
||||
}
|
||||
let hnsw_time = hnsw_start.elapsed();
|
||||
|
||||
let speedup = bf_time.as_secs_f64() / hnsw_time.as_secs_f64();
|
||||
|
||||
// Calculate recall
|
||||
let mut total_recall = 0.0;
|
||||
for query in &queries {
|
||||
let hnsw_results = index.search(query, k);
|
||||
let ground_truth = brute_force_knn(query, &vectors, k);
|
||||
total_recall += measure_recall_at_k(&hnsw_results, &ground_truth, k);
|
||||
}
|
||||
let avg_recall = total_recall / queries.len() as f32;
|
||||
|
||||
println!(" Brute Force: {:?} ({} queries)", bf_time, queries.len());
|
||||
println!(" HNSW: {:?} ({} queries)", hnsw_time, queries.len());
|
||||
println!(" Speedup: {:.1}x", speedup);
|
||||
println!(" Recall@{}: {:.3}", k, avg_recall);
|
||||
println!(
|
||||
" Target: {}x speedup ({})",
|
||||
targets::HNSW_SPEEDUP_VS_BRUTE_FORCE,
|
||||
if speedup >= targets::HNSW_SPEEDUP_VS_BRUTE_FORCE {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Latency Distribution Analysis
|
||||
// ============================================================================
|
||||
|
||||
/// Analyze query latency distribution
|
||||
fn analyze_latency_distribution() {
|
||||
println!("\n=== Query Latency Distribution Analysis ===\n");
|
||||
|
||||
let size = MEDIUM_INDEX;
|
||||
println!("Building index with {} vectors...", size);
|
||||
let index = setup_test_index(size);
|
||||
|
||||
let queries = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
let k = 10;
|
||||
|
||||
let mut latencies = Vec::with_capacity(queries.len());
|
||||
|
||||
for query in &queries {
|
||||
let start = Instant::now();
|
||||
let _ = index.search(query, k);
|
||||
latencies.push(start.elapsed());
|
||||
}
|
||||
|
||||
let stats = PerformanceStats::from_latencies(latencies);
|
||||
|
||||
println!("Query latency statistics (k={}, {} queries):", k, queries.len());
|
||||
println!("{}", stats.report());
|
||||
println!();
|
||||
println!("Performance targets:");
|
||||
println!(
|
||||
" p50 target: {}ms ({})",
|
||||
targets::QUERY_LATENCY_P50_MS,
|
||||
if stats.p50 <= Duration::from_millis(targets::QUERY_LATENCY_P50_MS) {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
println!(
|
||||
" p99 target: {}ms ({})",
|
||||
targets::QUERY_LATENCY_P99_MS,
|
||||
if stats.p99 <= Duration::from_millis(targets::QUERY_LATENCY_P99_MS) {
|
||||
"PASS"
|
||||
} else {
|
||||
"FAIL"
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Criterion Groups
|
||||
// ============================================================================
|
||||
|
||||
criterion_group!(
|
||||
name = search_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_hnsw_search, benchmark_hnsw_search_ef, benchmark_hnsw_vs_brute_force
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = insert_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_hnsw_insert_single, benchmark_hnsw_insert_batch
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = build_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = benchmark_hnsw_build, benchmark_hnsw_build_m_param
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
name = recall_benches;
|
||||
config = Criterion::default().with_output_color(true);
|
||||
targets = measure_recall
|
||||
);
|
||||
|
||||
criterion_main!(search_benches, insert_benches, build_benches, recall_benches);
|
||||
|
||||
// ============================================================================
|
||||
// Additional Analysis Functions (run separately)
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod analysis {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
#[ignore] // Run with: cargo test --release -- --ignored --nocapture
|
||||
fn run_speedup_analysis() {
|
||||
calculate_speedup_ratio();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn run_latency_analysis() {
|
||||
analyze_latency_distribution();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_target_recall_at_10() {
|
||||
let size = 5_000;
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
|
||||
let queries = generate_random_vectors(50, PERCH_EMBEDDING_DIM);
|
||||
|
||||
let mut total_recall = 0.0;
|
||||
for query in &queries {
|
||||
let hnsw_results = index.search(query, 10);
|
||||
let ground_truth = brute_force_knn(query, &vectors, 10);
|
||||
total_recall += measure_recall_at_k(&hnsw_results, &ground_truth, 10);
|
||||
}
|
||||
|
||||
let avg_recall = total_recall / queries.len() as f32;
|
||||
println!("Average Recall@10: {:.3}", avg_recall);
|
||||
assert!(
|
||||
avg_recall as f64 >= targets::RECALL_AT_10,
|
||||
"Recall@10 {} below target {}",
|
||||
avg_recall,
|
||||
targets::RECALL_AT_10
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_throughput() {
|
||||
let vectors = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
let start = Instant::now();
|
||||
for vec in &vectors {
|
||||
index.add(vec.clone());
|
||||
}
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
let throughput = vectors.len() as f64 / elapsed.as_secs_f64();
|
||||
println!("Insert throughput: {:.0} vectors/sec", throughput);
|
||||
|
||||
// Note: This is a simplified index, real HNSW should achieve higher throughput
|
||||
}
|
||||
}
|
||||
673
examples/vibecast-7sense/benches/utils.rs
Normal file
673
examples/vibecast-7sense/benches/utils.rs
Normal file
@@ -0,0 +1,673 @@
|
||||
//! Benchmark Utilities for 7sense Performance Testing
|
||||
//!
|
||||
//! This module provides common utilities for benchmarking:
|
||||
//! - Random vector generation
|
||||
//! - Test index setup
|
||||
//! - Recall calculation
|
||||
//! - Ground truth computation
|
||||
//! - Performance metrics
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Embedding dimensions for Perch 2.0 model
|
||||
pub const PERCH_EMBEDDING_DIM: usize = 1536;
|
||||
|
||||
/// Default HNSW parameters from ADR-004
|
||||
pub const DEFAULT_M: usize = 32;
|
||||
pub const DEFAULT_EF_CONSTRUCTION: usize = 200;
|
||||
pub const DEFAULT_EF_SEARCH: usize = 128;
|
||||
pub const HIGH_RECALL_EF_SEARCH: usize = 256;
|
||||
|
||||
/// Performance targets from ADR-004
|
||||
pub mod targets {
|
||||
use std::time::Duration;
|
||||
|
||||
/// HNSW Search Targets
|
||||
pub const HNSW_SPEEDUP_VS_BRUTE_FORCE: f64 = 150.0;
|
||||
pub const QUERY_LATENCY_P50_MS: u64 = 10;
|
||||
pub const QUERY_LATENCY_P99_MS: u64 = 50;
|
||||
pub const RECALL_AT_10: f64 = 0.95;
|
||||
pub const RECALL_AT_100: f64 = 0.98;
|
||||
|
||||
/// Embedding Inference Targets
|
||||
pub const EMBEDDING_SEGMENTS_PER_SECOND: u64 = 100;
|
||||
|
||||
/// Batch Ingestion Targets
|
||||
pub const BATCH_VECTORS_PER_MINUTE: u64 = 1_000_000;
|
||||
pub const INSERT_THROUGHPUT_PER_SECOND: u64 = 10_000;
|
||||
|
||||
/// Query Latency Targets
|
||||
pub const TOTAL_QUERY_LATENCY_MS: u64 = 100;
|
||||
|
||||
/// Build Time Targets
|
||||
pub const BUILD_TIME_1M_VECTORS: Duration = Duration::from_secs(30 * 60);
|
||||
|
||||
/// Quantization Targets
|
||||
pub const MAX_RECALL_LOSS_INT8: f64 = 0.03;
|
||||
}
|
||||
|
||||
/// Generate random f32 vectors for benchmarking
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `count` - Number of vectors to generate
|
||||
/// * `dims` - Dimensionality of each vector
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of random f32 vectors, normalized to unit length
|
||||
pub fn generate_random_vectors(count: usize, dims: usize) -> Vec<Vec<f32>> {
|
||||
use std::f32::consts::PI;
|
||||
|
||||
let mut vectors = Vec::with_capacity(count);
|
||||
|
||||
for i in 0..count {
|
||||
let mut vec = Vec::with_capacity(dims);
|
||||
|
||||
// Use a simple deterministic random generator for reproducibility
|
||||
let mut seed = (i as u64).wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
|
||||
for _ in 0..dims {
|
||||
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
let val = ((seed >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0;
|
||||
vec.push(val);
|
||||
}
|
||||
|
||||
// Normalize to unit length
|
||||
let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in vec.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
|
||||
vectors.push(vec);
|
||||
}
|
||||
|
||||
vectors
|
||||
}
|
||||
|
||||
/// Generate clustered random vectors for more realistic benchmarking
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `count` - Total number of vectors to generate
|
||||
/// * `dims` - Dimensionality of each vector
|
||||
/// * `num_clusters` - Number of clusters to create
|
||||
/// * `cluster_spread` - Standard deviation within clusters (0.0 to 1.0)
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of random f32 vectors organized around cluster centers
|
||||
pub fn generate_clustered_vectors(
|
||||
count: usize,
|
||||
dims: usize,
|
||||
num_clusters: usize,
|
||||
cluster_spread: f32,
|
||||
) -> Vec<Vec<f32>> {
|
||||
let mut vectors = Vec::with_capacity(count);
|
||||
|
||||
// Generate cluster centers
|
||||
let centers = generate_random_vectors(num_clusters, dims);
|
||||
|
||||
// Assign vectors to clusters
|
||||
for i in 0..count {
|
||||
let cluster_idx = i % num_clusters;
|
||||
let center = ¢ers[cluster_idx];
|
||||
|
||||
let mut vec = Vec::with_capacity(dims);
|
||||
|
||||
// Use deterministic random for offset
|
||||
let mut seed = (i as u64).wrapping_mul(2862933555777941757).wrapping_add(3);
|
||||
|
||||
for d in 0..dims {
|
||||
seed = seed.wrapping_mul(2862933555777941757).wrapping_add(3);
|
||||
let noise = ((seed >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0;
|
||||
let val = center[d] + noise * cluster_spread;
|
||||
vec.push(val);
|
||||
}
|
||||
|
||||
// Normalize to unit length
|
||||
let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in vec.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
|
||||
vectors.push(vec);
|
||||
}
|
||||
|
||||
vectors
|
||||
}
|
||||
|
||||
/// Compute L2 (Euclidean) distance between two vectors
|
||||
#[inline]
|
||||
pub fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
|
||||
debug_assert_eq!(a.len(), b.len());
|
||||
a.iter()
|
||||
.zip(b.iter())
|
||||
.map(|(x, y)| (x - y) * (x - y))
|
||||
.sum::<f32>()
|
||||
.sqrt()
|
||||
}
|
||||
|
||||
/// Compute L2 squared distance (faster, no sqrt)
|
||||
#[inline]
|
||||
pub fn l2_distance_squared(a: &[f32], b: &[f32]) -> f32 {
|
||||
debug_assert_eq!(a.len(), b.len());
|
||||
a.iter()
|
||||
.zip(b.iter())
|
||||
.map(|(x, y)| (x - y) * (x - y))
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Compute cosine similarity between two vectors
|
||||
#[inline]
|
||||
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
debug_assert_eq!(a.len(), b.len());
|
||||
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
||||
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm_a > 0.0 && norm_b > 0.0 {
|
||||
dot / (norm_a * norm_b)
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute brute-force k-nearest neighbors (ground truth)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `query` - Query vector
|
||||
/// * `dataset` - Dataset of vectors to search
|
||||
/// * `k` - Number of neighbors to find
|
||||
///
|
||||
/// # Returns
|
||||
/// Vector of (index, distance) pairs sorted by distance
|
||||
pub fn brute_force_knn(query: &[f32], dataset: &[Vec<f32>], k: usize) -> Vec<(usize, f32)> {
|
||||
let mut distances: Vec<(usize, f32)> = dataset
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, vec)| (i, l2_distance(query, vec)))
|
||||
.collect();
|
||||
|
||||
distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
distances.truncate(k);
|
||||
distances
|
||||
}
|
||||
|
||||
/// Measure recall@k for approximate nearest neighbor results
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `results` - Approximate results (index, distance) pairs
|
||||
/// * `ground_truth` - Exact brute-force results (index, distance) pairs
|
||||
/// * `k` - Number of top results to consider
|
||||
///
|
||||
/// # Returns
|
||||
/// Recall value between 0.0 and 1.0
|
||||
pub fn measure_recall_at_k(
|
||||
results: &[(usize, f32)],
|
||||
ground_truth: &[(usize, f32)],
|
||||
k: usize,
|
||||
) -> f32 {
|
||||
let k = k.min(results.len()).min(ground_truth.len());
|
||||
if k == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let result_set: HashSet<usize> = results.iter().take(k).map(|(idx, _)| *idx).collect();
|
||||
let truth_set: HashSet<usize> = ground_truth.iter().take(k).map(|(idx, _)| *idx).collect();
|
||||
|
||||
let intersection = result_set.intersection(&truth_set).count();
|
||||
intersection as f32 / k as f32
|
||||
}
|
||||
|
||||
/// Calculate percentile from a sorted slice of durations
|
||||
pub fn percentile(sorted_latencies: &[Duration], p: f64) -> Duration {
|
||||
if sorted_latencies.is_empty() {
|
||||
return Duration::ZERO;
|
||||
}
|
||||
let idx = ((sorted_latencies.len() as f64 - 1.0) * p / 100.0).round() as usize;
|
||||
sorted_latencies[idx.min(sorted_latencies.len() - 1)]
|
||||
}
|
||||
|
||||
/// Performance statistics from benchmark runs
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PerformanceStats {
|
||||
pub count: usize,
|
||||
pub total_time: Duration,
|
||||
pub min: Duration,
|
||||
pub max: Duration,
|
||||
pub mean: Duration,
|
||||
pub p50: Duration,
|
||||
pub p95: Duration,
|
||||
pub p99: Duration,
|
||||
pub p999: Duration,
|
||||
pub throughput_per_sec: f64,
|
||||
}
|
||||
|
||||
impl PerformanceStats {
|
||||
/// Calculate statistics from a collection of latency measurements
|
||||
pub fn from_latencies(mut latencies: Vec<Duration>) -> Self {
|
||||
if latencies.is_empty() {
|
||||
return Self {
|
||||
count: 0,
|
||||
total_time: Duration::ZERO,
|
||||
min: Duration::ZERO,
|
||||
max: Duration::ZERO,
|
||||
mean: Duration::ZERO,
|
||||
p50: Duration::ZERO,
|
||||
p95: Duration::ZERO,
|
||||
p99: Duration::ZERO,
|
||||
p999: Duration::ZERO,
|
||||
throughput_per_sec: 0.0,
|
||||
};
|
||||
}
|
||||
|
||||
latencies.sort();
|
||||
|
||||
let total_time: Duration = latencies.iter().sum();
|
||||
let count = latencies.len();
|
||||
let mean = total_time / count as u32;
|
||||
|
||||
Self {
|
||||
count,
|
||||
total_time,
|
||||
min: latencies[0],
|
||||
max: latencies[count - 1],
|
||||
mean,
|
||||
p50: percentile(&latencies, 50.0),
|
||||
p95: percentile(&latencies, 95.0),
|
||||
p99: percentile(&latencies, 99.0),
|
||||
p999: percentile(&latencies, 99.9),
|
||||
throughput_per_sec: count as f64 / total_time.as_secs_f64(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if stats meet p99 latency target
|
||||
pub fn meets_p99_target(&self, target_ms: u64) -> bool {
|
||||
self.p99 <= Duration::from_millis(target_ms)
|
||||
}
|
||||
|
||||
/// Check if stats meet throughput target
|
||||
pub fn meets_throughput_target(&self, target_per_sec: u64) -> bool {
|
||||
self.throughput_per_sec >= target_per_sec as f64
|
||||
}
|
||||
|
||||
/// Format as a readable report
|
||||
pub fn report(&self) -> String {
|
||||
format!(
|
||||
"Count: {}\n\
|
||||
Total Time: {:?}\n\
|
||||
Min: {:?}\n\
|
||||
Max: {:?}\n\
|
||||
Mean: {:?}\n\
|
||||
P50: {:?}\n\
|
||||
P95: {:?}\n\
|
||||
P99: {:?}\n\
|
||||
P99.9: {:?}\n\
|
||||
Throughput: {:.2} ops/sec",
|
||||
self.count,
|
||||
self.total_time,
|
||||
self.min,
|
||||
self.max,
|
||||
self.mean,
|
||||
self.p50,
|
||||
self.p95,
|
||||
self.p99,
|
||||
self.p999,
|
||||
self.throughput_per_sec
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple HNSW-like index for benchmarking
|
||||
/// This is a simplified implementation for benchmark purposes
|
||||
pub struct SimpleHnswIndex {
|
||||
vectors: Vec<Vec<f32>>,
|
||||
dims: usize,
|
||||
m: usize,
|
||||
ef_construction: usize,
|
||||
ef_search: usize,
|
||||
// Simplified graph structure: each vector has a list of neighbor indices
|
||||
graph: Vec<Vec<usize>>,
|
||||
}
|
||||
|
||||
impl SimpleHnswIndex {
|
||||
/// Create a new empty index
|
||||
pub fn new(dims: usize, m: usize, ef_construction: usize, ef_search: usize) -> Self {
|
||||
Self {
|
||||
vectors: Vec::new(),
|
||||
dims,
|
||||
m,
|
||||
ef_construction,
|
||||
ef_search,
|
||||
graph: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an index with default parameters for Perch embeddings
|
||||
pub fn new_default() -> Self {
|
||||
Self::new(
|
||||
PERCH_EMBEDDING_DIM,
|
||||
DEFAULT_M,
|
||||
DEFAULT_EF_CONSTRUCTION,
|
||||
DEFAULT_EF_SEARCH,
|
||||
)
|
||||
}
|
||||
|
||||
/// Get the number of vectors in the index
|
||||
pub fn len(&self) -> usize {
|
||||
self.vectors.len()
|
||||
}
|
||||
|
||||
/// Check if the index is empty
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.vectors.is_empty()
|
||||
}
|
||||
|
||||
/// Add a single vector to the index
|
||||
pub fn add(&mut self, vector: Vec<f32>) -> usize {
|
||||
assert_eq!(vector.len(), self.dims);
|
||||
let id = self.vectors.len();
|
||||
|
||||
// Find neighbors for the new vector
|
||||
let neighbors = if self.vectors.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
self.search_internal(&vector, self.m.min(self.vectors.len()))
|
||||
.into_iter()
|
||||
.map(|(idx, _)| idx)
|
||||
.collect()
|
||||
};
|
||||
|
||||
self.vectors.push(vector);
|
||||
self.graph.push(neighbors.clone());
|
||||
|
||||
// Update bidirectional connections
|
||||
for &neighbor_id in &neighbors {
|
||||
if self.graph[neighbor_id].len() < self.m * 2 {
|
||||
self.graph[neighbor_id].push(id);
|
||||
}
|
||||
}
|
||||
|
||||
id
|
||||
}
|
||||
|
||||
/// Batch add vectors to the index
|
||||
pub fn batch_add(&mut self, vectors: Vec<Vec<f32>>) -> Vec<usize> {
|
||||
vectors.into_iter().map(|v| self.add(v)).collect()
|
||||
}
|
||||
|
||||
/// Search for k nearest neighbors
|
||||
pub fn search(&self, query: &[f32], k: usize) -> Vec<(usize, f32)> {
|
||||
assert_eq!(query.len(), self.dims);
|
||||
if self.vectors.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
self.search_internal(query, k)
|
||||
}
|
||||
|
||||
/// Internal search implementation with simplified HNSW-like traversal
|
||||
fn search_internal(&self, query: &[f32], k: usize) -> Vec<(usize, f32)> {
|
||||
use std::collections::{BinaryHeap, HashSet};
|
||||
use std::cmp::Reverse;
|
||||
|
||||
let ef = self.ef_search.max(k);
|
||||
|
||||
// Start from a random entry point
|
||||
let entry_point = 0;
|
||||
|
||||
let mut visited: HashSet<usize> = HashSet::new();
|
||||
let mut candidates: BinaryHeap<Reverse<(ordered_float::OrderedFloat<f32>, usize)>> =
|
||||
BinaryHeap::new();
|
||||
let mut results: BinaryHeap<(ordered_float::OrderedFloat<f32>, usize)> = BinaryHeap::new();
|
||||
|
||||
let entry_dist = l2_distance(query, &self.vectors[entry_point]);
|
||||
candidates.push(Reverse((ordered_float::OrderedFloat(entry_dist), entry_point)));
|
||||
results.push((ordered_float::OrderedFloat(entry_dist), entry_point));
|
||||
visited.insert(entry_point);
|
||||
|
||||
while let Some(Reverse((dist, current))) = candidates.pop() {
|
||||
let worst_dist = if results.len() >= ef {
|
||||
results.peek().map(|(d, _)| d.0).unwrap_or(f32::MAX)
|
||||
} else {
|
||||
f32::MAX
|
||||
};
|
||||
|
||||
if dist.0 > worst_dist {
|
||||
break;
|
||||
}
|
||||
|
||||
// Explore neighbors
|
||||
for &neighbor in &self.graph[current] {
|
||||
if visited.insert(neighbor) {
|
||||
let neighbor_dist = l2_distance(query, &self.vectors[neighbor]);
|
||||
|
||||
if results.len() < ef || neighbor_dist < worst_dist {
|
||||
candidates.push(Reverse((
|
||||
ordered_float::OrderedFloat(neighbor_dist),
|
||||
neighbor,
|
||||
)));
|
||||
results.push((ordered_float::OrderedFloat(neighbor_dist), neighbor));
|
||||
|
||||
if results.len() > ef {
|
||||
results.pop();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Convert to output format and sort by distance
|
||||
let mut output: Vec<(usize, f32)> =
|
||||
results.into_iter().map(|(d, idx)| (idx, d.0)).collect();
|
||||
output.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
output.truncate(k);
|
||||
output
|
||||
}
|
||||
|
||||
/// Set ef_search parameter for queries
|
||||
pub fn set_ef_search(&mut self, ef: usize) {
|
||||
self.ef_search = ef;
|
||||
}
|
||||
}
|
||||
|
||||
/// Setup a test index with the specified number of vectors
|
||||
pub fn setup_test_index(size: usize) -> SimpleHnswIndex {
|
||||
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
||||
let mut index = SimpleHnswIndex::new_default();
|
||||
|
||||
for vec in vectors {
|
||||
index.add(vec);
|
||||
}
|
||||
|
||||
index
|
||||
}
|
||||
|
||||
/// Scalar quantizer for int8 compression
|
||||
pub struct ScalarQuantizer {
|
||||
mins: Vec<f32>,
|
||||
maxs: Vec<f32>,
|
||||
scales: Vec<f32>,
|
||||
dims: usize,
|
||||
}
|
||||
|
||||
impl ScalarQuantizer {
|
||||
/// Create a new quantizer for the specified dimensions
|
||||
pub fn new(dims: usize) -> Self {
|
||||
Self {
|
||||
mins: vec![f32::MAX; dims],
|
||||
maxs: vec![f32::MIN; dims],
|
||||
scales: vec![1.0; dims],
|
||||
dims,
|
||||
}
|
||||
}
|
||||
|
||||
/// Calibrate the quantizer from a sample of embeddings
|
||||
pub fn calibrate(&mut self, embeddings: &[Vec<f32>]) {
|
||||
// Find min/max per dimension
|
||||
for embedding in embeddings {
|
||||
for (d, &val) in embedding.iter().enumerate() {
|
||||
if val < self.mins[d] {
|
||||
self.mins[d] = val;
|
||||
}
|
||||
if val > self.maxs[d] {
|
||||
self.maxs[d] = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Compute scales
|
||||
for d in 0..self.dims {
|
||||
let range = self.maxs[d] - self.mins[d];
|
||||
if range > 0.0 {
|
||||
self.scales[d] = 255.0 / range;
|
||||
} else {
|
||||
self.scales[d] = 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantize a float32 embedding to int8
|
||||
pub fn quantize(&self, embedding: &[f32]) -> Vec<u8> {
|
||||
embedding
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(d, &val)| {
|
||||
let normalized = (val - self.mins[d]) * self.scales[d];
|
||||
normalized.round().clamp(0.0, 255.0) as u8
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Dequantize an int8 embedding back to float32
|
||||
pub fn dequantize(&self, quantized: &[u8]) -> Vec<f32> {
|
||||
quantized
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(d, &val)| (val as f32) / self.scales[d] + self.mins[d])
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Timer utility for measuring operations
|
||||
pub struct Timer {
|
||||
start: Instant,
|
||||
}
|
||||
|
||||
impl Timer {
|
||||
/// Start a new timer
|
||||
pub fn start() -> Self {
|
||||
Self {
|
||||
start: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get elapsed time
|
||||
pub fn elapsed(&self) -> Duration {
|
||||
self.start.elapsed()
|
||||
}
|
||||
|
||||
/// Stop and return elapsed time
|
||||
pub fn stop(self) -> Duration {
|
||||
self.start.elapsed()
|
||||
}
|
||||
}
|
||||
|
||||
/// Measure execution time of a closure
|
||||
pub fn measure_time<F, R>(f: F) -> (R, Duration)
|
||||
where
|
||||
F: FnOnce() -> R,
|
||||
{
|
||||
let start = Instant::now();
|
||||
let result = f();
|
||||
let duration = start.elapsed();
|
||||
(result, duration)
|
||||
}
|
||||
|
||||
/// Measure average execution time over multiple iterations
|
||||
pub fn measure_average<F>(iterations: usize, mut f: F) -> PerformanceStats
|
||||
where
|
||||
F: FnMut() -> (),
|
||||
{
|
||||
let latencies: Vec<Duration> = (0..iterations)
|
||||
.map(|_| {
|
||||
let start = Instant::now();
|
||||
f();
|
||||
start.elapsed()
|
||||
})
|
||||
.collect();
|
||||
|
||||
PerformanceStats::from_latencies(latencies)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_generate_random_vectors() {
|
||||
let vectors = generate_random_vectors(100, 1536);
|
||||
assert_eq!(vectors.len(), 100);
|
||||
assert_eq!(vectors[0].len(), 1536);
|
||||
|
||||
// Check normalization
|
||||
for vec in &vectors {
|
||||
let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
assert!((norm - 1.0).abs() < 1e-5);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_l2_distance() {
|
||||
let a = vec![1.0, 0.0, 0.0];
|
||||
let b = vec![0.0, 1.0, 0.0];
|
||||
let dist = l2_distance(&a, &b);
|
||||
assert!((dist - std::f32::consts::SQRT_2).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recall_at_k() {
|
||||
let results: Vec<(usize, f32)> = vec![(0, 0.1), (1, 0.2), (2, 0.3), (3, 0.4), (5, 0.5)];
|
||||
let ground_truth: Vec<(usize, f32)> =
|
||||
vec![(0, 0.1), (1, 0.2), (2, 0.3), (4, 0.4), (5, 0.5)];
|
||||
|
||||
let recall = measure_recall_at_k(&results, &ground_truth, 5);
|
||||
assert!((recall - 0.8).abs() < 1e-5); // 4 out of 5 match
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scalar_quantizer() {
|
||||
let vectors = generate_random_vectors(100, 128);
|
||||
let mut quantizer = ScalarQuantizer::new(128);
|
||||
quantizer.calibrate(&vectors);
|
||||
|
||||
for vec in &vectors {
|
||||
let quantized = quantizer.quantize(vec);
|
||||
let dequantized = quantizer.dequantize(&quantized);
|
||||
|
||||
// Check that dequantized is close to original
|
||||
let error: f32 = vec
|
||||
.iter()
|
||||
.zip(dequantized.iter())
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.sum::<f32>()
|
||||
/ vec.len() as f32;
|
||||
|
||||
assert!(error < 0.1); // Average error should be small
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_performance_stats() {
|
||||
let latencies: Vec<Duration> = (0..100)
|
||||
.map(|i| Duration::from_micros(100 + i * 10))
|
||||
.collect();
|
||||
|
||||
let stats = PerformanceStats::from_latencies(latencies);
|
||||
assert_eq!(stats.count, 100);
|
||||
assert!(stats.min <= stats.p50);
|
||||
assert!(stats.p50 <= stats.p95);
|
||||
assert!(stats.p95 <= stats.p99);
|
||||
assert!(stats.p99 <= stats.max);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user