diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs index 3e7dc270d..37ade8f81 100644 --- a/crates/milli/src/vector/mod.rs +++ b/crates/milli/src/vector/mod.rs @@ -584,6 +584,7 @@ pub struct ArroyStats { pub documents: RoaringBitmap, } /// One or multiple embeddings stored consecutively in a flat vector. +#[derive(Debug, PartialEq)] pub struct Embeddings { data: Vec, dimension: usize, @@ -734,15 +735,26 @@ impl EmbeddingConfig { } } -/// Map of embedder configurations. -/// -/// Each configuration is mapped to a name. +/// Map of runtime embedder data. #[derive(Clone, Default)] -pub struct EmbeddingConfigs(HashMap, Arc, bool)>); +pub struct RuntimeEmbedders(HashMap>); -impl EmbeddingConfigs { +pub struct RuntimeEmbedder { + pub embedder: Arc, + pub document_template: Prompt, + pub fragments: Vec, + pub is_quantized: bool, +} + +pub struct RuntimeFragment { + pub name: String, + pub id: u8, + pub template: JsonTemplate, +} + +impl RuntimeEmbedders { /// Create the map from its internal component.s - pub fn new(data: HashMap, Arc, bool)>) -> Self { + pub fn new(data: HashMap>) -> Self { Self(data) } @@ -751,24 +763,23 @@ impl EmbeddingConfigs { } /// Get an embedder configuration and template from its name. - pub fn get(&self, name: &str) -> Option<(Arc, Arc, bool)> { + pub fn get(&self, name: &str) -> Option> { self.0.get(name).cloned() } - pub fn inner_as_ref(&self) -> &HashMap, Arc, bool)> { + pub fn inner_as_ref(&self) -> &HashMap> { &self.0 } - pub fn into_inner(self) -> HashMap, Arc, bool)> { + pub fn into_inner(self) -> HashMap> { self.0 } } -impl IntoIterator for EmbeddingConfigs { - type Item = (String, (Arc, Arc, bool)); +impl IntoIterator for RuntimeEmbedders { + type Item = (String, Arc); - type IntoIter = - std::collections::hash_map::IntoIter, Arc, bool)>; + type IntoIter = std::collections::hash_map::IntoIter>; fn into_iter(self) -> Self::IntoIter { self.0.into_iter()