Skip to main content

borg/vectors/
model.rs

1//! Where vector recall's model and indexes live and whether they are there, without loading either:
2//! what the hooks check before a query, and what `borg doctor` reads (doctor.rs). No dependency on
3//! lessons or the engine.
4
5use std::path::{Path, PathBuf};
6
7/// The Hugging Face repository of the default model and the revision `--fetch` pins.
8pub const MODEL_REPO: &str = "minishlab/potion-multilingual-128M";
9pub const MODEL_REVISION: &str = "73908c3438cf03b6a01bcb9611d62b23d0726f08";
10/// The files `--fetch` downloads, with their sha256 at `MODEL_REVISION`.
11pub const MODEL_FILES: [(&str, &str); 3] = [
12    ("config.json", "595e4cab2093732efd5dbe084fd5c1826b5eea693b73b4c1fd971672867d2e54"),
13    ("tokenizer.json", "19f1909063da3cfe3bd83a782381f040dccea475f4816de11116444a73e1b6a1"),
14    ("model.safetensors", "14b5eb39cb4ce5666da8ad1f3dc6be4346e9b2d601c073302fa0a31bf7943397"),
15];
16
17/// Beside the `SeekStorm` files: which model the vectors come from and which lesson versions they hold.
18pub const MANIFEST: &str = "borg-vectors.json";
19
20/// Off only with `BORG_VECTORS=off` (or a build without the `vectors` feature).
21#[must_use]
22pub fn enabled() -> bool {
23    cfg!(feature = "vectors") && std::env::var("BORG_VECTORS").as_deref() != Ok("off")
24}
25
26/// The model directory: `BORG_VECTOR_MODEL`, default `~/.local/share/borg/models/potion-multilingual-128M`.
27#[must_use]
28pub fn model_dir() -> PathBuf {
29    std::env::var_os("BORG_VECTOR_MODEL")
30        .map_or_else(|| PathBuf::from(std::env::var_os("HOME").unwrap_or_default()).join(".local/share/borg/models/potion-multilingual-128M"), PathBuf::from)
31}
32
33/// Whether the model's files are there (no load).
34#[must_use]
35pub fn model_present(dir: &Path) -> bool {
36    MODEL_FILES.iter().all(|(f, _)| dir.join(f).is_file())
37}
38
39/// The vector index of a lessons file: `vectors/` beside it.
40#[must_use]
41pub fn index_dir(lessons: &Path) -> PathBuf {
42    lessons.with_file_name("vectors")
43}
44
45/// The model's identity in the manifest: its pinned revision (REVISION file) or the matrix's size.
46pub(crate) fn identity(model: &Path) -> String {
47    match std::fs::read_to_string(model.join("REVISION")) {
48        Ok(r) if !r.trim().is_empty() => r.trim().to_string(),
49        _ => format!("size:{}", std::fs::metadata(model.join("model.safetensors")).map_or(0, |m| m.len())),
50    }
51}
52
53/// Whether the index of `lessons` exists (`None` if not) and was built with `model`: the one rule for a
54/// current index, the engine's as doctor's and setup's.
55#[must_use]
56pub fn built_with(lessons: &Path, model: &Path) -> Option<bool> {
57    let index = index_dir(lessons);
58    let manifest: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(index.join(MANIFEST)).ok()?).ok()?;
59    let built = manifest["model"].as_str()?;
60    index.join("index.json").is_file().then(|| built == identity(model))
61}