wip-4b3
CI / test (push) Failing after 3s

This commit is contained in:
Omar Sobh
2026-08-07 06:55:46 -07:00
parent 211b8f8c81
commit 09df5235e0
2 changed files with 16 additions and 1 deletions
+3
View File
@@ -75,3 +75,6 @@ libhdf5-compare = ["hdf5"]
mpi-io = ["clawhdf5-io/mpi-io", "mpi"]
# Real MiniLM embeddings for longmemeval_bench, so the vector stage is not inert.
embeddings = ["candle-core", "candle-nn", "candle-transformers", "tokenizers"]
# CUDA-accelerated embedding. MiniLM on a CPU takes hours over the full
# longmemeval_s haystack; on a GPU it is minutes.
embeddings-cuda = ["embeddings", "candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
@@ -34,7 +34,19 @@ impl Embedder {
/// `config.json` is read when present; otherwise the published MiniLM-L6-v2
/// architecture constants are used, which are pinned rather than guessed.
pub fn load(dir: &Path) -> Result<Self, Box<dyn std::error::Error>> {
let device = Device::Cpu;
// CUDA when the feature is on and a device is actually present; the CPU
// path is correct but roughly two orders of magnitude slower, which is
// the difference between minutes and most of a day on the full haystack.
let device = match Device::new_cuda(0) {
Ok(d) => {
eprintln!("Embedder: CUDA device 0");
d
}
Err(e) => {
eprintln!("Embedder: CPU ({e})");
Device::Cpu
}
};
let weights = dir.join("model.safetensors");
let tok_path = dir.join("tokenizer.json");