use anyhow::{Context, Result, bail};
use fastembed::{
EmbeddingModel, InitOptionsUserDefined, Pooling, TextEmbedding, TextInitOptions,
TokenizerFiles, UserDefinedEmbeddingModel,
};
use std::path::{Path, PathBuf};
use std::str::FromStr;
pub const GRANITE_REPO: &str = "onnx-community/granite-embedding-small-english-r2-ONNX";
const GRANITE_ONNX: &str = "onnx/model_quantized.onnx";
const GRANITE_WEIGHTS: &str = "onnx/model_quantized.onnx_data";
pub const GRANITE_NAME: &str = "granite-embedding-small-english-r2-int8";
const GRANITE_DIM: usize = 384;
pub const THREADS_ENV: &str = "SEMLITH_EMBED_THREADS";
#[derive(Debug, Clone, PartialEq, Default)]
pub enum Model {
Builtin(EmbeddingModel),
#[default]
Granite,
}
impl std::fmt::Display for Model {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Model::Granite => f.write_str(GRANITE_NAME),
Model::Builtin(m) => write!(f, "{m}"),
}
}
}
impl FromStr for Model {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s == GRANITE_NAME {
return Ok(Model::Granite);
}
EmbeddingModel::from_str(s).map(Model::Builtin)
}
}
impl Model {
pub fn dim(&self) -> Result<usize> {
match self {
Model::Granite => Ok(GRANITE_DIM),
Model::Builtin(m) => Ok(TextEmbedding::get_model_info(m)
.map_err(|e| anyhow::anyhow!("{e}"))?
.dim),
}
}
pub fn query_text(&self, query: &str) -> String {
match self {
Model::Builtin(m) => {
let name = m.to_string();
if name.starts_with("BGE") && name.contains("EN") {
format!("Represent this sentence for searching relevant passages: {query}")
} else {
query.to_string()
}
}
Model::Granite => query.to_string(),
}
}
pub fn tokenizer(&self, cache_dir: &Path) -> Option<tokenizers::Tokenizer> {
let name = match self {
Model::Granite => GRANITE_REPO,
Model::Builtin(m) => return builtin_tokenizer(cache_dir, m),
};
let dir = snapshot_dir(cache_dir, name, GRANITE_REVISION)?;
tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).ok()
}
pub fn load(
&self,
cache_dir: PathBuf,
max_length: usize,
quiet: bool,
) -> Result<TextEmbedding> {
if airgap() && !is_cached(&cache_dir) {
bail!(
"{AIRGAP_ENV} is set and no model is cached at {} — \
pre-seed it with SEMLITH_MODEL_CACHE on a connected machine, \
or drop --airgap to let this run download it",
cache_dir.display()
);
}
match self {
Model::Builtin(m) => {
link_runtime()?;
let opts = TextInitOptions::new(m.clone())
.with_show_download_progress(!quiet)
.with_max_length(max_length)
.with_intra_threads(embed_threads())
.with_cache_dir(cache_dir);
TextEmbedding::try_new(opts).map_err(|e| anyhow::anyhow!("{e}"))
}
Model::Granite => load_granite(cache_dir, max_length, quiet),
}
}
}
pub fn link_runtime() -> Result<()> {
#[cfg(feature = "dynamic-ort")]
{
use std::sync::OnceLock;
static ONCE: OnceLock<Result<(), String>> = OnceLock::new();
return ONCE
.get_or_init(|| {
let Ok(exe) = std::env::current_exe() else {
return Err("the running binary could not be located".to_string());
};
let beside = exe
.parent()
.map(|dir| dir.join(RUNTIME_FILE))
.unwrap_or_else(|| PathBuf::from(RUNTIME_FILE));
if !beside.exists() {
return Err(format!(
"{} is not beside the semlith binary. The Linux release \
archive carries it next to `semlith`, and `install.sh` \
puts both in the same directory — a binary copied out of \
the archive on its own cannot embed anything. Unpack the \
archive again, or install with the one-liner in the \
README.",
beside.display()
));
}
ort::init_from(beside.to_string_lossy().as_ref())
.map_err(|e| format!("loading {}: {e}", beside.display()))?
.commit();
Ok(())
})
.clone()
.map_err(anyhow::Error::msg);
}
#[cfg(not(feature = "dynamic-ort"))]
Ok(())
}
#[cfg(feature = "dynamic-ort")]
const RUNTIME_FILE: &str = if cfg!(target_os = "macos") {
"libonnxruntime.dylib"
} else if cfg!(windows) {
"onnxruntime.dll"
} else {
"libonnxruntime.so"
};
fn load_granite(cache_dir: PathBuf, max_length: usize, quiet: bool) -> Result<TextEmbedding> {
link_runtime()?;
check_cache_dir(&cache_dir)?;
let cache = cache_dir.clone();
let checked = snapshot_dir(&cache, GRANITE_REPO, GRANITE_REVISION)
.is_some_and(|dir| already_verified(&dir, GRANITE_REVISION, GRANITE_FILES));
let repo = hf_hub::api::sync::ApiBuilder::new()
.with_cache_dir(cache_dir)
.with_progress(!quiet)
.build()
.context("building the Hugging Face client")?
.repo(hf_hub::Repo::with_revision(
GRANITE_REPO.to_string(),
hf_hub::RepoType::Model,
GRANITE_REVISION.to_string(),
));
let fetch = |name: &str| -> Result<Vec<u8>> {
let path = repo
.get(name)
.with_context(|| format!("fetching {name} from {GRANITE_REPO}"))?;
let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
if checked {
return Ok(bytes);
}
let expected = GRANITE_FILES
.iter()
.find(|(file, _)| *file == name)
.map(|(_, digest)| *digest)
.with_context(|| format!("{name} is not a file semlith pins a digest for"))?;
verify(name, &bytes, expected)?;
Ok(bytes)
};
let tokenizer_files = TokenizerFiles {
tokenizer_file: fetch("tokenizer.json")?,
config_file: fetch("config.json")?,
special_tokens_map_file: fetch("special_tokens_map.json")?,
tokenizer_config_file: fetch("tokenizer_config.json")?,
};
let model = UserDefinedEmbeddingModel::new(fetch(GRANITE_ONNX)?, tokenizer_files)
.with_pooling(Pooling::Cls)
.with_external_initializer(
Path::new(GRANITE_WEIGHTS)
.file_name()
.expect("weights constant has a file name")
.to_string_lossy()
.into_owned(),
fetch(GRANITE_WEIGHTS)?,
);
if !checked && let Some(dir) = snapshot_dir(&cache, GRANITE_REPO, GRANITE_REVISION) {
record_verified(&dir, GRANITE_REVISION, GRANITE_FILES);
}
let opts = InitOptionsUserDefined::new()
.with_max_length(max_length)
.with_intra_threads(embed_threads());
TextEmbedding::try_new_from_user_defined(model, opts)
.map_err(|e| anyhow::anyhow!("loading {GRANITE_NAME}: {e}"))
}
pub const AIRGAP_ENV: &str = "SEMLITH_AIRGAP";
fn builtin_tokenizer(cache_dir: &Path, model: &EmbeddingModel) -> Option<tokenizers::Tokenizer> {
let wanted = format!("{model}").to_lowercase().replace('/', "--");
let entries = std::fs::read_dir(cache_dir).ok()?;
for entry in entries.flatten() {
let name = entry.file_name().to_string_lossy().to_lowercase();
if !name.contains(&wanted) {
continue;
}
for candidate in walk_for(&entry.path(), "tokenizer.json") {
if let Ok(tokenizer) = tokenizers::Tokenizer::from_file(&candidate) {
return Some(tokenizer);
}
}
}
None
}
fn walk_for(root: &Path, name: &str) -> Vec<PathBuf> {
let mut found = Vec::new();
let direct = root.join(name);
if direct.is_file() {
found.push(direct);
}
let Ok(entries) = std::fs::read_dir(root) else {
return found;
};
for entry in entries.flatten() {
if !entry.path().is_dir() {
continue;
}
let nested = entry.path().join(name);
if nested.is_file() {
found.push(nested);
}
}
found
}
pub fn airgap() -> bool {
matches!(
std::env::var(AIRGAP_ENV).ok().as_deref(),
Some("1" | "true" | "yes")
)
}
pub fn refuse_if_airgapped(model: &str) -> Result<()> {
let cache = crate::model_cache_dir()?;
if airgap() && !model_is_cached(&cache, model) {
bail!(
"{AIRGAP_ENV} is set and {model} is not cached at {} — \
pre-seed it with SEMLITH_MODEL_CACHE on a connected machine, \
or drop --airgap to let this run download it",
cache.display()
);
}
Ok(())
}
pub fn model_is_cached(cache_dir: &Path, model: &str) -> bool {
let want = model.replace('/', "--").to_ascii_lowercase();
std::fs::read_dir(cache_dir)
.map(|entries| {
entries.flatten().any(|entry| {
let name = entry.file_name().to_string_lossy().to_ascii_lowercase();
name.ends_with(&want) && entry.path().join("snapshots").exists()
})
})
.unwrap_or(false)
}
pub fn is_cached(cache_dir: &Path) -> bool {
let snapshot = match snapshot_dir(cache_dir, GRANITE_REPO, GRANITE_REVISION) {
Some(dir) => dir,
None => return false,
};
GRANITE_FILES
.iter()
.all(|(name, _)| snapshot.join(name).exists())
}
fn snapshot_dir(cache_dir: &Path, repo: &str, revision: &str) -> Option<PathBuf> {
let dir = cache_dir
.join(format!("models--{}", repo.replace('/', "--")))
.join("snapshots")
.join(revision);
dir.is_dir().then_some(dir)
}
pub fn verify_cached(
cache_dir: &Path,
repo: &str,
revision: &str,
files: &[(&str, &str)],
) -> Result<()> {
check_cache_dir(cache_dir)?;
let _ = revision;
let snapshots = cache_dir
.join(format!("models--{}", repo.replace('/', "--")))
.join("snapshots");
let Ok(entries) = std::fs::read_dir(&snapshots) else {
return Ok(());
};
for entry in entries.flatten() {
let dir = entry.path();
if already_verified(&dir, revision, files) {
continue;
}
let mut read_all = true;
for (name, expected) in files {
let path = dir.join(name);
let Ok(bytes) = std::fs::read(&path) else {
read_all = false;
continue;
};
verify(&format!("{repo}/{name}"), &bytes, expected)?;
}
if read_all {
record_verified(&dir, revision, files);
}
}
Ok(())
}
pub fn embed_threads() -> usize {
if let Ok(raw) = std::env::var(THREADS_ENV)
&& let Ok(n) = raw.parse::<usize>()
&& n > 0
{
return n;
}
performance_cores().unwrap_or_else(|| {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
})
}
#[cfg(target_os = "macos")]
fn performance_cores() -> Option<usize> {
let mut out: i32 = 0;
let mut len = std::mem::size_of::<i32>();
let name = c"hw.perflevel0.logicalcpu";
let rc = unsafe {
libc::sysctlbyname(
name.as_ptr(),
&mut out as *mut i32 as *mut libc::c_void,
&mut len,
std::ptr::null_mut(),
0,
)
};
(rc == 0 && out > 0).then_some(out as usize)
}
#[cfg(not(target_os = "macos"))]
fn performance_cores() -> Option<usize> {
None
}
pub const GRANITE_REVISION: &str = "1dc7835ba0cb9c76a3618d0bf0c427c97671b3c8";
pub const GRANITE_FILES: &[(&str, &str)] = &[
(
"tokenizer.json",
"feeb83348dcb033bc6b9d2e1f7906ca9eb2d122845000c9416d894d7c2927149",
),
(
"config.json",
"1a1710c20911da8c96179716bf44058e54cca6fa7952cce77be83ef05edae3ee",
),
(
"special_tokens_map.json",
"ea97ecdbcc73713039d8d64dbb05e3689495c96657fbd9a18f5bed381be81049",
),
(
"tokenizer_config.json",
"ce06781b38bb393db68c9e0709bddd31ef5d88f2c6fbb3fd9f369778fb85e451",
),
(
GRANITE_ONNX,
"a3fad524afc3f060216a8ddbb1ac89c9b6498fba8995b5718bde879076a2e9ba",
),
(
GRANITE_WEIGHTS,
"1f4cf47e4adec7f7ae09db03d071ba8667e07f9a4203142c7efa8d37fe453597",
),
];
#[derive(serde::Serialize, serde::Deserialize)]
struct Verified {
revision: String,
files: std::collections::BTreeMap<String, (u64, u128, String)>,
}
const STAMP: &str = ".semlith-verified";
fn fingerprint(path: &Path) -> Option<(u64, u128)> {
let meta = std::fs::metadata(path).ok()?;
let modified = meta
.modified()
.ok()?
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_nanos();
Some((meta.len(), modified))
}
fn already_verified(dir: &Path, revision: &str, files: &[(&str, &str)]) -> bool {
let Ok(text) = std::fs::read_to_string(dir.join(STAMP)) else {
return false;
};
let Ok(stamp) = serde_json::from_str::<Verified>(&text) else {
return false;
};
if stamp.revision != revision {
return false;
}
files.iter().all(|(name, expected)| {
stamp
.files
.get(*name)
.zip(fingerprint(&dir.join(name)))
.is_some_and(|((size, mtime, digest), (now_size, now_mtime))| {
digest == expected && *size == now_size && *mtime == now_mtime
})
})
}
fn record_verified(dir: &Path, revision: &str, files: &[(&str, &str)]) {
let mut stamp = Verified {
revision: revision.to_string(),
files: std::collections::BTreeMap::new(),
};
for (name, expected) in files {
let Some((size, mtime)) = fingerprint(&dir.join(name)) else {
return;
};
stamp
.files
.insert((*name).to_string(), (size, mtime, (*expected).to_string()));
}
let Ok(body) = serde_json::to_vec(&stamp) else {
return;
};
let _ = crate::home::write_private(&dir.join(STAMP), &body);
}
pub fn digest(bytes: &[u8]) -> String {
use sha2::{Digest, Sha256};
format!("{:x}", Sha256::digest(bytes))
}
pub fn verify(name: &str, bytes: &[u8], expected: &str) -> Result<()> {
let actual = digest(bytes);
if actual != expected {
bail!(
"{name} does not match the digest semlith pins for it.\n \
expected {expected}\n got {actual}\n\
The model cache holds a file that is not the one this release was \
built against. Delete the cache and let semlith fetch it again; if \
it happens twice, the file upstream has changed and semlith needs a \
release rather than a retry."
);
}
Ok(())
}
pub fn check_cache_dir(cache: &Path) -> Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
let Ok(meta) = std::fs::metadata(cache) else {
return Ok(());
};
let me = unsafe { libc::getuid() };
if meta.uid() != me {
bail!(
"{} is owned by uid {} rather than by you. semlith will not load \
model weights from a directory somebody else owns — run `chown -R \
{me} {}`, or point {} somewhere you own.",
cache.display(),
meta.uid(),
cache.display(),
crate::MODEL_CACHE_ENV,
);
}
let mode = meta.permissions().mode() & 0o777;
if mode & 0o022 != 0 {
bail!(
"{} is mode {mode:o}, so other users on this machine can write to \
it. semlith will not load model weights from it — run `chmod 700 \
{}`.",
cache.display(),
cache.display(),
);
}
}
#[cfg(not(unix))]
let _ = cache;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn granite_round_trips_through_its_stored_name() {
let parsed: Model = GRANITE_NAME.parse().unwrap();
assert_eq!(parsed, Model::Granite);
assert_eq!(parsed.to_string(), GRANITE_NAME);
assert_eq!(parsed.dim().unwrap(), 384);
}
#[test]
fn stores_written_by_0_1_0_still_parse() {
let parsed: Model = "BGESmallENV15".parse().unwrap();
assert_eq!(parsed, Model::Builtin(EmbeddingModel::BGESmallENV15));
assert_eq!(parsed.to_string(), "BGESmallENV15");
assert_eq!(parsed.dim().unwrap(), 384);
}
#[test]
fn nonsense_model_names_are_rejected() {
assert!("not-a-model".parse::<Model>().is_err());
}
#[test]
fn only_bge_english_gets_the_instruction_prefix() {
let bge = Model::Builtin(EmbeddingModel::BGESmallENV15);
assert!(bge.query_text("hi").starts_with("Represent this sentence"));
assert!(bge.query_text("hi").ends_with("hi"));
assert_eq!(Model::Granite.query_text("hi"), "hi");
}
#[test]
fn thread_count_honours_the_override() {
unsafe { std::env::set_var(THREADS_ENV, "3") };
assert_eq!(embed_threads(), 3);
unsafe { std::env::remove_var(THREADS_ENV) };
assert!(embed_threads() >= 1);
}
}