146 lines
4.2 KiB
Rust
146 lines
4.2 KiB
Rust
use std::time::Instant;
|
|
|
|
use rand::{thread_rng, Rng};
|
|
use tracing::{metadata::LevelFilter, subscriber::set_global_default};
|
|
use tracing_subscriber::{fmt::format::FmtSpan, layer::SubscriberExt, EnvFilter, Layer, Registry};
|
|
use vectordb::Vector;
|
|
|
|
fn init_tracing() -> Result<(), Box<dyn std::error::Error>> {
|
|
let format_layer = tracing_subscriber::fmt::layer()
|
|
.with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
|
|
.with_filter(
|
|
EnvFilter::builder()
|
|
.with_default_directive(LevelFilter::INFO.into())
|
|
.from_env_lossy(),
|
|
);
|
|
|
|
let subscriber = Registry::default().with(format_layer);
|
|
|
|
set_global_default(subscriber)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
init_tracing()?;
|
|
|
|
let db = vectordb::VectorDb::open("./vectordb", 2)?;
|
|
|
|
let total_chunks = 32;
|
|
let chunk_size = 1024 * 8;
|
|
|
|
for _ in 0..total_chunks {
|
|
let vectors = (0..chunk_size)
|
|
.map(|_| thread_rng().gen::<[f32; 7]>().into())
|
|
.collect::<Vec<_>>();
|
|
|
|
tracing::info!("Inserting vectors");
|
|
let now = Instant::now();
|
|
db.insert_vectors_blocking(vectors)?;
|
|
tracing::info!("Inserting vectors - Took {:?}", now.elapsed());
|
|
}
|
|
|
|
let existing_vector: Vector<7> = thread_rng().gen::<[f32; 7]>().into();
|
|
|
|
db.insert_vector_blocking(existing_vector.clone())?;
|
|
|
|
let test_vecs = [existing_vector]
|
|
.into_iter()
|
|
.chain((0..5).map(|_| thread_rng().gen::<[f32; 7]>().into()))
|
|
.collect::<Vec<_>>();
|
|
|
|
let span = tracing::info_span!("Similarities");
|
|
let guard = span.enter();
|
|
for (i, v) in test_vecs.iter().enumerate() {
|
|
let span = tracing::info_span!("test", id = %i);
|
|
let guard = span.enter();
|
|
|
|
let similarities = db.find_similarities_blocking(v.clone(), None, 5)?;
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
tracing::info!(
|
|
"similar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
|
|
drop(guard);
|
|
drop(span);
|
|
}
|
|
drop(guard);
|
|
drop(span);
|
|
|
|
let span = tracing::info_span!("Furthest similarities");
|
|
let guard = span.enter();
|
|
for (i, v) in test_vecs.iter().enumerate() {
|
|
let span = tracing::info_span!("test", id = %i);
|
|
let guard = span.enter();
|
|
|
|
let similarities = db.find_furthest_similarities_blocking(v.clone(), None, 5)?;
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
tracing::info!(
|
|
"similar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
|
|
drop(guard);
|
|
drop(span);
|
|
}
|
|
drop(guard);
|
|
drop(span);
|
|
|
|
let span = tracing::info_span!("Dissimilarities");
|
|
let guard = span.enter();
|
|
for (i, v) in test_vecs.iter().enumerate() {
|
|
let span = tracing::info_span!("test", id = %i);
|
|
let guard = span.enter();
|
|
|
|
let similarities = db.find_dissimilarities_blocking(v.clone(), None, 5)?;
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
tracing::info!(
|
|
"dissimilar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
|
|
drop(guard);
|
|
drop(span);
|
|
}
|
|
drop(guard);
|
|
drop(span);
|
|
|
|
let span = tracing::info_span!("Closest dissimilarities");
|
|
let guard = span.enter();
|
|
for (i, v) in test_vecs.iter().enumerate() {
|
|
let span = tracing::info_span!("test", id = %i);
|
|
let guard = span.enter();
|
|
|
|
let similarities = db.find_closest_dissimilarities_blocking(v.clone(), None, 5)?;
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
tracing::info!(
|
|
"dissimilar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
|
|
drop(guard);
|
|
drop(span);
|
|
}
|
|
drop(guard);
|
|
drop(span);
|
|
|
|
Ok(())
|
|
}
|