feat: brute-force cosine vector index
This commit is contained in:
+28
-1
@@ -1 +1,28 @@
|
|||||||
// implemented in a later task
|
/// Exact cosine search. 90k×1024 f32 ≈ 370 MB held in RAM; a query
|
||||||
|
/// is 90k dot products, well under 50 ms. No ANN dependency (YAGNI).
|
||||||
|
pub struct VectorIndex {
|
||||||
|
vectors: Vec<Vec<f32>>,
|
||||||
|
norms: Vec<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorIndex {
|
||||||
|
pub fn from_vectors(vectors: Vec<Vec<f32>>) -> Self {
|
||||||
|
let norms = vectors.iter().map(|v| dot(v, v).sqrt().max(1e-8)).collect();
|
||||||
|
Self { vectors, norms }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// (entry_index, cosine_similarity) descending.
|
||||||
|
pub fn top_k(&self, query: &[f32], k: usize) -> Vec<(usize, f32)> {
|
||||||
|
let qn = dot(query, query).sqrt().max(1e-8);
|
||||||
|
let mut scored: Vec<(usize, f32)> = self.vectors.iter().enumerate()
|
||||||
|
.map(|(i, v)| (i, dot(v, query) / (self.norms[i] * qn)))
|
||||||
|
.collect();
|
||||||
|
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||||
|
scored.truncate(k);
|
||||||
|
scored
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dot(a: &[f32], b: &[f32]) -> f32 {
|
||||||
|
a.iter().zip(b).map(|(x, y)| x * y).sum()
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,14 @@
|
|||||||
|
use alpha_id::vector::VectorIndex;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cosine_topk_ranks_nearest_first() {
|
||||||
|
let vi = VectorIndex::from_vectors(vec![
|
||||||
|
vec![1.0, 0.0], // id 0
|
||||||
|
vec![0.0, 1.0], // id 1
|
||||||
|
vec![0.9, 0.1], // id 2
|
||||||
|
]);
|
||||||
|
let hits = vi.top_k(&[1.0, 0.0], 2);
|
||||||
|
assert_eq!(hits[0].0, 0);
|
||||||
|
assert_eq!(hits[1].0, 2);
|
||||||
|
assert!(hits[0].1 > hits[1].1);
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user