Correctly construct the Embeddings struct

This commit is contained in:
Clément Renault 2024-11-28 13:53:25 +01:00
parent 58eab9a018
commit cc4bd54669
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F
3 changed files with 21 additions and 8 deletions

View File

@ -293,6 +293,20 @@ impl ArroySetVectors {
});
Some(&vec[..])
}
/// Read all the embeddings and write them into an aligned `f32` Vec.
pub fn read_all_embeddings_into_vec<'v>(
&self,
frame: &FrameGrantR<'_>,
vec: &'v mut Vec<f32>,
) -> &'v [f32] {
vec.clear();
Self::remaining_bytes(frame).chunks_exact(mem::size_of::<f32>()).for_each(|bytes| {
let f = bytes.try_into().map(f32::from_ne_bytes).unwrap();
vec.push(f);
});
&vec[..]
}
}
#[derive(Debug, Clone, Copy)]

View File

@ -442,11 +442,12 @@ where
let LargeVectors { docid, embedder_id, .. } = large_vectors;
let (_, _, writer, dimensions) =
arroy_writers.get(&embedder_id).expect("requested a missing embedder");
writer.del_items(wtxn, *dimensions, docid)?;
let mut embeddings = Embeddings::new(*dimensions);
for embedding in large_vectors.read_embeddings() {
embeddings.push(embedding.to_vec()).unwrap();
}
writer.del_items(wtxn, *dimensions, docid)?;
writer.add_items(wtxn, docid, &embeddings)?;
}
}
@ -607,13 +608,11 @@ fn write_from_bbqueue(
let frame = frame_with_header.frame();
let (_, _, writer, dimensions) =
arroy_writers.get(&embedder_id).expect("requested a missing embedder");
let mut embeddings = Embeddings::new(*dimensions);
let all_embeddings = asvs.read_all_embeddings_into_vec(frame, aligned_embedding);
embeddings.append(all_embeddings.to_vec()).unwrap();
writer.del_items(wtxn, *dimensions, docid)?;
for index in 0.. {
match asvs.read_embedding_into_vec(frame, index, aligned_embedding) {
Some(embedding) => writer.add_item(wtxn, docid, embedding)?,
None => break,
}
}
writer.add_items(wtxn, docid, &embeddings)?;
}
}
}

View File

@ -475,7 +475,7 @@ impl<F> Embeddings<F> {
Ok(())
}
/// Append a flat vector of embeddings a the end of the embeddings.
/// Append a flat vector of embeddings at the end of the embeddings.
///
/// If `embeddings.len() % self.dimension != 0`, then the append operation fails.
pub fn append(&mut self, mut embeddings: Vec<F>) -> Result<(), Vec<F>> {