diff --git a/inference_core/Cargo.toml b/inference_core/Cargo.toml index d31645c..3e24541 100644 --- a/inference_core/Cargo.toml +++ b/inference_core/Cargo.toml @@ -12,6 +12,7 @@ tokenizers = { version = "0.14.1", default-features = false, features = ["progre ndarray = "0.15.6" uniffi = { version = "0.25" } +uuid = { version = "1.5.0", features = ["v4"] } [lib] crate-type = ["lib", "cdylib"] diff --git a/inference_core/src/cosine_similarity.rs b/inference_core/src/cosine_similarity.rs new file mode 100644 index 0000000..c33dea1 --- /dev/null +++ b/inference_core/src/cosine_similarity.rs @@ -0,0 +1,33 @@ +use crate::embedding::Embedding; +use crate::similarity::Similarity; + +pub struct CosineSimilarity; + +impl Similarity for CosineSimilarity { + fn similarity_score(&self, vector_a: &Embedding, vector_b: &Embedding) -> f32 { + if vector_a.len() != vector_b.len() { + panic!( + "Length of vector a ({}) must be equal to the length of vector b ({})", + vector_a.len(), + vector_b.len() + ); + } + + let dot_product: f32 = vector_a + .iter() + .zip(vector_b.iter()) // Use vector_b.iter() directly + .map(|(a, b)| a * b) + .sum(); + + let norm_a: f32 = vector_a.iter().map(|x| x * x).sum(); + let norm_b: f32 = vector_b.iter().map(|x| x * x).sum(); + + dot_product / (f32::sqrt(norm_a) * f32::sqrt(norm_b)) + } +} + +impl CosineSimilarity { + pub fn between(embedding: &Embedding, reference_embedding: &Embedding) -> f32 { + CosineSimilarity.similarity_score(embedding, reference_embedding) + } +} diff --git a/inference_core/src/document.rs b/inference_core/src/document.rs index 4da73fb..54bc0f8 100644 --- a/inference_core/src/document.rs +++ b/inference_core/src/document.rs @@ -1,17 +1,26 @@ use std::collections::HashMap; use crate::embedding::Embedding; -struct Document { +pub struct Document { id: String, metadata: Metadata, text: String, vector: Embedding, } -struct Metadata { +pub struct Metadata { metadata: HashMap, } +impl Metadata { + pub fn new() -> Self { + Self { + metadata: HashMap::new(), + } + } +} + + impl Document { fn from(string_value: String) -> Self { Self { diff --git a/inference_core/src/embedding_match.rs b/inference_core/src/embedding_match.rs new file mode 100644 index 0000000..0728ce2 --- /dev/null +++ b/inference_core/src/embedding_match.rs @@ -0,0 +1,42 @@ +use std::cmp::Ordering; +use crate::embedding::Embedding; + +#[derive(Debug, Clone)] +pub struct EmbeddingMatch { + score: f32, + embedding_id: String, + embedding: Embedding, + embedded: Embedded, +} + +impl EmbeddingMatch { + pub(crate) fn new(score: f32, embedding_id: String, embedding: Embedding, embedded: Embedded) -> Self { + EmbeddingMatch { + score, + embedding_id, + embedding, + embedded, + } + } +} + +impl PartialOrd for EmbeddingMatch { + fn partial_cmp(&self, other: &Self) -> Option { + self.score.partial_cmp(&other.score) + } +} + +impl PartialEq for EmbeddingMatch { + fn eq(&self, other: &Self) -> bool { + self.score == other.score + } +} + +impl Ord for EmbeddingMatch { + fn cmp(&self, other: &Self) -> Ordering { + self.score.partial_cmp(&other.score).unwrap_or(Ordering::Equal) + } +} + +impl Eq for EmbeddingMatch { +} \ No newline at end of file diff --git a/inference_core/src/embedding_store.rs b/inference_core/src/embedding_store.rs new file mode 100644 index 0000000..3a5f09c --- /dev/null +++ b/inference_core/src/embedding_store.rs @@ -0,0 +1,28 @@ +use crate::embedding::Embedding; +use crate::embedding_match::EmbeddingMatch; + +trait EmbeddingStore { + // Adds an embedding to the store and returns its unique identifier. + fn add(&mut self, embedding: Embedding) -> String; + + // Adds an embedding to the store with a specified identifier. + fn add_with_id(&mut self, id: String, embedding: Embedding); + + // Adds an embedding to the store and associates it with the provided embedded data. + fn add_with_embedded(&mut self, embedding: Embedding, embedded: Embedded) -> String; + + // Adds a list of embeddings to the store and returns a list of unique identifiers. + fn add_all(&mut self, embeddings: Vec) -> Vec; + + // Adds a list of embeddings to the store and associates them with a list of embedded data. + fn add_all_with_embedded(&mut self, embeddings: Vec, embedded: Vec) -> Vec; + + // Find relevant embeddings in the store based on a reference embedding, with a maximum number of results. + // An optional minimum score can be specified to filter results. + fn find_relevant( + &self, + reference_embedding: Embedding, + max_results: usize, + min_score: f32, + ) -> Vec>; +} diff --git a/inference_core/src/lib.rs b/inference_core/src/lib.rs index aaa082a..32e7fe0 100644 --- a/inference_core/src/lib.rs +++ b/inference_core/src/lib.rs @@ -2,6 +2,11 @@ pub mod embed; pub mod memory_store; pub mod document; pub mod embedding; +pub mod embedding_store; +mod embedding_match; +mod similarity; +mod relevance_score; +mod cosine_similarity; fn hello_name(name: String) -> String { format!("Hello {}", name) diff --git a/inference_core/src/memory_store.rs b/inference_core/src/memory_store.rs index e69de29..39efea9 100644 --- a/inference_core/src/memory_store.rs +++ b/inference_core/src/memory_store.rs @@ -0,0 +1,85 @@ +use std::collections::BinaryHeap; + +use crate::cosine_similarity::CosineSimilarity; +use crate::embedding::Embedding; +use crate::embedding_match::EmbeddingMatch; +use crate::relevance_score::RelevanceScore; + + +#[derive(Clone)] +struct Entry { + id: String, + embedding: Embedding, + embedded: Option, +} + +impl Entry { + fn new(id: String, embedding: Embedding, embedded: Option) -> Self { + Entry { id, embedding, embedded } + } +} + +pub struct InMemoryEmbeddingStore { + entries: Vec>, +} + +// Implement methods for InMemoryEmbeddingStore +impl InMemoryEmbeddingStore { + fn new() -> Self { + InMemoryEmbeddingStore { entries: Vec::new() } + } + + fn add(&mut self, embedding: Embedding) -> String { + let id = uuid::Uuid::new_v4().to_string(); + self.add_with_id(id.clone(), embedding); + id + } + + fn add_with_id(&mut self, id: String, embedding: Embedding) { + self.add_with_embedded(id, embedding, None); + } + + fn add_with_embedded(&mut self, id: String, embedding: Embedding, embedded: Option) -> String { + let entry = Entry::new(id.clone(), embedding, embedded); + self.entries.push(entry); + id + } + + fn add_all(&mut self, embeddings: Vec) -> Vec { + embeddings + .into_iter() + .map(|embedding| self.add(embedding)) + .collect() + } + + fn add_all_with_embedded(&mut self, embeddings: Vec, embedded: Vec) -> Vec { + assert_eq!(embeddings.len(), embedded.len(), "The list of embeddings and embedded must have the same size"); + + embeddings + .into_iter() + .zip(embedded) + .map(|(embedding, embedded)| self.add_with_embedded(uuid::Uuid::new_v4().to_string(), embedding, Some(embedded))) + .collect() + } + + fn find_relevant(&self, reference_embedding: Embedding, max_results: usize, min_score: f32) -> Vec> { + let mut matches = BinaryHeap::new(); + + for entry in &self.entries { + let cosine_similarity = CosineSimilarity::between(&entry.embedding, &reference_embedding); + let score = RelevanceScore::from_cosine_similarity(cosine_similarity); + + if score >= min_score { + matches.push(EmbeddingMatch::new(score, entry.id.clone(), entry.embedding.clone(), entry.embedded.clone().unwrap())); + + if matches.len() > max_results { + matches.pop(); + } + } + } + + let mut result: Vec<_> = matches.into_sorted_vec(); + result.reverse(); + result + } +} diff --git a/inference_core/src/relevance_score.rs b/inference_core/src/relevance_score.rs new file mode 100644 index 0000000..2c6a3e3 --- /dev/null +++ b/inference_core/src/relevance_score.rs @@ -0,0 +1,7 @@ +pub struct RelevanceScore; + +impl RelevanceScore { + pub fn from_cosine_similarity(cosine_similarity: f32) -> f32 { + (cosine_similarity + 1.0) / 2.0 + } +} diff --git a/inference_core/src/similarity.rs b/inference_core/src/similarity.rs new file mode 100644 index 0000000..a91777c --- /dev/null +++ b/inference_core/src/similarity.rs @@ -0,0 +1,5 @@ +use crate::embedding::Embedding; + +pub trait Similarity { + fn similarity_score(&self, set1: &Embedding, set2: &Embedding) -> f32; +}