Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c470244a6f | ||
|
|
d0db83812b | ||
|
|
5e4aa1c6bf |
@@ -149,6 +149,83 @@ vectors, and recall is measured against brute-force ground truth rather than
|
||||
against the f32 index, whose own approximation errors a re-scored search is
|
||||
entitled to get right.
|
||||
|
||||
### Search options: source filters, re-ranking, confidence
|
||||
|
||||
Measured 2026-09-24 on tank (AMD Ryzen 7 7800X3D). `HDF5Memory::search`
|
||||
with `SearchOptions`, clustered 384-dim data, k = 10, Hebbian boosting off.
|
||||
Filters keep 50%, 10% or 1% of the store at random, or two whole clusters
|
||||
chosen *away* from each query's own — the case an ANN index handles worst,
|
||||
because nothing it finds near the query is allowed. Recall is vector-only
|
||||
against an exact scan of the allowed records; latency is full hybrid search
|
||||
(vector + BM25). 200 queries, medians of three runs (recall was identical in
|
||||
every run).
|
||||
|
||||
```bash
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
|
||||
```
|
||||
|
||||
| N | options | filtered recall@10 | p50 ms | p99 ms |
|
||||
|---:|---|---:|---:|---:|
|
||||
| 10 000 | no filter | 1.0000 | 0.488 | 0.518 |
|
||||
| 10 000 | random 50% | 1.0000 | 0.520 | 0.563 |
|
||||
| 10 000 | random 10% | 1.0000 | 0.342 | 0.367 |
|
||||
| 10 000 | random 1% | 1.0000 | 0.220 | 0.239 |
|
||||
| 10 000 | 2 clusters away from the query | 1.0000 | 0.252 | 0.278 |
|
||||
| 10 000 | re-rank | — | 0.562 | 0.674 |
|
||||
| 10 000 | re-rank + confidence | — | 0.554 | 0.575 |
|
||||
| 100 000 | no filter | 0.9995 | 4.600 | 5.245 |
|
||||
| 100 000 | random 50% | 1.0000 | 4.770 | 5.614 |
|
||||
| 100 000 | random 10% | 1.0000 | 3.244 | 4.008 |
|
||||
| 100 000 | random 1% | 1.0000 | 2.288 | 2.943 |
|
||||
| 100 000 | 2 clusters away from the query | 1.0000 | 2.298 | 3.048 |
|
||||
| 100 000 | re-rank | — | 4.747 | 5.411 |
|
||||
| 100 000 | re-rank + confidence | — | 4.748 | 5.526 |
|
||||
|
||||
A filtered search finds the exact filtered top 10 and is never slower than an
|
||||
unfiltered one. The filter applies before ranking — over-fetching the index in
|
||||
proportion to what the filter removes, and scanning the allowed records
|
||||
exactly whenever that costs fewer distance evaluations than the index would
|
||||
(roughly `pool × M`). The first version compared the over-fetch to the store
|
||||
size instead, and that measured badly at 100K: the 1% filter took 5.9 ms at
|
||||
recall 0.9965 and the away-from-query filter 12.3 ms at 0.976, both through an
|
||||
index asked for ~16 000 candidates, where scanning the few hundred or thousand
|
||||
allowed records is exact and cheap. Re-ranking a 3k candidate pool and
|
||||
confidence rejection add about 3%.
|
||||
|
||||
### float16 embedding storage (`MemoryConfig::float16`)
|
||||
|
||||
Measured 2026-09-23 on tank (AMD Ryzen 7 7800X3D). The same clustered
|
||||
384-dim data in an `f32` store and a `float16` store, both with the default
|
||||
int8 index and Hebbian boosting off (so every query sees the same store).
|
||||
Recall is vector-only `hybrid_search` against an exact scan of the original
|
||||
`f32` vectors, 200 queries. Six runs, three with each store going first;
|
||||
medians. Nothing depended on the order.
|
||||
|
||||
```bash
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full --f16-first
|
||||
```
|
||||
|
||||
| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap | hybrid p50 ms |
|
||||
|---:|---|---:|---:|---:|---:|---:|---:|
|
||||
| 1 000 | f32 | 1.6 | 8 | 1.6 | 1.0000 | | 0.072 |
|
||||
| 1 000 | float16 | 0.8 | 6 | 1.9 | 0.9980 | 0.9980 | 0.072 |
|
||||
| 10 000 | f32 | 15.4 | 66 | 15.2 | 1.0000 | | 0.495 |
|
||||
| 10 000 | float16 | 8.2 | 45 | 18.1 | 1.0000 | 1.0000 | 0.494 |
|
||||
| 100 000 | f32 | 154.0 | 752 | 299.8 | 0.9940 | | 4.676 |
|
||||
| 100 000 | float16 | **80.8** | **512** | **252.1** | 0.9990 | 0.9940 | 4.654 |
|
||||
|
||||
The file is 48% smaller, checkpoints write less and open reads less. At
|
||||
small N opening is slightly slower (widening halves costs more than the I/O it
|
||||
saves: +3 ms at 10K). Recall does not move: half precision keeps about three
|
||||
significant digits, far finer than the gaps between neighbours on unit-length
|
||||
embeddings. Recall was identical in every run; the 0.999 against 0.994 at
|
||||
100K is two slightly different HNSW graphs, not an improvement to claim.
|
||||
|
||||
The in-memory cache holds the half-rounded values, so the store searches the
|
||||
same before and after a reopen; RAM use is unchanged (the cache is still
|
||||
`f32`). What `float16` saves is disk, and the I/O that goes with it.
|
||||
|
||||
### Opening a store (`read_from_disk`)
|
||||
|
||||
`HDF5Memory::open` memory-mapped the file, copied the whole mapping into a
|
||||
|
||||
@@ -3,6 +3,23 @@
|
||||
## Unreleased
|
||||
|
||||
### Upgrade Notes
|
||||
- **Files written by clawhdf5 now open in h5py and libhdf5.** Every `f32`
|
||||
dataset we wrote — including every agent store's embeddings — was refused
|
||||
with "sign bit position out of bounds", and every empty dataset with
|
||||
"invalid dataset size". Both were write-side bugs present in every release;
|
||||
clawhdf5's own reader was unaffected. An agent store is rewritten in full at
|
||||
each checkpoint, so it becomes readable at its next checkpoint on this
|
||||
version; other files with `f32` or empty datasets need rewriting. Details in
|
||||
`docs/known-issues.md`.
|
||||
- **`MemoryConfig::float16` now does what it says.** It was persisted and
|
||||
otherwise ignored; embeddings were always stored as `f32`. A store created
|
||||
with it on now writes half-precision embeddings (48% smaller files) and
|
||||
rounds embeddings to half precision as they are saved. A store that already
|
||||
had `float16 = true` rounds its embeddings when next opened and writes them
|
||||
as `float16` at its next checkpoint. Off by default.
|
||||
- **Breaking:** `MemoryError` gained `InvalidEntry`, returned when a
|
||||
`float16` store is given an embedding value beyond ±65504. Exhaustive
|
||||
matches need the new arm.
|
||||
- **The default build no longer compiles any C.** Deflate now defaults to the
|
||||
pure-Rust zlib-rs instead of zlib-ng, so building the core crates needs
|
||||
neither cmake nor a C compiler. Speed on HDF5 reads and writes is within 6%
|
||||
@@ -23,6 +40,62 @@
|
||||
`quantized_index = false`, or pass `create --f32-index` to the CLI, to opt
|
||||
out. The CLI's `--quantized-index` is still accepted but is now a no-op.
|
||||
|
||||
### Search
|
||||
- `clawhdf5-agent`: **`HDF5Memory::search` with `SearchOptions`** — source
|
||||
filtering, re-ranking and confidence rejection in the store's own search
|
||||
path. Re-ranking and confidence rejection used to be reachable only
|
||||
through the OpenClaw backend, which now calls `search` with both on.
|
||||
- `with_sources([..])` restricts a search to records from those source
|
||||
channels. It applies before ranking, so a filtered search still returns up
|
||||
to `k` results, normalised over what it can return. Measured at 100K: the
|
||||
exact filtered top 10 for filters keeping 50%, 10% and 1% of the store and
|
||||
for records far from the query, and never slower than an unfiltered search
|
||||
(2.3 ms for a 1% filter vs 4.6 ms unfiltered). See `BENCHMARKS.md`,
|
||||
"Search options".
|
||||
- `with_rerank(ReRankConfig)` re-ranks a pool of `max(3k, 10)` candidates
|
||||
(`rerank_pool` to change it) by relevance, recency, source authority and
|
||||
activation; `with_confidence(ConfidenceConfig)` drops low-confidence
|
||||
results; `at_time(now)` pins the clock for recency. About 3% on latency.
|
||||
- `hybrid_search` and `hybrid_search_with` are unchanged (tested bit for
|
||||
bit against `search` with default options).
|
||||
- `clawhdf5-agent`: the OpenClaw backend's search now boosts the Hebbian
|
||||
activation of the `k` results it returns, not of the whole `3k` candidate
|
||||
pool it re-ranks.
|
||||
|
||||
### Interop
|
||||
- `clawhdf5-format`: **every `f32` dataset was unreadable by h5py and
|
||||
libhdf5.** The float datatype encoder hard-coded the sign bit's position to
|
||||
63, correct only for `f64`; libhdf5 validates it and refused the dataset. It
|
||||
is now derived from the type (15 / 31 / 63). Our reader ignores the field,
|
||||
and the interop suites only wrote `f64`, which is how it went unnoticed.
|
||||
- `clawhdf5-format`: **every empty dataset was unreadable by h5py and
|
||||
libhdf5.** It was written with a real address and zero bytes, which trips
|
||||
libhdf5's `addr + size <= addr` overflow check. An empty contiguous dataset
|
||||
now gets the undefined address, as libhdf5 writes it. This affected every
|
||||
agent store without sessions or a knowledge graph.
|
||||
- New interop tests: `f32` and `float16` datasets in both directions (our
|
||||
`float16` rounding matches numpy's bit for bit on 4 020 probe values,
|
||||
including ties, subnormals and the overflow boundary), and an agent store —
|
||||
`f32` and `float16` — opened by h5py with every dataset decoded.
|
||||
|
||||
### Storage
|
||||
- `clawhdf5-format`: **half-precision datasets.**
|
||||
`DatasetBuilder::with_f16_data` writes IEEE binary16 (numpy `float16`),
|
||||
rounding to nearest-even; `make_f16_type`, and `clawhdf5_format::float16`
|
||||
with the conversions, which are checked against the `half` crate on 16.7M
|
||||
values and round-trip all 65 536 half values. Reading `float16` as `f32`
|
||||
gained a little-endian fast path.
|
||||
- `clawhdf5-agent`: **`MemoryConfig::float16` stores embeddings as half
|
||||
precision.** At 100K x 384 the file goes from 154.0 to 80.8 MiB (−48%), a
|
||||
checkpoint from 752 to 512 ms and open from 300 to 252 ms, with the same
|
||||
vector recall@10 against an exact scan (0.999 vs 0.994) and the same
|
||||
`hybrid_search` latency; at 10K open is 3 ms slower. The cache rounds each
|
||||
embedding as it is saved, so memory and file agree bit for bit and a store
|
||||
returns the same results before and after a reopen (tested). Out-of-range
|
||||
values are refused with `MemoryError::InvalidEntry` rather than stored as
|
||||
infinity; batches are all or nothing. CLI: `create --float16`. See
|
||||
`BENCHMARKS.md`, "float16 embedding storage".
|
||||
|
||||
### Build
|
||||
- **Pure-Rust default.** `clawhdf5-format`, `clawhdf5-filters` and the
|
||||
`clawhdf5` facade default to the `zlib-rs` deflate backend; `fast-deflate`
|
||||
|
||||
@@ -87,6 +87,24 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
`export` do). An unreadable WAL (torn header, bad magic) is quarantined to
|
||||
`<store>.h5.wal.corrupt-<ts>` rather than blocking `open()`; a WAL with an
|
||||
unknown *newer* version still fails and is left untouched.
|
||||
- `MemoryConfig::float16` (off by default, persisted; CLI `create --float16`)
|
||||
writes `/memory/embeddings` as IEEE half precision (48% smaller file at
|
||||
100K, same recall). `MemoryCache::half_precision` rounds each embedding as
|
||||
it enters the cache (push, update, WAL replay, and on load of a store still
|
||||
`f32` on disk), so memory and file agree bit for bit; the conversions live
|
||||
in `clawhdf5_format::float16` and must stay the single implementation.
|
||||
Values beyond ±65504 are `MemoryError::InvalidEntry`. Interop: every file
|
||||
must open in h5py — `f32` datasets and empty datasets did not until
|
||||
2026-09-23 (see `docs/known-issues.md`); the agent's `h5py_interop` test
|
||||
guards a whole store.
|
||||
- `HDF5Memory::search(query_emb, text, &SearchOptions)` is the full search
|
||||
path: optional source-channel filter (applied before ranking; exact scan of
|
||||
the allowed records whenever cheaper than `pool × M` index distance
|
||||
evaluations, and as the fallback when the pool comes back short), fusion,
|
||||
activation scaling, optional re-ranking and confidence rejection.
|
||||
`hybrid_search`/`hybrid_search_with` are thin wrappers; the OpenClaw
|
||||
backend is `search` with re-rank + confidence on. Measure changes with
|
||||
`search_harness --options-study`.
|
||||
- `MemoryConfig::compression` is off by default; when on, embeddings are
|
||||
deflate-compressed, or Zstd with the agent's `zstd` feature (links libzstd).
|
||||
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
||||
|
||||
@@ -82,6 +82,18 @@ breaking change, are in [CHANGELOG.md](CHANGELOG.md).
|
||||
100K × 384 store to 1.74× the raw vectors. At equal recall it is also faster
|
||||
than `f32`: 1.63× QPS on AVX2, 1.18× on a Raspberry Pi 5 (NEON `SDOT`).
|
||||
|
||||
**Interop and search (unreleased)**
|
||||
- **Files we write now open in h5py and libhdf5.** Every `f32` dataset —
|
||||
including every agent store's embeddings — and every empty dataset was
|
||||
refused by libhdf5. Both were write-side bugs in every release; agent stores
|
||||
fix themselves at their next checkpoint. See
|
||||
[docs/known-issues.md](docs/known-issues.md).
|
||||
- `MemoryConfig::float16` now stores half-precision embeddings (it was
|
||||
ignored): 48% smaller files at the same recall.
|
||||
- `HDF5Memory::search` with `SearchOptions`: filter by source channel (exact
|
||||
filtered top-k, never slower than unfiltered), and opt-in re-ranking and
|
||||
confidence rejection, which used to be OpenClaw-only.
|
||||
|
||||
**Tooling**
|
||||
- CI now runs the h5py/netCDF4 interop suites for real (they had been skipping
|
||||
silently) and runs an aarch64 job for the NEON kernels.
|
||||
@@ -244,6 +256,9 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
||||
| 10K | 17.0 MB | 1.7 KB | 2.7 MB (6.2x) |
|
||||
| 100K | 169.8 MB | 1.7 KB | 26.9 MB (6.2x) |
|
||||
|
||||
With `MemoryConfig::float16` the embeddings take half the space: an agent
|
||||
store of 100K × 384 records is 80.8 MiB instead of 154.0.
|
||||
|
||||
**In memory** — a store reopened from disk, 384-dim `f32`, measured with a
|
||||
counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)):
|
||||
|
||||
@@ -281,12 +296,14 @@ ClawhDF5's agent memory engine draws on 15+ recent papers on agentic memory syst
|
||||
└────────┬────────┘
|
||||
│
|
||||
┌─────────────────▼──────────────────┐
|
||||
│ HDF5Memory::hybrid_search │
|
||||
│ HDF5Memory::search │
|
||||
│ optional source-channel filter │
|
||||
│ HNSW vector + BM25 keyword │
|
||||
│ weighted fusion (0.4 / 0.6) │
|
||||
│ × √(Hebbian activation) │
|
||||
└─────────────────┬──────────────────┘
|
||||
│ OpenClaw backend adds:
|
||||
│ opt-in (SearchOptions);
|
||||
│ the OpenClaw backend turns both on
|
||||
┌─────────────────▼──────────────────┐
|
||||
│ Multi-factor re-ranking │
|
||||
│ relevance · recency · authority · │
|
||||
@@ -322,8 +339,8 @@ directly; the store persists the records, sessions and graph they work over.
|
||||
| **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy (Levenshtein) entity resolution |
|
||||
| **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring, novelty, and time-decay |
|
||||
| **`hybrid`** | Vector + BM25 fusion. Default is a min-max-normalised weighted sum, vector 0.4 / keyword 0.6 (`hybrid::DEFAULT_FUSION`, tuned on LongMemEval); RRF is available via `Fusion::Rrf` / `hybrid_search_with`. The vector stage uses the HNSW index by default (`hnsw` feature); disable with `--no-default-features --features float16` for an exact linear scan |
|
||||
| **`reranker`** | Multi-factor re-ranking: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Used by the OpenClaw backend |
|
||||
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches (OpenClaw backend) |
|
||||
| **`reranker`** | Multi-factor re-ranking: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Opt-in via `SearchOptions::with_rerank`; on in the OpenClaw backend |
|
||||
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches. Opt-in via `SearchOptions::with_confidence`; on in the OpenClaw backend |
|
||||
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
||||
| **`multimodal`** | Cross-modal search across text/image/audio/video embeddings |
|
||||
| **`provenance`** | Source attribution and an unkeyed FNV-1a content hash per record, held in memory for the session, for detecting accidental corruption (not tamper-proof) |
|
||||
@@ -389,6 +406,31 @@ for result in results {
|
||||
}
|
||||
```
|
||||
|
||||
### Search Options
|
||||
|
||||
```rust
|
||||
use clawhdf5_agent::SearchOptions;
|
||||
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||
use clawhdf5_agent::reranker::ReRankConfig;
|
||||
|
||||
// Only memories from these source channels; still a full page of k results.
|
||||
let work = memory.search(
|
||||
&query_embedding,
|
||||
"deadline",
|
||||
&SearchOptions::new(5).with_sources(["slack", "email"]),
|
||||
);
|
||||
|
||||
// Re-rank by relevance, recency, source authority and activation, then drop
|
||||
// low-confidence results — the pipeline the OpenClaw backend runs.
|
||||
let careful = memory.search(
|
||||
&query_embedding,
|
||||
"user preferences",
|
||||
&SearchOptions::new(5)
|
||||
.with_rerank(ReRankConfig::default())
|
||||
.with_confidence(ConfidenceConfig::default()),
|
||||
);
|
||||
```
|
||||
|
||||
### Knowledge Graph
|
||||
|
||||
```rust
|
||||
@@ -538,7 +580,7 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `float16` | **yes** | Half-precision cosine kernel (`cosine_similarity_f16`). The store itself always writes `f32` embeddings; `MemoryConfig::float16` is recorded in `/meta` but not yet applied |
|
||||
| `float16` | **yes** | Half-precision cosine kernel (`cosine_similarity_f16`). Half-precision *storage* is the `MemoryConfig::float16` setting below, and needs no feature |
|
||||
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
|
||||
| `parallel` | **yes** | Parallel HNSW bulk build (same graph, ~3× faster on 16 cores) and Rayon brute-force search strategies |
|
||||
| `zstd` | no | Compress embeddings with Zstd instead of deflate when `MemoryConfig::compression` is on (links libzstd) |
|
||||
@@ -568,6 +610,14 @@ setting existed keep their `f32` index; opt out for new stores with
|
||||
`quantized_index = false` or `clawhdf5-cli create --f32-index`. See
|
||||
[BENCHMARKS.md § Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index).
|
||||
|
||||
`MemoryConfig::float16` (off by default; CLI `create --float16`) stores the
|
||||
embeddings on disk as IEEE half precision (numpy `float16`): at 100K × 384 the
|
||||
file drops from 154 to 81 MiB, checkpoints and opens get faster, and vector
|
||||
recall and search latency do not change. Embeddings are rounded as they are
|
||||
saved, so the store searches the same before and after a reopen; values must
|
||||
lie within ±65504. See
|
||||
[BENCHMARKS.md § float16 embedding storage](BENCHMARKS.md#float16-embedding-storage-memoryconfigfloat16).
|
||||
|
||||
### `clawhdf5-format`
|
||||
|
||||
| Flag | Default | Description |
|
||||
@@ -656,8 +706,9 @@ agent_memory.h5
|
||||
│ └── ann_generation (ties the .ann sidecar to this checkpoint)
|
||||
├── /memory
|
||||
│ ├── chunks: string[N]
|
||||
│ ├── embeddings: f32[N × D] (chunked; deflate, or Zstd with the
|
||||
│ │ `zstd` feature, when compression is on)
|
||||
│ ├── embeddings: f32[N × D], or f16 for a `float16` store
|
||||
│ │ (chunked; deflate, or Zstd with the `zstd`
|
||||
│ │ feature, when compression is on)
|
||||
│ ├── source_channel: string[N]
|
||||
│ ├── timestamps: f64[N]
|
||||
│ ├── session_ids: string[N]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 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.
|
||||
///
|
||||
@@ -149,6 +150,11 @@ pub struct MemoryCache {
|
||||
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 {
|
||||
@@ -164,9 +170,44 @@ impl MemoryCache {
|
||||
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")]
|
||||
@@ -202,6 +243,7 @@ impl MemoryCache {
|
||||
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);
|
||||
@@ -240,6 +282,7 @@ impl MemoryCache {
|
||||
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);
|
||||
@@ -439,4 +482,60 @@ mod tests {
|
||||
.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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -65,31 +65,34 @@ pub fn hybrid_search_fused(
|
||||
) -> Vec<(usize, f32)> {
|
||||
// Get raw scores from both systems. Request all results so normalization
|
||||
// covers the full distribution.
|
||||
// Use parallel search when rayon feature is enabled and vector count > 10K.
|
||||
let vec_scores = {
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
if vectors.count() > 10_000 {
|
||||
vector_search::parallel_cosine_batch(
|
||||
query_embedding,
|
||||
vectors,
|
||||
tombstones,
|
||||
vectors.count(),
|
||||
)
|
||||
} else {
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
}
|
||||
}
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
{
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
}
|
||||
};
|
||||
let vec_scores = exact_vector_scores(query_embedding, vectors, tombstones);
|
||||
let kw_scores = bm25_index.scores(query_text);
|
||||
|
||||
fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
|
||||
/// Cosine similarity of `query_embedding` to every vector whose `skip` byte is
|
||||
/// 0 (a tombstone, or any other exclusion mask). Parallel above 10K vectors
|
||||
/// when the `parallel` feature is on.
|
||||
pub fn exact_vector_scores(
|
||||
query_embedding: &[f32],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
skip: &[u8],
|
||||
) -> Vec<(usize, f32)> {
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
if vectors.count() > 10_000 {
|
||||
return vector_search::parallel_cosine_batch(
|
||||
query_embedding,
|
||||
vectors,
|
||||
skip,
|
||||
vectors.count(),
|
||||
);
|
||||
}
|
||||
}
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, skip)
|
||||
}
|
||||
|
||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||
///
|
||||
/// Both score sets are independently min-max normalized to [0, 1] and combined
|
||||
|
||||
@@ -63,6 +63,7 @@ use std::path::{Path, PathBuf};
|
||||
use cache::MemoryCache;
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
use clawhdf5_format::float16::round_to_f16;
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
||||
@@ -71,6 +72,7 @@ use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
||||
use knowledge::KnowledgeCache;
|
||||
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
||||
pub use search::SearchOptions;
|
||||
use session::SessionCache;
|
||||
|
||||
// --- Error type ---
|
||||
@@ -83,6 +85,9 @@ pub enum MemoryError {
|
||||
NotFound(String),
|
||||
/// Another `HDF5Memory` (in this or another process) has the store open.
|
||||
Locked(String),
|
||||
/// A record the store cannot hold as given, e.g. an embedding value
|
||||
/// outside the half-precision range of a `float16` store.
|
||||
InvalidEntry(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for MemoryError {
|
||||
@@ -93,6 +98,7 @@ impl std::fmt::Display for MemoryError {
|
||||
MemoryError::Schema(e) => write!(f, "schema error: {e}"),
|
||||
MemoryError::NotFound(e) => write!(f, "not found: {e}"),
|
||||
MemoryError::Locked(e) => write!(f, "store is locked: {e}"),
|
||||
MemoryError::InvalidEntry(e) => write!(f, "invalid entry: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -124,6 +130,12 @@ pub struct MemoryConfig {
|
||||
pub embedding_dim: usize,
|
||||
pub chunk_size: usize,
|
||||
pub overlap: usize,
|
||||
/// Store embeddings as IEEE half precision (numpy `float16`): half the
|
||||
/// bytes of the embeddings dataset on disk. Every embedding is rounded to
|
||||
/// the nearest half as it enters the store, in memory as well as on disk,
|
||||
/// so search results are the same before and after a reopen. Values must
|
||||
/// lie within ±65504; a save outside that is `MemoryError::InvalidEntry`.
|
||||
/// Fixed when the store is created (persisted in `/meta`).
|
||||
pub float16: bool,
|
||||
pub compression: bool,
|
||||
pub compression_level: u32,
|
||||
@@ -317,7 +329,8 @@ impl HDF5Memory {
|
||||
/// Create a new HDF5 memory file with the given configuration.
|
||||
pub fn create(config: MemoryConfig) -> Result<Self> {
|
||||
let lock = store_lock::StoreLock::acquire(&config.path)?;
|
||||
let cache = MemoryCache::new(config.embedding_dim);
|
||||
let mut cache = MemoryCache::new(config.embedding_dim);
|
||||
cache.set_half_precision(config.float16);
|
||||
let sessions = SessionCache::new();
|
||||
let knowledge = KnowledgeCache::new();
|
||||
|
||||
@@ -1070,7 +1083,29 @@ impl HDF5Memory {
|
||||
/// Upsert: if an active entry with the same tags (key) exists, update it in-place.
|
||||
/// Otherwise append a new entry. Use this for key-based memory stores where
|
||||
/// the same key should not create duplicates.
|
||||
/// A `float16` store holds embeddings as IEEE half precision, which has no
|
||||
/// finite value beyond ±65504. Refuse such an embedding rather than
|
||||
/// silently store infinity. (Values that are already infinite or NaN are
|
||||
/// stored as they are, as in an `f32` store.)
|
||||
fn check_embedding(&self, embedding: &[f32]) -> Result<()> {
|
||||
if !self.config.float16 {
|
||||
return Ok(());
|
||||
}
|
||||
let overflow = embedding
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find(|&(_, &v)| v.is_finite() && round_to_f16(v).is_infinite());
|
||||
match overflow {
|
||||
None => Ok(()),
|
||||
Some((i, v)) => Err(MemoryError::InvalidEntry(format!(
|
||||
"embedding[{i}] = {v} is outside the half-precision range (±65504) \
|
||||
of this float16 store"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn save_or_update(&mut self, entry: MemoryEntry) -> Result<usize> {
|
||||
self.check_embedding(&entry.embedding)?;
|
||||
if let Some(existing_idx) = self.cache.find_by_tags(&entry.tags) {
|
||||
if let Some(ref mut w) = self.wal {
|
||||
let wal_entry = wal::WalEntry {
|
||||
@@ -1126,6 +1161,7 @@ impl HDF5Memory {
|
||||
|
||||
impl AgentMemory for HDF5Memory {
|
||||
fn save(&mut self, entry: MemoryEntry) -> Result<usize> {
|
||||
self.check_embedding(&entry.embedding)?;
|
||||
if let Some(ref mut w) = self.wal {
|
||||
let wal_entry = wal::WalEntry {
|
||||
entry_type: wal::WalEntryType::Save,
|
||||
@@ -1167,6 +1203,10 @@ impl AgentMemory for HDF5Memory {
|
||||
}
|
||||
|
||||
fn save_batch(&mut self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>> {
|
||||
// All or nothing: check every entry before storing any.
|
||||
for entry in &entries {
|
||||
self.check_embedding(&entry.embedding)?;
|
||||
}
|
||||
let mut indices = Vec::with_capacity(entries.len());
|
||||
for entry in entries {
|
||||
let idx = self.cache.push(
|
||||
@@ -1327,6 +1367,9 @@ impl HDF5Memory {
|
||||
})?;
|
||||
let view = memory_strategy::CacheStoreView::new(&self.cache, &self.knowledge);
|
||||
let output = strat.evaluate(&exchange, &view);
|
||||
for e in &output.entries {
|
||||
self.check_embedding(&e.embedding)?;
|
||||
}
|
||||
for e in &output.entries {
|
||||
self.cache.push(
|
||||
e.chunk.clone(),
|
||||
@@ -1411,6 +1454,16 @@ impl HDF5Memory {
|
||||
let mut promoted = 0;
|
||||
|
||||
for key in candidates {
|
||||
// Check before taking, so a rejected entry stays in the ephemeral
|
||||
// tier rather than being lost.
|
||||
if let Some(emb) = self
|
||||
.ephemeral
|
||||
.as_ref()
|
||||
.and_then(|s| s.get_entry(&key))
|
||||
.and_then(|e| e.embedding.as_deref())
|
||||
{
|
||||
self.check_embedding(emb)?;
|
||||
}
|
||||
let entry = match self
|
||||
.ephemeral
|
||||
.as_mut()
|
||||
|
||||
@@ -13,9 +13,8 @@ use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::{
|
||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry,
|
||||
confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence},
|
||||
reranker::{ReRankConfig, RerankInput, rerank},
|
||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions,
|
||||
confidence::ConfidenceConfig, reranker::ReRankConfig,
|
||||
};
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
@@ -524,71 +523,27 @@ impl ClawhdfBackend {
|
||||
|
||||
impl MemoryBackend for ClawhdfBackend {
|
||||
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
||||
/// confidence-filter.
|
||||
/// confidence-filter — [`HDF5Memory::search`] with both stages on.
|
||||
fn search(
|
||||
&mut self,
|
||||
query_text: &str,
|
||||
query_embedding: &[f32],
|
||||
k: usize,
|
||||
) -> Vec<MemorySearchResult> {
|
||||
// 1. Hybrid retrieval (vector + BM25, fused by score).
|
||||
let candidates = k.saturating_mul(3).max(10);
|
||||
let raw = self.memory.hybrid_search_with(
|
||||
query_embedding,
|
||||
query_text,
|
||||
crate::hybrid::DEFAULT_FUSION,
|
||||
candidates,
|
||||
);
|
||||
|
||||
if raw.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let now = Self::now_secs();
|
||||
|
||||
// 2. Re-rank using temporal recency, source authority, Hebbian weight.
|
||||
let rerank_inputs: Vec<RerankInput> = raw
|
||||
.iter()
|
||||
.map(|r| RerankInput {
|
||||
index: r.index,
|
||||
timestamp: r.timestamp,
|
||||
source_channel: r.source_channel.clone(),
|
||||
raw_activation: r.activation,
|
||||
relevance: r.score,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let reranked = rerank(&rerank_inputs, &self.rerank_config, now);
|
||||
|
||||
// 3. Confidence rejection.
|
||||
let scored: Vec<ScoredResult> = reranked
|
||||
.iter()
|
||||
.map(|r| ScoredResult {
|
||||
index: r.index,
|
||||
score: r.combined_score,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let confident = reject_low_confidence(&scored, &self.confidence_config);
|
||||
|
||||
// 4. Map back to MemorySearchResult; preserve raw text via index lookup.
|
||||
let raw_by_idx: HashMap<usize, &crate::SearchResult> =
|
||||
raw.iter().map(|r| (r.index, r)).collect();
|
||||
|
||||
confident
|
||||
let options = SearchOptions::new(k)
|
||||
.with_rerank(self.rerank_config)
|
||||
.with_confidence(self.confidence_config.clone())
|
||||
.at_time(Self::now_secs());
|
||||
self.memory
|
||||
.search(query_embedding, query_text, &options)
|
||||
.into_iter()
|
||||
.take(k)
|
||||
.filter_map(|sr| {
|
||||
let r = raw_by_idx.get(&sr.index)?;
|
||||
let path = r.source_channel.clone();
|
||||
Some(MemorySearchResult {
|
||||
text: r.chunk.clone(),
|
||||
score: sr.score,
|
||||
path: path.clone(),
|
||||
line_range: None,
|
||||
timestamp: Some(r.timestamp),
|
||||
source: path,
|
||||
})
|
||||
.map(|r| MemorySearchResult {
|
||||
text: r.chunk,
|
||||
score: r.score,
|
||||
path: r.source_channel.clone(),
|
||||
line_range: None,
|
||||
timestamp: Some(r.timestamp),
|
||||
source: r.source_channel,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -159,20 +159,27 @@ fn build_memory_group(
|
||||
// chunks: fixed-length string array
|
||||
write_string_dataset(&mut group, "chunks", &cache.chunks);
|
||||
|
||||
// embeddings: f32 [N x D]
|
||||
// embeddings: [N x D], f32 — or IEEE half precision for a `float16`
|
||||
// store. The cache already holds half-rounded values then, so this
|
||||
// conversion is exact and a reopened store sees the same numbers.
|
||||
let n = cache.embeddings.len() as u64;
|
||||
let d = cache.embedding_dim as u64;
|
||||
let flat = cache.flat_embeddings();
|
||||
{
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(flat)
|
||||
.with_shape(&[n, d]);
|
||||
let ds = group.create_dataset("embeddings");
|
||||
let elem_bytes: u64 = if config.float16 {
|
||||
ds.with_f16_data(flat);
|
||||
2
|
||||
} else {
|
||||
ds.with_f32_data(flat);
|
||||
4
|
||||
};
|
||||
ds.with_shape(&[n, d]);
|
||||
|
||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||
if n > 0 && d > 0 {
|
||||
let target_chunk_bytes: u64 = 256 * 1024;
|
||||
let rows_per_chunk = (target_chunk_bytes / (d * 4)).max(1).min(n);
|
||||
let rows_per_chunk = (target_chunk_bytes / (d * elem_bytes)).max(1).min(n);
|
||||
ds.with_chunks(&[rows_per_chunk, d]);
|
||||
|
||||
// Compression. Shuffle is applied automatically (auto-shuffle
|
||||
@@ -513,7 +520,16 @@ pub fn validate_and_load(
|
||||
};
|
||||
|
||||
// Load /memory group
|
||||
let memory_cache = load_memory_group(file, embedding_dim)?;
|
||||
let mut memory_cache = load_memory_group(file, embedding_dim)?;
|
||||
// A float16 store's cache holds half-rounded embeddings. Embeddings read
|
||||
// from an f16 dataset already are; a float16 store whose last checkpoint
|
||||
// predates half-precision storage is still f32 on disk and is rounded
|
||||
// here.
|
||||
if config.float16 && embeddings_are_f16(file) {
|
||||
memory_cache.half_precision = true;
|
||||
} else {
|
||||
memory_cache.set_half_precision(config.float16);
|
||||
}
|
||||
|
||||
// Load /sessions group
|
||||
let session_cache = load_sessions_group(file)?;
|
||||
@@ -756,6 +772,13 @@ fn read_string_dataset_from_group(
|
||||
.map_err(|e| MemoryError::Hdf5(format!("cannot read strings from {name}: {e}")))
|
||||
}
|
||||
|
||||
/// Whether `/memory/embeddings` is stored as IEEE half precision.
|
||||
fn embeddings_are_f16(file: &clawhdf5::File) -> bool {
|
||||
file.dataset("memory/embeddings")
|
||||
.and_then(|ds| ds.dtype())
|
||||
.is_ok_and(|dt| matches!(dt, clawhdf5::DType::Other(ref s) if s == "float16"))
|
||||
}
|
||||
|
||||
fn read_f32_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<f32>, MemoryError> {
|
||||
let ds = group
|
||||
.dataset(name)
|
||||
|
||||
@@ -2,18 +2,107 @@
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::bm25;
|
||||
use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence};
|
||||
use crate::hybrid;
|
||||
use crate::reranker::{ReRankConfig, RerankInput, rerank};
|
||||
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};
|
||||
|
||||
/// Options for [`HDF5Memory::search`].
|
||||
///
|
||||
/// [`SearchOptions::new`] is plain hybrid search with the tuned default
|
||||
/// fusion — the same as `hybrid_search_with(.., hybrid::DEFAULT_FUSION, k)`.
|
||||
/// Every stage beyond that is opt-in.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchOptions {
|
||||
/// Number of results to return.
|
||||
pub k: usize,
|
||||
/// How the vector and keyword stages are combined.
|
||||
pub fusion: hybrid::Fusion,
|
||||
/// Only consider records whose `source_channel` is one of these. The
|
||||
/// filter applies *before* ranking, so a filtered search still returns up
|
||||
/// to `k` results and scores are normalised over the records it can
|
||||
/// return. `None` searches everything; an empty list matches nothing.
|
||||
pub source_channels: Option<Vec<String>>,
|
||||
/// Re-rank a candidate pool by retrieval relevance, recency, source
|
||||
/// authority and activation — the pipeline the OpenClaw backend runs.
|
||||
pub rerank: Option<ReRankConfig>,
|
||||
/// Candidates retrieved for re-ranking; 0 means `max(3k, 10)`.
|
||||
pub rerank_pool: usize,
|
||||
/// Drop low-confidence results (after re-ranking, when that is on).
|
||||
pub confidence: Option<ConfidenceConfig>,
|
||||
/// The time recency is measured from, in seconds since the epoch.
|
||||
/// `None` uses the system clock.
|
||||
pub now: Option<f64>,
|
||||
}
|
||||
|
||||
impl SearchOptions {
|
||||
pub fn new(k: usize) -> Self {
|
||||
Self {
|
||||
k,
|
||||
fusion: hybrid::DEFAULT_FUSION,
|
||||
source_channels: None,
|
||||
rerank: None,
|
||||
rerank_pool: 0,
|
||||
confidence: None,
|
||||
now: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_fusion(mut self, fusion: hybrid::Fusion) -> Self {
|
||||
self.fusion = fusion;
|
||||
self
|
||||
}
|
||||
|
||||
/// Search only records from these source channels.
|
||||
pub fn with_sources<S: Into<String>>(mut self, channels: impl IntoIterator<Item = S>) -> Self {
|
||||
self.source_channels = Some(channels.into_iter().map(Into::into).collect());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_rerank(mut self, config: ReRankConfig) -> Self {
|
||||
self.rerank = Some(config);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_confidence(mut self, config: ConfidenceConfig) -> Self {
|
||||
self.confidence = Some(config);
|
||||
self
|
||||
}
|
||||
|
||||
/// Measure recency from `now` (seconds since the epoch) instead of the
|
||||
/// system clock — for reproducible results and tests.
|
||||
pub fn at_time(mut self, now: f64) -> Self {
|
||||
self.now = Some(now);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SearchOptions {
|
||||
fn default() -> Self {
|
||||
Self::new(10)
|
||||
}
|
||||
}
|
||||
|
||||
impl HDF5Memory {
|
||||
/// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
|
||||
/// Vector + keyword scoring stage of [`HDF5Memory::search`].
|
||||
///
|
||||
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
||||
/// previous behaviour, also used as the correctness oracle in tests). With
|
||||
/// `hnsw` enabled and an index available, the vector candidates come from an
|
||||
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
||||
/// with BM25 via the shared [`hybrid::merge_vector_keyword`].
|
||||
///
|
||||
/// `exclude`, when given, marks records that must not be returned (1 =
|
||||
/// excluded; it covers tombstones too). The index is over-fetched in
|
||||
/// proportion to how much the mask removes. Surfacing `pool` candidates
|
||||
/// costs the index roughly `pool × M` distance evaluations, while an exact
|
||||
/// scan of the allowed records costs one each — so whenever that scan is
|
||||
/// the cheaper of the two it is used instead, and it is also the fallback
|
||||
/// if the pool comes back with too few allowed hits (the allowed records
|
||||
/// sit away from the query). A filtered search never comes back short.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
@@ -22,16 +111,30 @@ impl HDF5Memory {
|
||||
bm25: &bm25::BM25Index,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
exclude: Option<&[u8]>,
|
||||
) -> Vec<(usize, f32)> {
|
||||
self.ensure_hnsw_fresh();
|
||||
let n = self.cache.len();
|
||||
// Over-fetch so the merge sees a useful vector pool. `ef` is
|
||||
// configurable, but the pool the fusion stage sees is not tied to it:
|
||||
// a caller lowering `ef` for speed should not silently narrow what
|
||||
// fusion has to work with.
|
||||
let mut pool = (k * 8).max(64);
|
||||
let mut allowed = n;
|
||||
if let Some(ex) = exclude {
|
||||
allowed = ex.iter().filter(|&&e| e == 0).count();
|
||||
if allowed == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
// Expect `pool` allowed hits if the filter is independent of the
|
||||
// query's neighbourhood.
|
||||
pool = pool.saturating_mul(n).div_ceil(allowed);
|
||||
if allowed <= pool.saturating_mul(self.hnsw_m()) {
|
||||
return self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex);
|
||||
}
|
||||
}
|
||||
match self.hnsw.as_ref() {
|
||||
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
||||
// distance from the index converts back to similarity (1 - d).
|
||||
// `ef` is configurable, but the pool the fusion stage sees is
|
||||
// not tied to it: a caller lowering `ef` for speed should not
|
||||
// silently narrow what fusion has to work with.
|
||||
let pool = (k * 8).max(64);
|
||||
let ef = self.hnsw_ef_search(k).max(pool);
|
||||
let candidates = index.search(query_embedding, pool, ef);
|
||||
// A quantised index returns approximate distances, and no
|
||||
@@ -42,6 +145,7 @@ impl HDF5Memory {
|
||||
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
|
||||
let vec_scores: Vec<(usize, f32)> = candidates
|
||||
.into_iter()
|
||||
.filter(|(id, _)| exclude.is_none_or(|ex| ex[*id] == 0))
|
||||
.map(|(id, dist)| {
|
||||
let score = if exact {
|
||||
crate::vector_search::cosine_similarity(
|
||||
@@ -56,10 +160,54 @@ impl HDF5Memory {
|
||||
.collect();
|
||||
// Fusion normalises over every keyword match, so it needs all
|
||||
// the scores — but not ranked.
|
||||
let kw_scores = bm25.scores(query_text);
|
||||
let mut kw_scores = bm25.scores(query_text);
|
||||
if let Some(ex) = exclude {
|
||||
if vec_scores.len() < k.min(allowed) {
|
||||
// The allowed records are not where the index looked.
|
||||
return self.exact_masked_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
ex,
|
||||
);
|
||||
}
|
||||
kw_scores.retain(|(id, _)| ex[*id] == 0);
|
||||
}
|
||||
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
_ => hybrid::hybrid_search_fused(
|
||||
_ => match exclude {
|
||||
Some(ex) => {
|
||||
self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex)
|
||||
}
|
||||
None => hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
exclude: Option<&[u8]>,
|
||||
) -> Vec<(usize, f32)> {
|
||||
match exclude {
|
||||
Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex),
|
||||
None => hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
@@ -72,25 +220,33 @@ impl HDF5Memory {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
/// Exact hybrid search over the records `exclude` leaves (0 = allowed).
|
||||
fn exact_masked_search(
|
||||
&self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
exclude: &[u8],
|
||||
) -> Vec<(usize, f32)> {
|
||||
hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
)
|
||||
let vec_scores =
|
||||
hybrid::exact_vector_scores(query_embedding, &self.cache.embeddings, exclude);
|
||||
let mut kw_scores = bm25.scores(query_text);
|
||||
kw_scores.retain(|(id, _)| exclude.get(*id) == Some(&0));
|
||||
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
|
||||
/// The exclusion mask for a source-channel filter: 1 for a tombstoned
|
||||
/// record or one from a channel not in `channels`.
|
||||
fn source_mask(&self, channels: &[String]) -> Vec<u8> {
|
||||
let allowed: HashSet<&str> = channels.iter().map(String::as_str).collect();
|
||||
self.cache
|
||||
.source_channels
|
||||
.iter()
|
||||
.zip(&self.cache.tombstones)
|
||||
.map(|(ch, &t)| u8::from(t != 0 || !allowed.contains(ch.as_str())))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
||||
@@ -125,12 +281,53 @@ impl HDF5Memory {
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
self.search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&SearchOptions::new(k).with_fusion(fusion),
|
||||
)
|
||||
}
|
||||
|
||||
/// Hybrid search with optional source filtering, re-ranking and
|
||||
/// confidence rejection — see [`SearchOptions`].
|
||||
///
|
||||
/// Stages, in order: vector + keyword retrieval over the records the
|
||||
/// source filter allows; fusion; scaling by Hebbian activation; re-ranking
|
||||
/// (if on) of a `rerank_pool` of candidates; confidence rejection (if on);
|
||||
/// the top `k`. The records returned with a positive score get their
|
||||
/// Hebbian boost.
|
||||
pub fn search(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
options: &SearchOptions,
|
||||
) -> Vec<SearchResult> {
|
||||
let k = options.k;
|
||||
let fetch = match options.rerank {
|
||||
Some(_) if options.rerank_pool > 0 => options.rerank_pool.max(k),
|
||||
Some(_) => k.saturating_mul(3).max(10),
|
||||
None => k,
|
||||
};
|
||||
let exclude = options
|
||||
.source_channels
|
||||
.as_deref()
|
||||
.map(|channels| self.source_mask(channels));
|
||||
|
||||
// The keyword index lives for the life of the store and is updated
|
||||
// incrementally. Take it out for the duration of the call so the
|
||||
// vector stage can borrow `self` mutably, then put it back.
|
||||
self.ensure_bm25_fresh();
|
||||
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
||||
let scored = self.vector_keyword_search(query_embedding, query_text, &bm25, fusion, k);
|
||||
let scored = self.vector_keyword_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&bm25,
|
||||
options.fusion,
|
||||
fetch,
|
||||
exclude.as_deref(),
|
||||
);
|
||||
self.bm25 = Some(bm25);
|
||||
|
||||
let mut results: Vec<SearchResult> = scored
|
||||
.into_iter()
|
||||
.map(|(idx, score)| {
|
||||
@@ -154,6 +351,25 @@ impl HDF5Memory {
|
||||
.then(a.index.cmp(&b.index))
|
||||
});
|
||||
|
||||
if let Some(config) = &options.rerank {
|
||||
results = Self::rerank_results(results, config, options.now);
|
||||
}
|
||||
if let Some(config) = &options.confidence {
|
||||
let scored: Vec<ScoredResult> = results
|
||||
.iter()
|
||||
.map(|r| ScoredResult {
|
||||
index: r.index,
|
||||
score: r.score,
|
||||
})
|
||||
.collect();
|
||||
let keep: HashSet<usize> = reject_low_confidence(&scored, config)
|
||||
.into_iter()
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
results.retain(|r| keep.contains(&r.index));
|
||||
}
|
||||
results.truncate(k);
|
||||
|
||||
// Only reinforce records that actually matched. When fewer than `k`
|
||||
// records are relevant, the rest of the list is zero-score filler;
|
||||
// boosting it would teach the store that arbitrary records are
|
||||
@@ -164,11 +380,45 @@ impl HDF5Memory {
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
self.apply_hebbian_boost(&hit_indices);
|
||||
self.bm25 = Some(bm25);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Reorder by the re-ranker's combined score, which also becomes each
|
||||
/// result's `score`.
|
||||
fn rerank_results(
|
||||
results: Vec<SearchResult>,
|
||||
config: &ReRankConfig,
|
||||
now: Option<f64>,
|
||||
) -> Vec<SearchResult> {
|
||||
let now = now.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs_f64())
|
||||
.unwrap_or(0.0)
|
||||
});
|
||||
let inputs: Vec<RerankInput> = results
|
||||
.iter()
|
||||
.map(|r| RerankInput {
|
||||
index: r.index,
|
||||
timestamp: r.timestamp,
|
||||
source_channel: r.source_channel.clone(),
|
||||
raw_activation: r.activation,
|
||||
relevance: r.score,
|
||||
})
|
||||
.collect();
|
||||
let mut by_index: std::collections::HashMap<usize, SearchResult> =
|
||||
results.into_iter().map(|r| (r.index, r)).collect();
|
||||
rerank(&inputs, config, now)
|
||||
.into_iter()
|
||||
.filter_map(|rr| {
|
||||
let mut r = by_index.remove(&rr.index)?;
|
||||
r.score = rr.combined_score;
|
||||
Some(r)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Reinforce the records a query returned. The new weights are persisted by
|
||||
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
||||
/// by rewriting the whole store inside the query, which is what made
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
//! `MemoryConfig::float16`: embeddings stored as IEEE half precision.
|
||||
//!
|
||||
//! The setting used to be recorded in `/meta` and otherwise ignored — the
|
||||
//! embeddings dataset was always `f32`. These tests pin what it now does: the
|
||||
//! dataset is `float16`, the in-memory cache holds exactly the values the file
|
||||
//! holds (so search results survive a reopen bit for bit), and a value half
|
||||
//! precision cannot represent is refused rather than stored as infinity.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, MemoryError};
|
||||
use clawhdf5_format::float16::round_to_f16;
|
||||
use tempfile::TempDir;
|
||||
|
||||
const DIM: usize = 64;
|
||||
|
||||
/// Deterministic, embedding-like unit vectors.
|
||||
fn embedding(seed: u64) -> Vec<f32> {
|
||||
let mut x = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||
let v: Vec<f32> = (0..DIM)
|
||||
.map(|_| {
|
||||
x ^= x << 13;
|
||||
x ^= x >> 7;
|
||||
x ^= x << 17;
|
||||
(x >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
||||
})
|
||||
.collect();
|
||||
let norm = v.iter().map(|a| a * a).sum::<f32>().sqrt();
|
||||
v.iter().map(|a| a / norm).collect()
|
||||
}
|
||||
|
||||
fn entry(i: u64) -> MemoryEntry {
|
||||
MemoryEntry {
|
||||
chunk: format!("memory number {i} about topic {}", i % 7),
|
||||
embedding: embedding(i),
|
||||
source_channel: "test".into(),
|
||||
timestamp: i as f64,
|
||||
session_id: "s".into(),
|
||||
tags: format!("t{i}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn config(dir: &TempDir, name: &str, float16: bool) -> MemoryConfig {
|
||||
let mut c = MemoryConfig::new(dir.path().join(name), "agent", DIM);
|
||||
c.float16 = float16;
|
||||
c
|
||||
}
|
||||
|
||||
fn embeddings_dtype_and_values(path: &Path) -> (String, Vec<f32>) {
|
||||
let file = clawhdf5::File::open(path).unwrap();
|
||||
let ds = file.dataset("memory/embeddings").unwrap();
|
||||
(format!("{:?}", ds.dtype().unwrap()), ds.read_f32().unwrap())
|
||||
}
|
||||
|
||||
fn search_bits(m: &mut HDF5Memory, q: u64) -> Vec<(usize, u32)> {
|
||||
m.hybrid_search(&embedding(q), "memory topic 3", 0.4, 0.6, 10)
|
||||
.iter()
|
||||
.map(|r| (r.index, r.score.to_bits()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn float16_store_writes_half_precision_and_reopens_identically() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
// Two identical stores. Search is not read-only (it boosts the Hebbian
|
||||
// activation of what it returns, and checkpoints persist that), so each
|
||||
// is queried exactly once: one live, one after a checkpoint and reopen.
|
||||
let live_cfg = config(&dir, "live.h5", true);
|
||||
let cfg = config(&dir, "f16.h5", true);
|
||||
let path: PathBuf = cfg.path.clone();
|
||||
|
||||
let mut live = HDF5Memory::create(live_cfg).unwrap();
|
||||
live.save_batch((0..200).map(entry).collect()).unwrap();
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
m.save_batch((0..200).map(entry).collect()).unwrap();
|
||||
drop(m);
|
||||
|
||||
// On disk: a genuine float16 dataset holding the rounded inputs.
|
||||
let (dtype, values) = embeddings_dtype_and_values(&path);
|
||||
assert_eq!(dtype, "Other(\"float16\")");
|
||||
let expected: Vec<u32> = (0..200)
|
||||
.flat_map(|i| embedding(i).into_iter().map(|v| round_to_f16(v).to_bits()))
|
||||
.collect();
|
||||
let got: Vec<u32> = values.iter().map(|v| v.to_bits()).collect();
|
||||
assert_eq!(got, expected);
|
||||
|
||||
// Reopened, the store answers exactly as the live one does: the cache
|
||||
// held the half-rounded values before the checkpoint.
|
||||
let mut reopened = HDF5Memory::open(&path).unwrap();
|
||||
for q in 0..5 {
|
||||
assert_eq!(
|
||||
search_bits(&mut live, 1000 + q),
|
||||
search_bits(&mut reopened, 1000 + q),
|
||||
"query {q}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn float16_halves_the_embeddings_on_disk() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut sizes = Vec::new();
|
||||
for float16 in [false, true] {
|
||||
let cfg = config(&dir, &format!("s{float16}.h5"), float16);
|
||||
let path = cfg.path.clone();
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
m.save_batch((0..2000).map(entry).collect()).unwrap();
|
||||
drop(m);
|
||||
sizes.push(std::fs::metadata(&path).unwrap().len());
|
||||
}
|
||||
let embedding_bytes_f32 = (2000 * DIM * 4) as u64;
|
||||
let saved = sizes[0] - sizes[1];
|
||||
// Half of the f32 embeddings, give or take metadata and alignment.
|
||||
assert!(
|
||||
saved.abs_diff(embedding_bytes_f32 / 2) < 16 * 1024,
|
||||
"f32 {} B, f16 {} B, saved {saved} B, expected ~{} B",
|
||||
sizes[0],
|
||||
sizes[1],
|
||||
embedding_bytes_f32 / 2
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn f32_store_is_unchanged() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let cfg = config(&dir, "f32.h5", false);
|
||||
let path = cfg.path.clone();
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
m.save_batch((0..50).map(entry).collect()).unwrap();
|
||||
drop(m);
|
||||
let (dtype, values) = embeddings_dtype_and_values(&path);
|
||||
assert_eq!(dtype, "F32");
|
||||
let expected: Vec<f32> = (0..50).flat_map(embedding).collect();
|
||||
assert_eq!(values, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn out_of_range_values_are_refused_not_stored_as_infinity() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut cfg = config(&dir, "range.h5", true);
|
||||
cfg.wal_enabled = true;
|
||||
let path = cfg.path.clone();
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
m.save(entry(1)).unwrap();
|
||||
|
||||
let mut bad = entry(2);
|
||||
bad.embedding[5] = 70_000.0;
|
||||
match m.save(bad.clone()) {
|
||||
Err(MemoryError::InvalidEntry(msg)) => assert!(msg.contains("embedding[5]"), "{msg}"),
|
||||
other => panic!("expected InvalidEntry, got {other:?}"),
|
||||
}
|
||||
assert!(matches!(
|
||||
m.save_or_update(bad.clone()),
|
||||
Err(MemoryError::InvalidEntry(_))
|
||||
));
|
||||
// A batch is all or nothing.
|
||||
assert!(matches!(
|
||||
m.save_batch(vec![entry(3), bad.clone(), entry(4)]),
|
||||
Err(MemoryError::InvalidEntry(_))
|
||||
));
|
||||
assert_eq!(m.count(), 1);
|
||||
|
||||
// The largest finite half, and values that round down to it, are fine.
|
||||
let mut edge = entry(5);
|
||||
edge.embedding[0] = 65504.0;
|
||||
edge.embedding[1] = -65519.0;
|
||||
m.save(edge).unwrap();
|
||||
assert_eq!(m.count(), 2);
|
||||
drop(m);
|
||||
|
||||
// Nothing rejected reached the WAL or the file.
|
||||
let m = HDF5Memory::open(&path).unwrap();
|
||||
assert_eq!(m.count(), 2);
|
||||
|
||||
// An f32 store takes the same value as it always did.
|
||||
let mut m32 = HDF5Memory::create(config(&dir, "range32.h5", false)).unwrap();
|
||||
m32.save(bad).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wal_replay_rounds_like_a_live_save() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut cfg = config(&dir, "wal.h5", true);
|
||||
cfg.wal_enabled = true;
|
||||
cfg.wal_max_entries = 10_000; // keep everything in the WAL
|
||||
let path = cfg.path.clone();
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
for i in 0..30 {
|
||||
m.save(entry(i)).unwrap();
|
||||
}
|
||||
let live = search_bits(&mut m, 77);
|
||||
|
||||
// Crash image: the .h5 is still the empty checkpoint; everything is in
|
||||
// the WAL, which holds the caller's f32 values.
|
||||
let crash = TempDir::new().unwrap();
|
||||
let image = crash.path().join("image.h5");
|
||||
std::fs::copy(&path, &image).unwrap();
|
||||
std::fs::copy(
|
||||
path.with_extension("h5.wal"),
|
||||
image.with_extension("h5.wal"),
|
||||
)
|
||||
.unwrap();
|
||||
drop(m);
|
||||
|
||||
let mut recovered = HDF5Memory::open(&image).unwrap();
|
||||
assert_eq!(recovered.count(), 30);
|
||||
assert_eq!(search_bits(&mut recovered, 77), live);
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
//! An agent store is a standard HDF5 file: h5py can open it and read every
|
||||
//! dataset.
|
||||
//!
|
||||
//! It could not: the float datatype's sign-bit position was hard-coded for
|
||||
//! f64, so every f32 dataset (embeddings, norms, activation weights) made
|
||||
//! libhdf5 refuse the file with "sign bit position out of bounds".
|
||||
|
||||
use std::process::Command;
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
|
||||
fn python() -> String {
|
||||
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
|
||||
}
|
||||
|
||||
fn h5py_available() -> bool {
|
||||
Command::new(python())
|
||||
.args(["-c", "import h5py"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_reads_every_dataset_of_an_agent_store() {
|
||||
if !h5py_available() {
|
||||
assert!(
|
||||
std::env::var("CLAWHDF5_REQUIRE_INTEROP").as_deref() != Ok("1"),
|
||||
"CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available"
|
||||
);
|
||||
eprintln!("SKIP: python3 with h5py not available");
|
||||
return;
|
||||
}
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
for float16 in [false, true] {
|
||||
let path = dir.path().join(format!("store_{float16}.h5"));
|
||||
let mut cfg = MemoryConfig::new(path.clone(), "agent", 8);
|
||||
cfg.float16 = float16;
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
// save_batch checkpoints, so the records are in the .h5, not the WAL.
|
||||
m.save_batch(
|
||||
(0..20)
|
||||
.map(|i| MemoryEntry {
|
||||
chunk: format!("memory {i}"),
|
||||
embedding: (0..8).map(|j| ((i * 8 + j) as f32).sin()).collect(),
|
||||
source_channel: "test".into(),
|
||||
timestamp: i as f64,
|
||||
session_id: "s".into(),
|
||||
tags: String::new(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.unwrap();
|
||||
drop(m);
|
||||
|
||||
// Exact expected values, as bits: numpy's sin need not match Rust's
|
||||
// to the last place.
|
||||
let bits = (0..160)
|
||||
.map(|k| (k as f32).sin().to_bits().to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
want = np.float16 if {py_bool} else np.float32
|
||||
with h5py.File("{path}", "r") as f:
|
||||
names = []
|
||||
f.visititems(lambda n, o: names.append(n) if isinstance(o, h5py.Dataset) else None)
|
||||
for n in names:
|
||||
f[n][()] # every dataset must decode
|
||||
e = f["memory/embeddings"]
|
||||
assert e.dtype == want, e.dtype
|
||||
assert e.shape == (20, 8), e.shape
|
||||
ref = np.array([{bits}], dtype=np.uint32).view(np.float32).astype(want).reshape(20, 8)
|
||||
assert (e[()] == ref).all()
|
||||
assert f["memory/norms"].dtype == np.float32
|
||||
print(len(names))
|
||||
"#,
|
||||
py_bool = if float16 { "True" } else { "False" },
|
||||
path = path.display()
|
||||
);
|
||||
let out = Command::new(python())
|
||||
.args(["-c", &script])
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
out.status.success(),
|
||||
"float16={float16}: {}",
|
||||
String::from_utf8_lossy(&out.stderr)
|
||||
);
|
||||
let n: usize = String::from_utf8_lossy(&out.stdout).trim().parse().unwrap();
|
||||
assert!(n >= 10, "only {n} datasets");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
//! `HDF5Memory::search` with `SearchOptions`: source filtering, re-ranking and
|
||||
//! confidence rejection in the store's own search path.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||
use clawhdf5_agent::reranker::ReRankConfig;
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions, hybrid};
|
||||
use tempfile::TempDir;
|
||||
|
||||
const DIM: usize = 32;
|
||||
const N: usize = 3000;
|
||||
const CLUSTERS: usize = 20;
|
||||
|
||||
struct Rng(u64);
|
||||
impl Rng {
|
||||
fn next(&mut self) -> u64 {
|
||||
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
let mut z = self.0;
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
fn unit(&mut self) -> f32 {
|
||||
(self.next() >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize(v: &mut [f32]) {
|
||||
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
v.iter_mut().for_each(|x| *x /= n);
|
||||
}
|
||||
|
||||
struct Data {
|
||||
vectors: Vec<Vec<f32>>,
|
||||
cluster: Vec<usize>,
|
||||
centres: Vec<Vec<f32>>,
|
||||
}
|
||||
|
||||
fn data() -> Data {
|
||||
let mut rng = Rng(42);
|
||||
let centres: Vec<Vec<f32>> = (0..CLUSTERS)
|
||||
.map(|_| {
|
||||
let mut c: Vec<f32> = (0..DIM).map(|_| rng.unit()).collect();
|
||||
normalize(&mut c);
|
||||
c
|
||||
})
|
||||
.collect();
|
||||
let mut vectors = Vec::new();
|
||||
let mut cluster = Vec::new();
|
||||
for i in 0..N {
|
||||
let c = i % CLUSTERS;
|
||||
let mut v: Vec<f32> = centres[c].iter().map(|x| x + rng.unit() * 0.3).collect();
|
||||
normalize(&mut v);
|
||||
vectors.push(v);
|
||||
cluster.push(c);
|
||||
}
|
||||
Data {
|
||||
vectors,
|
||||
cluster,
|
||||
centres,
|
||||
}
|
||||
}
|
||||
|
||||
/// Channel of record `i` for a filter keeping `percent`% of the store at
|
||||
/// random (independent of the vectors).
|
||||
fn random_channel(i: usize, rng_seed: u64, percent: u64) -> String {
|
||||
let mut r = Rng(rng_seed ^ (i as u64 * 7919));
|
||||
if r.next() % 100 < percent {
|
||||
"keep".into()
|
||||
} else {
|
||||
"other".into()
|
||||
}
|
||||
}
|
||||
|
||||
fn build(data: &Data, channel: impl Fn(usize) -> String) -> (TempDir, HDF5Memory) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM);
|
||||
cfg.hebbian_boost = 0.0; // every query sees the same store
|
||||
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||
let entries = data
|
||||
.vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| MemoryEntry {
|
||||
chunk: format!("record {i} cluster {}", data.cluster[i]),
|
||||
embedding: v.clone(),
|
||||
source_channel: channel(i),
|
||||
timestamp: i as f64,
|
||||
session_id: "s".into(),
|
||||
tags: format!("t{i}"),
|
||||
})
|
||||
.collect();
|
||||
m.save_batch(entries).unwrap();
|
||||
(dir, m)
|
||||
}
|
||||
|
||||
/// Exact top-k by cosine among the records `allowed` keeps.
|
||||
fn exact_top(data: &Data, q: &[f32], k: usize, allowed: impl Fn(usize) -> bool) -> Vec<usize> {
|
||||
let mut s: Vec<(usize, f32)> = (0..N)
|
||||
.filter(|&i| allowed(i))
|
||||
.map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||
.collect();
|
||||
s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
s.into_iter().take(k).map(|(i, _)| i).collect()
|
||||
}
|
||||
|
||||
fn query(data: &Data, i: usize) -> Vec<f32> {
|
||||
let mut rng = Rng(1000 + i as u64);
|
||||
let mut q: Vec<f32> = data.centres[i % CLUSTERS]
|
||||
.iter()
|
||||
.map(|x| x + rng.unit() * 0.3)
|
||||
.collect();
|
||||
normalize(&mut q);
|
||||
q
|
||||
}
|
||||
|
||||
fn vector_only(k: usize) -> SearchOptions {
|
||||
SearchOptions::new(k).with_fusion(hybrid::Fusion::Weighted {
|
||||
vector: 1.0,
|
||||
keyword: 0.0,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_filter_returns_only_allowed_records_and_a_full_page() {
|
||||
let d = data();
|
||||
// At N = 3000 and k = 10 the index serves a filter only when that is
|
||||
// cheaper than scanning the allowed records: pool = 80 * N / allowed
|
||||
// candidates at ~M = 16 distances each, against `allowed` distances. So
|
||||
// 90% goes through the index, 50% and 1% to the exact scan.
|
||||
for percent in [90, 50, 1] {
|
||||
let (_dir, mut m) = build(&d, |i| random_channel(i, 5, percent));
|
||||
let allowed = |i: usize| random_channel(i, 5, percent) == "keep";
|
||||
let mut hits = 0;
|
||||
for qi in 0..40 {
|
||||
let q = query(&d, qi);
|
||||
let got = m.search(&q, "", &vector_only(10).with_sources(["keep"]));
|
||||
assert_eq!(got.len(), 10, "{percent}%: short page");
|
||||
assert!(got.iter().all(|r| r.source_channel == "keep"));
|
||||
let want: HashSet<usize> = exact_top(&d, &q, 10, allowed).into_iter().collect();
|
||||
hits += got.iter().filter(|r| want.contains(&r.index)).count();
|
||||
}
|
||||
let recall = hits as f64 / 400.0;
|
||||
let floor = if percent == 90 { 0.95 } else { 1.0 };
|
||||
assert!(recall >= floor, "{percent}%: recall@10 {recall}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_away_from_the_query_falls_back_to_an_exact_scan() {
|
||||
// Channel = cluster, and the filter keeps two clusters (10% of the
|
||||
// store) that are not the query's: the index's neighbourhood of the
|
||||
// query holds none of them. The search must still return the exact
|
||||
// top 10 among the allowed records, not a short or empty page.
|
||||
let d = data();
|
||||
let (_dir, mut m) = build(&d, |i| format!("c{}", d.cluster[i]));
|
||||
for qi in 0..20 {
|
||||
let q = query(&d, qi);
|
||||
let a = format!("c{}", (qi + 7) % CLUSTERS);
|
||||
let b = format!("c{}", (qi + 13) % CLUSTERS);
|
||||
let got: Vec<usize> = m
|
||||
.search(
|
||||
&q,
|
||||
"",
|
||||
&vector_only(10).with_sources([a.clone(), b.clone()]),
|
||||
)
|
||||
.iter()
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
let want = exact_top(&d, &q, 10, |i| {
|
||||
let c = format!("c{}", d.cluster[i]);
|
||||
c == a || c == b
|
||||
});
|
||||
assert_eq!(got, want, "query {qi}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_edge_cases() {
|
||||
let d = data();
|
||||
let (_dir, mut m) = build(&d, |i| random_channel(i, 9, 50));
|
||||
let q = query(&d, 0);
|
||||
assert!(
|
||||
m.search(
|
||||
&q,
|
||||
"cluster",
|
||||
&SearchOptions::new(10).with_sources(Vec::<String>::new())
|
||||
)
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
m.search(
|
||||
&q,
|
||||
"cluster",
|
||||
&SearchOptions::new(10).with_sources(["nope"])
|
||||
)
|
||||
.is_empty()
|
||||
);
|
||||
// Keyword matches from other channels are filtered too.
|
||||
let got = m.search(
|
||||
&q,
|
||||
"record cluster",
|
||||
&SearchOptions::new(50).with_sources(["keep"]),
|
||||
);
|
||||
assert_eq!(got.len(), 50);
|
||||
assert!(got.iter().all(|r| r.source_channel == "keep"));
|
||||
// Deleted records never come back, filtered or not.
|
||||
let first = got[0].index;
|
||||
m.delete(first).unwrap();
|
||||
let again = m.search(
|
||||
&q,
|
||||
"record cluster",
|
||||
&SearchOptions::new(50).with_sources(["keep"]),
|
||||
);
|
||||
assert!(again.iter().all(|r| r.index != first));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plain_options_equal_hybrid_search_with() {
|
||||
// Two identical stores, so neither query sees the other's boosts.
|
||||
let d = data();
|
||||
let (_a, mut a) = build(&d, |i| random_channel(i, 3, 50));
|
||||
let (_b, mut b) = build(&d, |i| random_channel(i, 3, 50));
|
||||
for qi in 0..10 {
|
||||
let q = query(&d, qi);
|
||||
let x: Vec<(usize, u32)> = a
|
||||
.search(&q, "record cluster 3", &SearchOptions::new(10))
|
||||
.iter()
|
||||
.map(|r| (r.index, r.score.to_bits()))
|
||||
.collect();
|
||||
let y: Vec<(usize, u32)> = b
|
||||
.hybrid_search_with(&q, "record cluster 3", hybrid::DEFAULT_FUSION, 10)
|
||||
.iter()
|
||||
.map(|r| (r.index, r.score.to_bits()))
|
||||
.collect();
|
||||
assert_eq!(x, y);
|
||||
}
|
||||
}
|
||||
|
||||
fn small_store(entries: &[(&str, &str, f64)]) -> (TempDir, HDF5Memory) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("r.h5"), "a", 4)).unwrap();
|
||||
m.save_batch(
|
||||
entries
|
||||
.iter()
|
||||
.map(|(chunk, channel, ts)| MemoryEntry {
|
||||
chunk: chunk.to_string(),
|
||||
embedding: vec![1.0, 0.0, 0.0, 0.0],
|
||||
source_channel: channel.to_string(),
|
||||
timestamp: *ts,
|
||||
session_id: "s".into(),
|
||||
tags: String::new(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.unwrap();
|
||||
(dir, m)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rerank_breaks_relevance_ties_by_recency() {
|
||||
// Identical text and vectors, so retrieval ties; re-ranking must put the
|
||||
// newer record first and report the combined score.
|
||||
let now = 1_000_000.0;
|
||||
let (_d, mut m) = small_store(&[
|
||||
("user prefers dark mode", "chat", now - 30.0 * 86_400.0),
|
||||
("user prefers dark mode", "chat", now - 60.0),
|
||||
]);
|
||||
let q = [1.0, 0.0, 0.0, 0.0];
|
||||
let plain = m.search(&q, "dark mode", &SearchOptions::new(2));
|
||||
assert_eq!(plain[0].index, 0, "ties break by index without re-ranking");
|
||||
let reranked = m.search(
|
||||
&q,
|
||||
"dark mode",
|
||||
&SearchOptions::new(2)
|
||||
.with_rerank(ReRankConfig::default())
|
||||
.at_time(now),
|
||||
);
|
||||
assert_eq!(reranked[0].index, 1);
|
||||
assert!(reranked[0].score > reranked[1].score);
|
||||
assert_ne!(reranked[0].score.to_bits(), plain[0].score.to_bits());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confidence_rejects_when_nothing_is_good_enough() {
|
||||
let (_d, mut m) = small_store(&[("alpha", "chat", 0.0), ("beta", "chat", 0.0)]);
|
||||
let q = [1.0, 0.0, 0.0, 0.0];
|
||||
let strict = ConfidenceConfig {
|
||||
min_score: 10.0,
|
||||
..ConfidenceConfig::default()
|
||||
};
|
||||
assert!(
|
||||
m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(strict))
|
||||
.is_empty()
|
||||
);
|
||||
let lenient = ConfidenceConfig {
|
||||
min_score: 0.0,
|
||||
min_gap: f32::INFINITY,
|
||||
max_results: 1,
|
||||
};
|
||||
assert_eq!(
|
||||
m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(lenient))
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_returned_results_are_reinforced() {
|
||||
// With re-ranking, a pool of max(3k, 10) candidates is retrieved; only
|
||||
// the k returned should gain activation.
|
||||
let d = data();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("h.h5");
|
||||
let mut m = HDF5Memory::create(MemoryConfig::new(path, "a", DIM)).unwrap();
|
||||
m.save_batch(
|
||||
(0..200)
|
||||
.map(|i| MemoryEntry {
|
||||
chunk: format!("record {i}"),
|
||||
embedding: d.vectors[i].clone(),
|
||||
source_channel: "chat".into(),
|
||||
timestamp: i as f64,
|
||||
session_id: "s".into(),
|
||||
tags: String::new(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.unwrap();
|
||||
let q = query(&d, 0);
|
||||
let got = m.search(
|
||||
&q,
|
||||
"record",
|
||||
&SearchOptions::new(3).with_rerank(ReRankConfig::default()),
|
||||
);
|
||||
assert_eq!(got.len(), 3);
|
||||
let returned: HashSet<usize> = got.iter().map(|r| r.index).collect();
|
||||
// A second plain search reports each record's current activation.
|
||||
let all = m.search(&q, "record", &SearchOptions::new(200));
|
||||
for r in &all {
|
||||
let boosted = r.activation > 1.0;
|
||||
assert_eq!(boosted, returned.contains(&r.index), "record {}", r.index);
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,8 @@
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full # + 100K
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
|
||||
//! ```
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
@@ -88,6 +90,9 @@ static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::n
|
||||
/// the memory) instead of f32, to price the recall it costs.
|
||||
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// `--f16-first`: in `--float16-study`, run the float16 store first.
|
||||
static F16_FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// `--rerank`: re-score the candidate pool against the exact vectors before
|
||||
/// taking the top K.
|
||||
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
@@ -483,6 +488,319 @@ fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
}));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Search options study: source filters, re-ranking, confidence rejection
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// `--options-study`: what `HDF5Memory::search`'s options cost and whether a
|
||||
/// filtered search finds the right records. Filters keep 50%, 10% or 1% of
|
||||
/// the store at random, or two whole clusters away from the query (the case
|
||||
/// the index cannot serve, which falls back to an exact scan). Recall is
|
||||
/// vector-only against an exact scan of the allowed records; latency is full
|
||||
/// hybrid search. Hebbian boosting is off.
|
||||
fn options_study(n: usize) {
|
||||
use clawhdf5_agent::SearchOptions;
|
||||
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||
use clawhdf5_agent::hybrid::Fusion;
|
||||
use clawhdf5_agent::reranker::ReRankConfig;
|
||||
|
||||
let data = make_dataset(n, 0x0B7 ^ n as u64);
|
||||
let n_clusters = data.cluster_of.iter().max().map_or(1, |m| m + 1);
|
||||
let mut rng = Rng(5);
|
||||
let bucket_of: Vec<usize> = (0..n).map(|_| rng.below(100)).collect();
|
||||
let bucket = &bucket_of;
|
||||
let query_texts: Vec<String> = data
|
||||
.query_cluster
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
let exact_top = |q: &[f32], allowed: &dyn Fn(usize) -> bool| -> Vec<usize> {
|
||||
let mut s: Vec<(usize, f32)> = (0..n)
|
||||
.filter(|&i| allowed(i))
|
||||
.map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||
.collect();
|
||||
s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
s.into_iter().take(K).map(|(i, _)| i).collect()
|
||||
};
|
||||
|
||||
// Two stores: channel = random bucket, and channel = cluster.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut stores = Vec::new();
|
||||
for by_cluster in [false, true] {
|
||||
let mut rng = Rng(3);
|
||||
let entries: Vec<MemoryEntry> = data
|
||||
.vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| MemoryEntry {
|
||||
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||
embedding: v.clone(),
|
||||
source_channel: if by_cluster {
|
||||
format!("c{}", data.cluster_of[i])
|
||||
} else {
|
||||
format!("b{}", bucket[i])
|
||||
},
|
||||
timestamp: i as f64,
|
||||
session_id: format!("s{}", i % 50),
|
||||
tags: format!("t{i}"),
|
||||
})
|
||||
.collect();
|
||||
let mut config = MemoryConfig::new(
|
||||
dir.path().join(format!("opt_{by_cluster}.h5")),
|
||||
"bench",
|
||||
DIM,
|
||||
);
|
||||
config.hebbian_boost = 0.0;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
mem.save_batch(entries).unwrap();
|
||||
std::hint::black_box(mem.search(&data.queries[0], "", &SearchOptions::new(K)));
|
||||
stores.push(mem);
|
||||
}
|
||||
|
||||
let vector_only = SearchOptions::new(K).with_fusion(Fusion::Weighted {
|
||||
vector: 1.0,
|
||||
keyword: 0.0,
|
||||
});
|
||||
// (label, store, channels for query i, allowed(i, record))
|
||||
type Case<'a> = (
|
||||
String,
|
||||
usize,
|
||||
Box<dyn Fn(usize) -> Option<Vec<String>> + 'a>,
|
||||
Box<dyn Fn(usize, usize) -> bool + 'a>,
|
||||
);
|
||||
let mut cases: Vec<Case> = vec![(
|
||||
"no filter".into(),
|
||||
0,
|
||||
Box::new(|_| None),
|
||||
Box::new(|_, _| true),
|
||||
)];
|
||||
for pct in [50usize, 10, 1] {
|
||||
cases.push((
|
||||
format!("random {pct}%"),
|
||||
0,
|
||||
Box::new(move |_| Some((0..pct).map(|b| format!("b{b}")).collect())),
|
||||
Box::new(move |_, i| bucket[i] < pct),
|
||||
));
|
||||
}
|
||||
let d = &data;
|
||||
let away = move |qi: usize| {
|
||||
let qc = d.query_cluster[qi];
|
||||
[
|
||||
(qc + n_clusters / 3) % n_clusters,
|
||||
(qc + 2 * n_clusters / 3) % n_clusters,
|
||||
]
|
||||
};
|
||||
cases.push((
|
||||
"2 clusters away from the query".into(),
|
||||
1,
|
||||
Box::new(move |qi| Some(away(qi).iter().map(|c| format!("c{c}")).collect())),
|
||||
Box::new(move |qi, i| away(qi).contains(&d.cluster_of[i])),
|
||||
));
|
||||
|
||||
for (label, store, channels, allowed) in &cases {
|
||||
let mem = &mut stores[*store];
|
||||
let mut hits = 0;
|
||||
let mut kept = 0;
|
||||
for (qi, q) in data.queries.iter().enumerate() {
|
||||
let mut opts = vector_only.clone();
|
||||
opts.source_channels = channels(qi);
|
||||
let got = mem.search(q, "", &opts);
|
||||
let want = exact_top(q, &|i| allowed(qi, i));
|
||||
kept += want.len();
|
||||
hits += got.iter().filter(|r| want.contains(&r.index)).count();
|
||||
}
|
||||
let latency = summarize(
|
||||
(0..N_QUERIES)
|
||||
.map(|qi| {
|
||||
let mut opts = SearchOptions::new(K);
|
||||
opts.source_channels = channels(qi);
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
println!(
|
||||
"| {n} | {label} | {:.4} | {:.3} | {:.3} |",
|
||||
hits as f64 / kept.max(1) as f64,
|
||||
millis(latency.p50),
|
||||
millis(latency.p99),
|
||||
);
|
||||
}
|
||||
|
||||
let mem = &mut stores[0];
|
||||
for (label, opts) in [
|
||||
(
|
||||
"re-rank",
|
||||
SearchOptions::new(K).with_rerank(ReRankConfig::default()),
|
||||
),
|
||||
(
|
||||
"re-rank + confidence",
|
||||
SearchOptions::new(K)
|
||||
.with_rerank(ReRankConfig::default())
|
||||
.with_confidence(ConfidenceConfig::default()),
|
||||
),
|
||||
] {
|
||||
let latency = summarize(
|
||||
(0..N_QUERIES)
|
||||
.map(|qi| {
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
println!(
|
||||
"| {n} | {label} | — | {:.3} | {:.3} |",
|
||||
millis(latency.p50),
|
||||
millis(latency.p99)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// float16 study: what does half-precision embedding storage cost?
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// `--float16-study`: the same data in an `f32` store and a `float16` store.
|
||||
/// Reports file size, checkpoint and open time, vector-search recall@10
|
||||
/// against an exact scan of the *original* f32 vectors, how often the two
|
||||
/// stores return the same top 10, and `hybrid_search` latency. Hebbian
|
||||
/// boosting is off, so every query sees the same store.
|
||||
fn float16_study(n: usize) {
|
||||
let data = make_dataset(n, 0xF16 ^ n as u64);
|
||||
let mut rng = Rng(11);
|
||||
let query_texts: Vec<String> = data
|
||||
.query_cluster
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
|
||||
// Exact top K by cosine (the vectors are unit length) on the f32 inputs.
|
||||
let exact: Vec<Vec<usize>> = data
|
||||
.queries
|
||||
.iter()
|
||||
.map(|q| {
|
||||
let mut scored: Vec<(usize, f32)> = data
|
||||
.vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, v.iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
scored.into_iter().take(K).map(|(i, _)| i).collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut per_variant: Vec<(bool, Vec<Vec<usize>>)> = Vec::new();
|
||||
// `--f16-first` swaps the order, to check the numbers do not depend on
|
||||
// which store runs first (page cache, allocator, CPU frequency).
|
||||
let order = if F16_FIRST.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
[true, false]
|
||||
} else {
|
||||
[false, true]
|
||||
};
|
||||
for float16 in order {
|
||||
let path = dir.path().join(format!("f16study_{float16}.h5"));
|
||||
let mut rng = Rng(3);
|
||||
let entries: Vec<MemoryEntry> = data
|
||||
.vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| MemoryEntry {
|
||||
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||
embedding: v.clone(),
|
||||
source_channel: "bench".into(),
|
||||
timestamp: i as f64,
|
||||
session_id: format!("s{}", i % 50),
|
||||
tags: format!("t{i}"),
|
||||
})
|
||||
.collect();
|
||||
let mut config = MemoryConfig::new(path.clone(), "bench", DIM);
|
||||
config.float16 = float16;
|
||||
config.hebbian_boost = 0.0;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
mem.save_batch(entries).unwrap();
|
||||
// Build the indexes, then time a checkpoint that writes everything.
|
||||
std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
|
||||
let t = Instant::now();
|
||||
mem.flush_wal().unwrap();
|
||||
let checkpoint = t.elapsed();
|
||||
drop(mem);
|
||||
let file_bytes = std::fs::metadata(&path).unwrap().len();
|
||||
|
||||
// Median of three opens.
|
||||
let mut opens: Vec<Duration> = (0..3)
|
||||
.map(|_| {
|
||||
let t = Instant::now();
|
||||
let m = HDF5Memory::open(&path).unwrap();
|
||||
let d = t.elapsed();
|
||||
drop(m);
|
||||
d
|
||||
})
|
||||
.collect();
|
||||
opens.sort();
|
||||
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||
|
||||
// Vector-only search: empty text, all weight on the vector stage.
|
||||
let results: Vec<Vec<usize>> = data
|
||||
.queries
|
||||
.iter()
|
||||
.map(|q| {
|
||||
mem.hybrid_search(q, "", 1.0, 0.0, K)
|
||||
.iter()
|
||||
.map(|r| r.index)
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
let hits: usize = results
|
||||
.iter()
|
||||
.zip(&exact)
|
||||
.map(|(got, want)| got.iter().filter(|i| want.contains(i)).count())
|
||||
.sum();
|
||||
let recall = hits as f64 / (K * data.queries.len()) as f64;
|
||||
|
||||
let latency = summarize(
|
||||
(0..N_QUERIES)
|
||||
.map(|i| {
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.hybrid_search(
|
||||
&data.queries[i],
|
||||
&query_texts[i],
|
||||
0.4,
|
||||
0.6,
|
||||
K,
|
||||
));
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
let overlap = match per_variant.first() {
|
||||
Some((_, other)) => {
|
||||
let same: usize = results
|
||||
.iter()
|
||||
.zip(other)
|
||||
.map(|(a, b)| a.iter().filter(|i| b.contains(i)).count())
|
||||
.sum();
|
||||
format!("{:.4}", same as f64 / (K * data.queries.len()) as f64)
|
||||
}
|
||||
None => "—".into(),
|
||||
};
|
||||
println!(
|
||||
"| {n} | {} | {:.1} | {:.0} | {:.1} | {recall:.4} | {overlap} | {:.3} |",
|
||||
if float16 { "float16" } else { "f32" },
|
||||
mib(file_bytes),
|
||||
millis(checkpoint),
|
||||
millis(opens[1]),
|
||||
millis(latency.p50),
|
||||
);
|
||||
per_variant.push((float16, results));
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fusion study: does capping the keyword candidate pool change the ranking?
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -642,6 +960,37 @@ fn main() {
|
||||
}
|
||||
return;
|
||||
}
|
||||
if args.iter().any(|a| a == "--options-study") {
|
||||
println!("## Search options ({DIM}-dim, k = {K}, Hebbian boost off)\n");
|
||||
println!("| N | options | filtered recall@10 | p50 ms | p99 ms |");
|
||||
println!("|---:|---|---:|---:|---:|");
|
||||
for &n in if full {
|
||||
&[10_000, 100_000][..]
|
||||
} else {
|
||||
&[10_000][..]
|
||||
} {
|
||||
options_study(n);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if args.iter().any(|a| a == "--f16-first") {
|
||||
F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
if args.iter().any(|a| a == "--float16-study") {
|
||||
println!("## float16 embedding storage ({DIM}-dim, int8 index, Hebbian boost off)\n");
|
||||
println!(
|
||||
"| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap with the other | hybrid p50 ms |"
|
||||
);
|
||||
println!("|---:|---|---:|---:|---:|---:|---:|---:|");
|
||||
for &n in if full {
|
||||
&[1_000, 10_000, 100_000][..]
|
||||
} else {
|
||||
&[1_000, 10_000][..]
|
||||
} {
|
||||
float16_study(n);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if args.iter().any(|a| a == "--int8") {
|
||||
INT8.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
println!("(int8-quantised index vectors)");
|
||||
|
||||
@@ -36,6 +36,11 @@ enum Commands {
|
||||
/// Accepted for compatibility; int8 is now the default
|
||||
#[arg(long, hide = true, conflicts_with = "f32_index")]
|
||||
quantized_index: bool,
|
||||
/// Store embeddings on disk as IEEE half precision (float16): half
|
||||
/// the bytes, about three significant digits; values must lie within
|
||||
/// ±65504
|
||||
#[arg(long)]
|
||||
float16: bool,
|
||||
},
|
||||
/// Save a memory entry (reads JSON from stdin or --json)
|
||||
Save {
|
||||
@@ -102,9 +107,11 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
wal,
|
||||
f32_index,
|
||||
quantized_index: _,
|
||||
float16,
|
||||
} => {
|
||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
||||
config.wal_enabled = wal;
|
||||
config.float16 = float16;
|
||||
// Only ever switch *off* the library default: assigning the flag
|
||||
// outright would force every CLI-created store back to f32 unless
|
||||
// the caller knew to ask for int8.
|
||||
@@ -120,6 +127,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
"embedding_dim": dim,
|
||||
"wal_enabled": wal,
|
||||
"quantized_index": config_quantized,
|
||||
"float16": float16,
|
||||
"count": mem.count(),
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||
|
||||
@@ -24,6 +24,7 @@ libaec-sys = { path = "../libaec-sys", version = "0.1", optional = true }
|
||||
pco = { version = "1.0", optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
half = { workspace = true }
|
||||
serde_json = "1"
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.7.0" }
|
||||
|
||||
@@ -1076,6 +1076,21 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
||||
) {
|
||||
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||
}
|
||||
// Little-endian half precision (numpy float16): widen directly.
|
||||
if matches!(
|
||||
datatype,
|
||||
Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: DatatypeByteOrder::LittleEndian,
|
||||
..
|
||||
}
|
||||
) {
|
||||
let (halves, _) = raw[..count * 2].as_chunks::<2>();
|
||||
return Ok(halves
|
||||
.iter()
|
||||
.map(|&b| f16_bits_to_f32(u16::from_le_bytes(b)))
|
||||
.collect());
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
let mut result = Vec::with_capacity(count);
|
||||
@@ -1622,36 +1637,7 @@ fn read_f16_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
f16_bits_to_f32(u16::from_le_bytes(buf))
|
||||
}
|
||||
|
||||
/// Convert the bit pattern of an IEEE-754 half (binary16) to an `f32`.
|
||||
fn f16_bits_to_f32(h: u16) -> f32 {
|
||||
let h = h as u32;
|
||||
let sign = (h & 0x8000) << 16;
|
||||
let exp = (h >> 10) & 0x1f;
|
||||
let mant = h & 0x3ff;
|
||||
let bits = if exp == 0 {
|
||||
if mant == 0 {
|
||||
sign // signed zero
|
||||
} else {
|
||||
// Subnormal: normalize into an f32 normal.
|
||||
let mut e: i32 = -1;
|
||||
let mut m = mant;
|
||||
loop {
|
||||
e += 1;
|
||||
m <<= 1;
|
||||
if m & 0x400 != 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let m = m & 0x3ff;
|
||||
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
||||
}
|
||||
} else if exp == 0x1f {
|
||||
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
||||
} else {
|
||||
sign | ((exp + (127 - 15)) << 23) | (mant << 13)
|
||||
};
|
||||
f32::from_bits(bits)
|
||||
}
|
||||
use crate::float16::f16_bits_to_f32;
|
||||
|
||||
fn read_f32_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
let mut buf = [0u8; 4];
|
||||
|
||||
@@ -640,7 +640,8 @@ impl Datatype {
|
||||
mantissa_size,
|
||||
exponent_bias,
|
||||
} => {
|
||||
let mut bf0 = 0x20u8; // bit 5: sign location bit (standard IEEE 754)
|
||||
// Bits 4-5: mantissa normalization = 2 (implied leading 1, IEEE 754).
|
||||
let mut bf0 = 0x20u8;
|
||||
match byte_order {
|
||||
DatatypeByteOrder::BigEndian => {
|
||||
bf0 |= 0x01;
|
||||
@@ -650,9 +651,14 @@ impl Datatype {
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
// bf[1] bits 0-1: mantissa normalization = 2 (MSB not stored, IEEE 754)
|
||||
let bf1 = 0x3fu8; // matching what h5py generates
|
||||
let mut buf = Self::build_header(1, 1, [bf0, bf1, 0], *size);
|
||||
// Bits 8-15: the sign bit's position, the top bit of the value.
|
||||
// This was hard-coded to 63, which is right only for f64: the
|
||||
// HDF5 library rejects any other float with "sign bit position
|
||||
// out of bounds", so every f32 dataset and attribute we wrote
|
||||
// was unreadable by h5py and libhdf5.
|
||||
let sign_location =
|
||||
(u32::from(*bit_offset) + u32::from(*bit_precision)).saturating_sub(1) as u8;
|
||||
let mut buf = Self::build_header(1, 1, [bf0, sign_location, 0], *size);
|
||||
buf.extend_from_slice(&bit_offset.to_le_bytes());
|
||||
buf.extend_from_slice(&bit_precision.to_le_bytes());
|
||||
buf.push(*exponent_location);
|
||||
@@ -818,6 +824,24 @@ fn build_dt_header(class: u8, version: u8, bf: [u8; 3], size: u32) -> Vec<u8> {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn float_sign_location_is_the_top_bit_of_the_value() {
|
||||
// The HDF5 library rejects a float whose sign position is not inside
|
||||
// its precision; this was hard-coded to 63, so every f32 we wrote was
|
||||
// unreadable by h5py. Byte 2 of the message is the sign position.
|
||||
use crate::type_builders::{make_f16_type, make_f32_type, make_f64_type};
|
||||
for (dt, sign) in [
|
||||
(make_f16_type(), 15),
|
||||
(make_f32_type(), 31),
|
||||
(make_f64_type(), 63),
|
||||
] {
|
||||
let bytes = dt.serialize();
|
||||
assert_eq!(bytes[2], sign, "{dt:?}");
|
||||
let (parsed, _) = Datatype::parse(&bytes).unwrap();
|
||||
assert_eq!(parsed, dt);
|
||||
}
|
||||
}
|
||||
|
||||
// Helper to build a fixed-point datatype message
|
||||
fn build_fixed_point(
|
||||
size: u32,
|
||||
|
||||
@@ -86,6 +86,12 @@ pub(crate) fn build_dataset_oh(
|
||||
let mut dl = Vec::new();
|
||||
dl.push(4); // version
|
||||
dl.push(1); // class = contiguous
|
||||
// An empty dataset has no storage: its address must be the undefined
|
||||
// address, as libhdf5 writes it. A real address with size 0 trips
|
||||
// libhdf5's `addr + size <= addr` overflow check, and it refuses the
|
||||
// dataset as "invalid dataset size, likely file corruption" — which made
|
||||
// every store with no sessions or knowledge graph unreadable by h5py.
|
||||
let data_addr = if data_size == 0 { u64::MAX } else { data_addr };
|
||||
dl.extend_from_slice(&data_addr.to_le_bytes());
|
||||
dl.extend_from_slice(&data_size.to_le_bytes());
|
||||
w.add_message(MessageType::DataLayout, dl);
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
//! IEEE-754 half precision (binary16) conversions.
|
||||
//!
|
||||
//! Pure integer bit manipulation, so it works under `no_std` and needs no
|
||||
//! `libm`. The writer ([`crate::type_builders::DatasetBuilder::with_f16_data`]),
|
||||
//! the reader and `clawhdf5-agent`'s half-precision embedding store all use
|
||||
//! these two functions, so a value rounded in memory is bit-for-bit the value
|
||||
//! that reads back from the file.
|
||||
|
||||
/// Largest finite half-precision value. Anything larger in magnitude rounds
|
||||
/// to infinity.
|
||||
pub const F16_MAX: f32 = 65504.0;
|
||||
|
||||
/// Convert an `f32` to the bit pattern of the nearest half-precision value,
|
||||
/// rounding ties to even (the IEEE default, and what numpy and the `half`
|
||||
/// crate do).
|
||||
///
|
||||
/// Values beyond ±[`F16_MAX`] become ±infinity, values too small for a
|
||||
/// subnormal become signed zero, and NaN stays NaN (quiet, payload
|
||||
/// truncated).
|
||||
pub fn f32_to_f16_bits(value: f32) -> u16 {
|
||||
let x = value.to_bits();
|
||||
let sign = (x >> 16) & 0x8000;
|
||||
let exp = x & 0x7F80_0000;
|
||||
let man = x & 0x007F_FFFF;
|
||||
|
||||
// Infinity and NaN.
|
||||
if exp == 0x7F80_0000 {
|
||||
let quiet_nan = if man == 0 { 0 } else { 0x0200 };
|
||||
return (sign | 0x7C00 | quiet_nan | (man >> 13)) as u16;
|
||||
}
|
||||
|
||||
let half_exp = ((exp >> 23) as i32) - 127 + 15;
|
||||
|
||||
// Too large: infinity.
|
||||
if half_exp >= 0x1F {
|
||||
return (sign | 0x7C00) as u16;
|
||||
}
|
||||
|
||||
// Subnormal half, or zero.
|
||||
if half_exp <= 0 {
|
||||
if 14 - half_exp > 24 {
|
||||
return sign as u16;
|
||||
}
|
||||
let man = man | 0x0080_0000; // implicit leading bit
|
||||
let shift = (14 - half_exp) as u32;
|
||||
let mut half_man = man >> shift;
|
||||
let round_bit = 1u32 << (shift - 1);
|
||||
// Round half to even: up if above half, or exactly half and odd.
|
||||
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||
half_man += 1;
|
||||
}
|
||||
return (sign | half_man) as u16;
|
||||
}
|
||||
|
||||
// Normal half. A mantissa carry correctly rolls into the exponent (and
|
||||
// from the largest finite value into infinity).
|
||||
let half = sign | ((half_exp as u32) << 10) | (man >> 13);
|
||||
let round_bit = 0x0000_1000;
|
||||
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||
(half + 1) as u16
|
||||
} else {
|
||||
half as u16
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert the bit pattern of a half-precision value to `f32` (exact: every
|
||||
/// half value is representable as an `f32`).
|
||||
pub fn f16_bits_to_f32(h: u16) -> f32 {
|
||||
let h = h as u32;
|
||||
let sign = (h & 0x8000) << 16;
|
||||
let exp = (h >> 10) & 0x1f;
|
||||
let mant = h & 0x3ff;
|
||||
let bits = if exp == 0 {
|
||||
if mant == 0 {
|
||||
sign // signed zero
|
||||
} else {
|
||||
// Subnormal: normalize into an f32 normal.
|
||||
let mut e: i32 = -1;
|
||||
let mut m = mant;
|
||||
loop {
|
||||
e += 1;
|
||||
m <<= 1;
|
||||
if m & 0x400 != 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let m = m & 0x3ff;
|
||||
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
||||
}
|
||||
} else if exp == 0x1f {
|
||||
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
||||
} else {
|
||||
sign | ((exp + 127 - 15) << 23) | (mant << 13)
|
||||
};
|
||||
f32::from_bits(bits)
|
||||
}
|
||||
|
||||
/// Round an `f32` to the nearest half-precision value, returned as `f32`.
|
||||
pub fn round_to_f16(value: f32) -> f32 {
|
||||
f16_bits_to_f32(f32_to_f16_bits(value))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn every_half_value_round_trips() {
|
||||
for bits in 0..=u16::MAX {
|
||||
let v = f16_bits_to_f32(bits);
|
||||
if v.is_nan() {
|
||||
assert!(f16_bits_to_f32(f32_to_f16_bits(v)).is_nan(), "{bits:#06x}");
|
||||
} else {
|
||||
assert_eq!(f32_to_f16_bits(v), bits, "{bits:#06x} -> {v}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_the_half_crate() {
|
||||
// Every 257th f32 bit pattern (~16.7M values) covers every exponent,
|
||||
// the subnormal range, both signs, ties and the overflow boundary.
|
||||
let mut bits: u32 = 0;
|
||||
loop {
|
||||
let v = f32::from_bits(bits);
|
||||
let ours = f32_to_f16_bits(v);
|
||||
let theirs = half::f16::from_f32(v);
|
||||
if v.is_nan() {
|
||||
assert!(theirs.is_nan() && f16_bits_to_f32(ours).is_nan());
|
||||
} else {
|
||||
assert_eq!(ours, theirs.to_bits(), "{bits:#010x} ({v:e})");
|
||||
assert_eq!(f16_bits_to_f32(ours).to_bits(), theirs.to_f32().to_bits());
|
||||
}
|
||||
match bits.checked_add(257) {
|
||||
Some(b) => bits = b,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rounds_ties_to_even_and_saturates_to_infinity() {
|
||||
// 1 + 2^-11 is exactly halfway between 1.0 and the next half (1 + 2^-10).
|
||||
assert_eq!(round_to_f16(1.0 + 2f32.powi(-11)), 1.0);
|
||||
assert_eq!(
|
||||
round_to_f16(1.0 + 3.0 * 2f32.powi(-11)),
|
||||
1.0 + 2.0 * 2f32.powi(-10)
|
||||
);
|
||||
assert_eq!(round_to_f16(F16_MAX), F16_MAX);
|
||||
assert_eq!(round_to_f16(65520.0), f32::INFINITY); // halfway to 2^16 rounds up
|
||||
assert_eq!(round_to_f16(-1e9), f32::NEG_INFINITY);
|
||||
assert_eq!(round_to_f16(1e-9).to_bits(), 0);
|
||||
assert_eq!(round_to_f16(-1e-9).to_bits(), (-0.0f32).to_bits());
|
||||
}
|
||||
}
|
||||
@@ -72,6 +72,7 @@ pub mod filter_pipeline;
|
||||
pub mod filters;
|
||||
mod filters_szip;
|
||||
pub mod fixed_array;
|
||||
pub mod float16;
|
||||
pub mod fractal_heap;
|
||||
pub mod global_heap;
|
||||
pub mod group_info;
|
||||
|
||||
@@ -56,6 +56,21 @@ pub fn make_f64_type() -> Datatype {
|
||||
}
|
||||
}
|
||||
|
||||
/// IEEE-754 half precision (binary16), little-endian — numpy's `float16`.
|
||||
pub fn make_f16_type() -> Datatype {
|
||||
Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: DatatypeByteOrder::LittleEndian,
|
||||
bit_offset: 0,
|
||||
bit_precision: 16,
|
||||
exponent_location: 10,
|
||||
exponent_size: 5,
|
||||
mantissa_location: 0,
|
||||
mantissa_size: 10,
|
||||
exponent_bias: 15,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn make_f32_type() -> Datatype {
|
||||
Datatype::FloatingPoint {
|
||||
size: 4,
|
||||
@@ -478,6 +493,24 @@ impl DatasetBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Store `data` as IEEE half precision (numpy `float16`), rounding each
|
||||
/// value to the nearest half ([`crate::float16::f32_to_f16_bits`]).
|
||||
/// Half the bytes of [`Self::with_f32_data`], at about three significant
|
||||
/// decimal digits; values beyond ±65504 become ±infinity. Reading it back
|
||||
/// with `read_f32` yields the rounded values exactly.
|
||||
pub fn with_f16_data(&mut self, data: &[f32]) -> &mut Self {
|
||||
self.datatype = Some(make_f16_type());
|
||||
let mut b = Vec::with_capacity(data.len() * 2);
|
||||
for &v in data {
|
||||
b.extend_from_slice(&crate::float16::f32_to_f16_bits(v).to_le_bytes());
|
||||
}
|
||||
self.data = Some(b);
|
||||
if self.shape.is_none() {
|
||||
self.shape = Some(vec![data.len() as u64]);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_i32_data(&mut self, data: &[i32]) -> &mut Self {
|
||||
self.datatype = Some(make_i32_type());
|
||||
let mut b = Vec::with_capacity(data.len() * 4);
|
||||
|
||||
@@ -1387,3 +1387,166 @@ with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Half precision (float16) in both directions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Values that exercise rounding: ties, subnormals, the overflow boundary and
|
||||
/// ordinary embedding-sized components.
|
||||
fn f16_probe_values() -> Vec<f32> {
|
||||
let mut v = vec![
|
||||
0.0,
|
||||
-0.0,
|
||||
1.0,
|
||||
-1.0,
|
||||
0.5,
|
||||
1.0 + 2f32.powi(-11),
|
||||
1.0 + 3.0 * 2f32.powi(-11),
|
||||
65504.0,
|
||||
65519.0,
|
||||
65520.0,
|
||||
-70000.0,
|
||||
6.0e-8,
|
||||
3.0e-8,
|
||||
1.0e-9,
|
||||
1.0e-5,
|
||||
0.1,
|
||||
0.333_333,
|
||||
1234.567,
|
||||
f32::INFINITY,
|
||||
f32::NEG_INFINITY,
|
||||
];
|
||||
// A deterministic spread of embedding-like values.
|
||||
let mut x = 0x2545_F491u32;
|
||||
for _ in 0..4000 {
|
||||
x ^= x << 13;
|
||||
x ^= x >> 17;
|
||||
x ^= x << 5;
|
||||
v.push((x as f32 / u32::MAX as f32 - 0.5) * 0.4);
|
||||
}
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clawhdf5_writes_f16_h5py_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("ours_f16.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values = f16_probe_values();
|
||||
|
||||
let mut fb = FileBuilder::new();
|
||||
fb.create_dataset("plain").with_f16_data(&values);
|
||||
fb.create_dataset("chunked")
|
||||
.with_f16_data(&values)
|
||||
.with_shape(&[values.len() as u64])
|
||||
.with_chunks(&[512])
|
||||
.with_deflate(6);
|
||||
fb.write(&path).unwrap();
|
||||
|
||||
// h5py must see a genuine float16 dataset, and our rounding must agree
|
||||
// with numpy's own float32 -> float16 conversion bit for bit.
|
||||
let input = values
|
||||
.iter()
|
||||
.map(|v| format!("{:?}", v.to_bits()))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
src = np.array([{input}], dtype=np.uint32).view(np.float32)
|
||||
expected = src.astype(np.float16).view(np.uint16)
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
for name in ("plain", "chunked"):
|
||||
d = f[name]
|
||||
assert d.dtype == np.float16, (name, d.dtype)
|
||||
got = d[:].view(np.uint16)
|
||||
bad = np.nonzero(got != expected)[0]
|
||||
assert bad.size == 0, (name, bad[:5], got[bad[:5]], expected[bad[:5]])
|
||||
print("ok")
|
||||
"#
|
||||
);
|
||||
assert_eq!(run_python_output(&script), "ok");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_writes_f16_clawhdf5_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("h5py_f16.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values = f16_probe_values();
|
||||
let input = values
|
||||
.iter()
|
||||
.map(|v| format!("{:?}", v.to_bits()))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
src = np.array([{input}], dtype=np.uint32).view(np.float32).astype(np.float16)
|
||||
with h5py.File("{path_str}", "w") as f:
|
||||
f.create_dataset("plain", data=src)
|
||||
f.create_dataset("chunked", data=src, chunks=(512,), compression="gzip", shuffle=True)
|
||||
f.create_dataset("big_endian", data=src.astype(">f2"))
|
||||
"#
|
||||
);
|
||||
run_python(&script);
|
||||
|
||||
let expected: Vec<u32> = values
|
||||
.iter()
|
||||
.map(|&v| clawhdf5_format::float16::round_to_f16(v).to_bits())
|
||||
.collect();
|
||||
let file = File::open(&path).unwrap();
|
||||
for name in ["plain", "chunked", "big_endian"] {
|
||||
let ds = file.dataset(name).unwrap();
|
||||
assert_eq!(
|
||||
ds.dtype().unwrap(),
|
||||
DType::Other("float16".into()),
|
||||
"{name}"
|
||||
);
|
||||
let got: Vec<u32> = ds.read_f32().unwrap().iter().map(|v| v.to_bits()).collect();
|
||||
assert_eq!(got, expected, "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clawhdf5_writes_f32_h5py_reads() {
|
||||
// Every f32 dataset used to be unreadable by h5py ("sign bit position out
|
||||
// of bounds"): the float datatype's sign position was hard-coded for f64.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("ours_f32.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values: Vec<f32> = vec![1.5, -2.25, 3.0e-7, 65536.5, f32::MAX, -0.0];
|
||||
|
||||
let mut fb = FileBuilder::new();
|
||||
fb.create_dataset("plain").with_f32_data(&values);
|
||||
fb.create_dataset("chunked")
|
||||
.with_f32_data(&values)
|
||||
.with_shape(&[values.len() as u64])
|
||||
.with_chunks(&[4])
|
||||
.with_deflate(6);
|
||||
fb.write(&path).unwrap();
|
||||
|
||||
let bits = values
|
||||
.iter()
|
||||
.map(|v| v.to_bits().to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
expected = np.array([{bits}], dtype=np.uint32)
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
for name in ("plain", "chunked"):
|
||||
d = f[name]
|
||||
assert d.dtype == np.float32, (name, d.dtype)
|
||||
assert (d[:].view(np.uint32) == expected).all(), (name, d[:])
|
||||
print("ok")
|
||||
"#
|
||||
);
|
||||
assert_eq!(run_python_output(&script), "ok");
|
||||
}
|
||||
|
||||
@@ -227,3 +227,55 @@ wrong structure. All four are fixed and covered by interop tests against
|
||||
HDF5 2.0 at sizes that cross each boundary, including paged data blocks.
|
||||
|
||||
Files written by this crate are unaffected — this was purely a read-path bug.
|
||||
|
||||
## Every `f32` dataset we wrote was unreadable by h5py / libhdf5
|
||||
|
||||
**Status:** fixed 2026-09-23, after v2.7.0. **Every
|
||||
release up to and including v2.7.0 is affected** — the encoder was already
|
||||
wrong in v2.1.0.
|
||||
|
||||
The floating-point datatype message carries the position of the sign bit
|
||||
(bits 8–15 of its class bit field). `clawhdf5-format` wrote 63 for every
|
||||
float, which is correct only for `f64`. libhdf5 validates the field, so opening
|
||||
any `f32` dataset written by this crate failed:
|
||||
|
||||
```
|
||||
KeyError: 'Unable to synchronously open object (sign bit position out of bounds)'
|
||||
```
|
||||
|
||||
That covers every agent store (`/memory/embeddings`, `norms` and
|
||||
`activation_weights` are `f32`). `clawhdf5` itself ignores the field on read,
|
||||
and the interop suites only ever wrote `f64` from our side, so nothing here
|
||||
noticed.
|
||||
|
||||
**Fix:** the sign position is computed from the type (`bit_offset +
|
||||
bit_precision - 1`: 15, 31, 63 for half, single, double). Regression tests:
|
||||
`float_sign_location_is_the_top_bit_of_the_value` (byte level),
|
||||
`clawhdf5_writes_f32_h5py_reads` and the agent's
|
||||
`h5py_reads_every_dataset_of_an_agent_store`.
|
||||
|
||||
**Existing files:** an agent store is rewritten in full at every checkpoint, so
|
||||
it becomes readable by h5py at its next checkpoint with a fixed build. Other
|
||||
files with `f32` datasets need to be rewritten.
|
||||
|
||||
## Empty datasets we wrote were unreadable by h5py / libhdf5
|
||||
|
||||
**Status:** fixed 2026-09-23, after v2.7.0. Every
|
||||
release up to and including v2.7.0 is affected.
|
||||
|
||||
A dataset with no elements was written with a real file address and a storage
|
||||
size of 0. libhdf5 guards contiguous storage with an overflow check
|
||||
(`addr + size <= addr`) that is always true when the size is 0, so it rejected
|
||||
the dataset:
|
||||
|
||||
```
|
||||
KeyError: 'Unable to synchronously open object (invalid dataset size, likely file corruption)'
|
||||
```
|
||||
|
||||
In practice: every agent store without sessions or a knowledge graph — the
|
||||
`/sessions` and `/knowledge_graph` datasets are empty until something is added
|
||||
— could not be read by h5py even once the `f32` bug above was fixed. Found by
|
||||
the same agent-store interop test.
|
||||
|
||||
**Fix:** an empty contiguous dataset gets the undefined address (all `0xff`),
|
||||
which is what libhdf5 itself writes.
|
||||
|
||||
Reference in New Issue
Block a user