osobhandClaude Opus 5.5
d0db83812b
The setting was persisted in /meta and otherwise ignored: embeddings
were always written as f32. It now does what it says.
clawhdf5-format:
- `DatasetBuilder::with_f16_data` writes IEEE binary16 (numpy float16),
rounding to nearest-even, and `make_f16_type`.
- `clawhdf5_format::float16` holds the f32 <-> f16 conversions, the one
implementation the writer, the reader and the agent all use. Checked
against the `half` crate on 16.7M f32 values and round-trips all 65536
half values; the h5py interop tests confirm the rounding matches
numpy's bit for bit (4020 values incl. ties, subnormals, overflow).
- Reading little-endian float16 as f32 has a fast path.
clawhdf5-agent:
- A float16 store writes /memory/embeddings as half precision, and
`MemoryCache::half_precision` rounds each embedding as it enters the
cache (save, update, WAL replay, and on load of a store still f32 on
disk), so memory and file agree bit for bit and a store searches the
same before and after a reopen (tested).
- Values beyond +-65504 are refused with the new
`MemoryError::InvalidEntry` rather than stored as infinity, on every
save path; batches are all or nothing, and a rejected ephemeral entry
stays in the ephemeral tier. Breaking for exhaustive matches.
- CLI: `create --float16`. Off by default.
Measured on tank, 384-dim, six runs alternating order, medians
(search_harness --float16-study --full): at 100K the file goes from
154.0 to 80.8 MiB (-48%), checkpoint 752 -> 512 ms, open 300 -> 252 ms;
vector recall@10 against an exact scan and hybrid_search latency do not
change. At 10K open is 3 ms slower. Also a test that h5py opens a whole
agent store, f32 and float16, and decodes every dataset.
Docs: README, BENCHMARKS.md ("float16 embedding storage"), CHANGELOG
(including the h5py interop fixes in the previous commit), CLAUDE.md.
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
542 lines
17 KiB
Rust
542 lines
17 KiB
Rust
//! In-memory cache for memory entries, sessions, and knowledge graph.
|
|
|
|
use crate::vector_search;
|
|
use clawhdf5_format::float16::round_to_f16;
|
|
|
|
/// Every entry's embedding, in one contiguous `[N x dim]` buffer.
|
|
///
|
|
/// Rows are always exactly `dim` long: a shorter one is zero-padded, a longer
|
|
/// one truncated. The previous `Vec<Vec<f32>>` allowed ragged rows, which
|
|
/// silently misaligned the flattened copy that the batched kernels read — a
|
|
/// single wrong-length embedding shifted every row after it. Padding makes
|
|
/// that unrepresentable. A record stored without an embedding therefore holds
|
|
/// a zero row, and is told apart by its norm being zero rather than by length.
|
|
///
|
|
/// This used to be two fields — a `Vec<Vec<f32>>` and a flattened copy kept in
|
|
/// lock-step — which stored the whole corpus twice and cost one heap
|
|
/// allocation per entry on top. At 100k 384-dim entries that duplicate was
|
|
/// ~150 MiB. Indexing yields a `&[f32]` row, so `embeddings[i]` still reads
|
|
/// the same way.
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct Embeddings {
|
|
flat: Vec<f32>,
|
|
dim: usize,
|
|
}
|
|
|
|
impl Embeddings {
|
|
pub fn new(dim: usize) -> Self {
|
|
Self {
|
|
flat: Vec::new(),
|
|
dim,
|
|
}
|
|
}
|
|
|
|
/// Number of embeddings.
|
|
pub fn len(&self) -> usize {
|
|
self.flat.len().checked_div(self.dim).unwrap_or(0)
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.len() == 0
|
|
}
|
|
|
|
/// The whole buffer, `[N x dim]` row-major — what batched kernels read.
|
|
pub fn as_flat(&self) -> &[f32] {
|
|
&self.flat
|
|
}
|
|
|
|
pub fn dim(&self) -> usize {
|
|
self.dim
|
|
}
|
|
|
|
/// Row `i`, or `None` if out of range.
|
|
pub fn get(&self, i: usize) -> Option<&[f32]> {
|
|
let start = i.checked_mul(self.dim)?;
|
|
self.flat.get(start..start.checked_add(self.dim)?)
|
|
}
|
|
|
|
pub fn iter(&self) -> impl ExactSizeIterator<Item = &[f32]> {
|
|
self.flat.chunks_exact(self.dim.max(1))
|
|
}
|
|
|
|
/// Append one embedding. A row whose length doesn't match `dim` is padded
|
|
/// or truncated, so the buffer stays rectangular whatever a caller passes.
|
|
pub fn push(&mut self, embedding: &[f32]) {
|
|
if self.dim == 0 {
|
|
return;
|
|
}
|
|
let take = embedding.len().min(self.dim);
|
|
self.flat.extend_from_slice(&embedding[..take]);
|
|
self.flat.resize(self.flat.len() + (self.dim - take), 0.0);
|
|
}
|
|
|
|
/// Replace row `i`. Out-of-range indices are ignored.
|
|
pub fn set(&mut self, i: usize, embedding: &[f32]) {
|
|
let Some(start) = i.checked_mul(self.dim) else {
|
|
return;
|
|
};
|
|
if start + self.dim > self.flat.len() {
|
|
return;
|
|
}
|
|
let take = embedding.len().min(self.dim);
|
|
self.flat[start..start + take].copy_from_slice(&embedding[..take]);
|
|
self.flat[start + take..start + self.dim].fill(0.0);
|
|
}
|
|
|
|
/// Keep only the rows `keep` returns true for, preserving order.
|
|
pub fn retain(&mut self, mut keep: impl FnMut(usize) -> bool) {
|
|
if self.dim == 0 {
|
|
return;
|
|
}
|
|
let mut write = 0usize;
|
|
for read in 0..self.len() {
|
|
if keep(read) {
|
|
if write != read {
|
|
let (dst, src) = (write * self.dim, read * self.dim);
|
|
self.flat.copy_within(src..src + self.dim, dst);
|
|
}
|
|
write += 1;
|
|
}
|
|
}
|
|
self.flat.truncate(write * self.dim);
|
|
}
|
|
|
|
/// Replace the contents with `rows`.
|
|
pub fn reset_from(&mut self, dim: usize, rows: impl IntoIterator<Item = Vec<f32>>) {
|
|
self.dim = dim;
|
|
self.flat.clear();
|
|
for row in rows {
|
|
self.push(&row);
|
|
}
|
|
}
|
|
|
|
/// Adopt an already-flat buffer, trimming any partial trailing row.
|
|
pub fn set_flat(&mut self, dim: usize, mut flat: Vec<f32>) {
|
|
self.dim = dim;
|
|
match flat.len().checked_div(dim) {
|
|
Some(rows) => flat.truncate(rows * dim),
|
|
None => flat.clear(),
|
|
}
|
|
self.flat = flat;
|
|
}
|
|
}
|
|
|
|
impl PartialEq for Embeddings {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.dim == other.dim && self.flat == other.flat
|
|
}
|
|
}
|
|
|
|
impl std::ops::Index<usize> for Embeddings {
|
|
type Output = [f32];
|
|
|
|
fn index(&self, i: usize) -> &[f32] {
|
|
self.get(i).expect("embedding index out of range")
|
|
}
|
|
}
|
|
|
|
/// In-memory cache for the /memory group data.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MemoryCache {
|
|
pub chunks: Vec<String>,
|
|
pub embeddings: Embeddings,
|
|
pub source_channels: Vec<String>,
|
|
pub timestamps: Vec<f64>,
|
|
pub session_ids: Vec<String>,
|
|
pub tags: Vec<String>,
|
|
pub tombstones: Vec<u8>,
|
|
pub embedding_dim: usize,
|
|
/// Pre-computed L2 norms for each embedding.
|
|
pub norms: Vec<f32>,
|
|
/// Hebbian activation weights (default 1.0 per entry).
|
|
pub activation_weights: Vec<f32>,
|
|
/// Round every embedding to IEEE half precision as it enters the cache,
|
|
/// so the cache holds exactly what a `float16` store writes to disk. Set
|
|
/// it with [`MemoryCache::set_half_precision`], which also rounds the
|
|
/// rows already held.
|
|
pub half_precision: bool,
|
|
}
|
|
|
|
impl MemoryCache {
|
|
pub fn new(embedding_dim: usize) -> Self {
|
|
Self {
|
|
chunks: Vec::new(),
|
|
embeddings: Embeddings::new(embedding_dim),
|
|
source_channels: Vec::new(),
|
|
timestamps: Vec::new(),
|
|
session_ids: Vec::new(),
|
|
tags: Vec::new(),
|
|
tombstones: Vec::new(),
|
|
embedding_dim,
|
|
norms: Vec::new(),
|
|
activation_weights: Vec::new(),
|
|
half_precision: false,
|
|
}
|
|
}
|
|
|
|
/// Switch half-precision rounding on or off. Turning it on rounds every
|
|
/// embedding already held (and recomputes norms where one changed) —
|
|
/// e.g. a `float16` store whose last checkpoint predates half-precision
|
|
/// storage and so is still `f32` on disk.
|
|
pub fn set_half_precision(&mut self, on: bool) {
|
|
self.half_precision = on;
|
|
if !on {
|
|
return;
|
|
}
|
|
for i in 0..self.embeddings.len() {
|
|
let row = &self.embeddings[i];
|
|
if row
|
|
.iter()
|
|
.all(|&v| round_to_f16(v).to_bits() == v.to_bits())
|
|
{
|
|
continue;
|
|
}
|
|
let rounded: Vec<f32> = row.iter().map(|&v| round_to_f16(v)).collect();
|
|
self.norms[i] = vector_search::compute_norm(&rounded);
|
|
self.embeddings.set(i, &rounded);
|
|
}
|
|
}
|
|
|
|
/// The embedding as the cache will hold it: rounded to half precision
|
|
/// when [`Self::half_precision`] is on, otherwise unchanged.
|
|
fn stored_form(&self, mut embedding: Vec<f32>) -> Vec<f32> {
|
|
if self.half_precision {
|
|
for v in &mut embedding {
|
|
*v = round_to_f16(*v);
|
|
}
|
|
}
|
|
embedding
|
|
}
|
|
|
|
/// Kept for callers that used to have to re-flatten after a bulk load.
|
|
/// The buffer is always flat now, so there is nothing to rebuild.
|
|
#[deprecated(note = "embeddings are stored flat; this is a no-op")]
|
|
pub fn rebuild_flat(&mut self) {}
|
|
|
|
/// The embeddings as one contiguous `[N x dim]` buffer.
|
|
pub fn flat_embeddings(&self) -> &[f32] {
|
|
self.embeddings.as_flat()
|
|
}
|
|
|
|
/// Total number of entries (including tombstoned).
|
|
pub fn len(&self) -> usize {
|
|
self.chunks.len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.chunks.is_empty()
|
|
}
|
|
|
|
/// Number of active (non-tombstoned) entries.
|
|
pub fn count_active(&self) -> usize {
|
|
self.tombstones.iter().filter(|&&t| t == 0).count()
|
|
}
|
|
|
|
/// Push a new entry, returns its index.
|
|
pub fn push(
|
|
&mut self,
|
|
chunk: String,
|
|
embedding: Vec<f32>,
|
|
source_channel: String,
|
|
timestamp: f64,
|
|
session_id: String,
|
|
tags: String,
|
|
) -> usize {
|
|
let idx = self.chunks.len();
|
|
let embedding = self.stored_form(embedding);
|
|
let norm = vector_search::compute_norm(&embedding);
|
|
self.chunks.push(chunk);
|
|
self.embeddings.push(&embedding);
|
|
self.source_channels.push(source_channel);
|
|
self.timestamps.push(timestamp);
|
|
self.session_ids.push(session_id);
|
|
self.tags.push(tags);
|
|
self.tombstones.push(0);
|
|
self.norms.push(norm);
|
|
self.activation_weights.push(1.0);
|
|
idx
|
|
}
|
|
|
|
/// Find an active (non-tombstoned) entry by tags (used as key for dedup).
|
|
/// Returns the index of the first matching active entry, or None.
|
|
pub fn find_by_tags(&self, tags: &str) -> Option<usize> {
|
|
if tags.is_empty() {
|
|
return None;
|
|
}
|
|
for (i, t) in self.tags.iter().enumerate() {
|
|
if self.tombstones.get(i).copied().unwrap_or(1) == 0 && t == tags {
|
|
return Some(i);
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
/// Update an existing entry in-place (for upsert dedup).
|
|
pub fn update(
|
|
&mut self,
|
|
idx: usize,
|
|
chunk: String,
|
|
embedding: Vec<f32>,
|
|
source_channel: String,
|
|
timestamp: f64,
|
|
session_id: String,
|
|
) {
|
|
if idx < self.chunks.len() {
|
|
let embedding = self.stored_form(embedding);
|
|
let norm = vector_search::compute_norm(&embedding);
|
|
self.chunks[idx] = chunk;
|
|
self.embeddings.set(idx, &embedding);
|
|
self.source_channels[idx] = source_channel;
|
|
self.timestamps[idx] = timestamp;
|
|
self.session_ids[idx] = session_id;
|
|
self.norms[idx] = norm;
|
|
self.activation_weights[idx] = 1.0; // reset activation on update
|
|
}
|
|
}
|
|
|
|
/// Mark an entry as deleted (tombstoned).
|
|
pub fn mark_deleted(&mut self, id: usize) -> bool {
|
|
if id < self.tombstones.len() && self.tombstones[id] == 0 {
|
|
self.tombstones[id] = 1;
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Fraction of entries that are tombstoned.
|
|
pub fn tombstone_fraction(&self) -> f32 {
|
|
if self.chunks.is_empty() {
|
|
return 0.0;
|
|
}
|
|
let tombstoned = self.tombstones.iter().filter(|&&t| t == 1).count();
|
|
tombstoned as f32 / self.chunks.len() as f32
|
|
}
|
|
|
|
/// Remove all tombstoned entries, returns number removed.
|
|
/// Also returns a mapping from old indices to new indices (None if removed).
|
|
/// Recomputes norms for remaining entries.
|
|
pub fn compact(&mut self) -> (usize, Vec<Option<usize>>) {
|
|
let old_len = self.chunks.len();
|
|
let mut index_map = vec![None; old_len];
|
|
let mut new_idx = 0usize;
|
|
|
|
let mut new_chunks = Vec::new();
|
|
let mut new_embeddings = Vec::new();
|
|
let mut new_source_channels = Vec::new();
|
|
let mut new_timestamps = Vec::new();
|
|
let mut new_session_ids = Vec::new();
|
|
let mut new_tags = Vec::new();
|
|
let mut new_tombstones = Vec::new();
|
|
let mut new_norms = Vec::new();
|
|
let mut new_activation_weights = Vec::new();
|
|
|
|
for (i, slot) in index_map.iter_mut().enumerate() {
|
|
if self.tombstones[i] == 0 {
|
|
*slot = Some(new_idx);
|
|
new_idx += 1;
|
|
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
|
new_chunks.push(self.chunks[i].clone());
|
|
new_embeddings.push(self.embeddings[i].to_vec());
|
|
new_source_channels.push(self.source_channels[i].clone());
|
|
new_timestamps.push(self.timestamps[i]);
|
|
new_session_ids.push(self.session_ids[i].clone());
|
|
new_tags.push(self.tags[i].clone());
|
|
new_tombstones.push(0u8);
|
|
new_norms.push(norm);
|
|
new_activation_weights.push(self.activation_weights[i]);
|
|
}
|
|
}
|
|
|
|
let removed = old_len - new_chunks.len();
|
|
self.chunks = new_chunks;
|
|
self.embeddings
|
|
.reset_from(self.embedding_dim, new_embeddings);
|
|
self.source_channels = new_source_channels;
|
|
self.timestamps = new_timestamps;
|
|
self.session_ids = new_session_ids;
|
|
self.tags = new_tags;
|
|
self.tombstones = new_tombstones;
|
|
self.norms = new_norms;
|
|
self.activation_weights = new_activation_weights;
|
|
|
|
(removed, index_map)
|
|
}
|
|
|
|
/// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
|
|
/// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
|
|
pub fn flat_embeddings_owned(&self) -> Vec<f32> {
|
|
self.embeddings.as_flat().to_vec()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
|
fn assert_flat_in_sync(cache: &MemoryCache) {
|
|
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
|
assert_eq!(cache.embeddings.as_flat(), expected);
|
|
}
|
|
|
|
#[test]
|
|
fn push_keeps_flat_buffer_in_sync() {
|
|
let mut cache = MemoryCache::new(3);
|
|
cache.push(
|
|
"a".into(),
|
|
vec![1.0, 2.0, 3.0],
|
|
"chan".into(),
|
|
0.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.push(
|
|
"b".into(),
|
|
vec![4.0, 5.0, 6.0],
|
|
"chan".into(),
|
|
1.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
assert_flat_in_sync(&cache);
|
|
assert_eq!(
|
|
cache.embeddings.as_flat(),
|
|
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn update_keeps_flat_buffer_in_sync() {
|
|
let mut cache = MemoryCache::new(3);
|
|
cache.push(
|
|
"a".into(),
|
|
vec![1.0, 2.0, 3.0],
|
|
"chan".into(),
|
|
0.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.push(
|
|
"b".into(),
|
|
vec![4.0, 5.0, 6.0],
|
|
"chan".into(),
|
|
1.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.update(
|
|
0,
|
|
"a2".into(),
|
|
vec![7.0, 8.0, 9.0],
|
|
"chan".into(),
|
|
2.0,
|
|
"s1".into(),
|
|
);
|
|
assert_flat_in_sync(&cache);
|
|
assert_eq!(
|
|
cache.embeddings.as_flat(),
|
|
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
|
"update must overwrite the correct flat slice, not just append"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn compact_keeps_flat_buffer_in_sync() {
|
|
let mut cache = MemoryCache::new(2);
|
|
cache.push(
|
|
"a".into(),
|
|
vec![1.0, 1.0],
|
|
"chan".into(),
|
|
0.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.push(
|
|
"b".into(),
|
|
vec![2.0, 2.0],
|
|
"chan".into(),
|
|
1.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.push(
|
|
"c".into(),
|
|
vec![3.0, 3.0],
|
|
"chan".into(),
|
|
2.0,
|
|
"s1".into(),
|
|
String::new(),
|
|
);
|
|
cache.mark_deleted(1);
|
|
cache.compact();
|
|
assert_flat_in_sync(&cache);
|
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 1.0, 3.0, 3.0]);
|
|
}
|
|
|
|
#[test]
|
|
fn rebuild_flat_matches_manual_flatten() {
|
|
let mut cache = MemoryCache::new(2);
|
|
cache
|
|
.embeddings
|
|
.reset_from(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
|
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 2.0, 3.0, 4.0]);
|
|
}
|
|
|
|
#[test]
|
|
fn set_half_precision_rounds_existing_rows_and_their_norms() {
|
|
// A store with float16 set whose checkpoint is still f32 on disk
|
|
// loads full-precision rows; switching rounding on must bring them to
|
|
// exactly what the next checkpoint will write.
|
|
let mut cache = MemoryCache::new(3);
|
|
cache.push(
|
|
"a".into(),
|
|
vec![0.1, 0.2, 0.3],
|
|
"c".into(),
|
|
0.0,
|
|
"s".into(),
|
|
"".into(),
|
|
);
|
|
cache.push(
|
|
"b".into(),
|
|
vec![0.5, 0.25, 1.0],
|
|
"c".into(),
|
|
0.0,
|
|
"s".into(),
|
|
"".into(),
|
|
);
|
|
let exact_norm = cache.norms[0];
|
|
|
|
cache.set_half_precision(true);
|
|
let row0: Vec<f32> = [0.1f32, 0.2, 0.3]
|
|
.iter()
|
|
.map(|&v| round_to_f16(v))
|
|
.collect();
|
|
assert_eq!(&cache.embeddings[0], row0.as_slice());
|
|
assert_eq!(cache.norms[0], vector_search::compute_norm(&row0));
|
|
assert_ne!(cache.norms[0], exact_norm);
|
|
// Already representable: untouched.
|
|
assert_eq!(&cache.embeddings[1], &[0.5, 0.25, 1.0]);
|
|
|
|
// New rows are rounded as they arrive, and updates too.
|
|
cache.push(
|
|
"c".into(),
|
|
vec![0.1, 0.0, 0.0],
|
|
"c".into(),
|
|
0.0,
|
|
"s".into(),
|
|
"".into(),
|
|
);
|
|
assert_eq!(cache.embeddings[2][0], round_to_f16(0.1));
|
|
cache.update(
|
|
2,
|
|
"c".into(),
|
|
vec![0.3, 0.0, 0.0],
|
|
"c".into(),
|
|
0.0,
|
|
"s".into(),
|
|
);
|
|
assert_eq!(cache.embeddings[2][0], round_to_f16(0.3));
|
|
}
|
|
}
|