goosedump 0.12.49

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

//! Local text and embedding models.
//!
//! GGUF weights download into the standard Hugging Face Hub cache on first use. See
//! `docs/models/` for role, inference, and benchmark notes.

mod bge;
mod gguf;
mod integrity;
mod kernels;
mod profile;
mod qwen3;
mod tokenizer;
mod types;
mod wordpiece;

use std::cmp::Ordering;
use std::env;
#[cfg(all(feature = "bench", target_os = "linux"))]
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};
use std::time::Duration;

use anyhow::{Context as _, bail};
use rayon::{ThreadPool, ThreadPoolBuilder};

pub(crate) use self::bge::EMBEDDING_DIMENSIONS;
use self::bge::EmbeddingModel;
use self::integrity::verify_cached_sha256;
use self::qwen3::{Generation, GenerationStopReason, PhaseProfile, TextModel};
#[cfg(feature = "bench")]
use crate::alloc_profile::AllocationProfile;

/// Hugging Face source of the fixed text model.
const TEXTGEN_WEIGHTS_REPO: &str = "prism-ml/Bonsai-4B-gguf";
const TEXTGEN_WEIGHTS_FILE: &str = "Bonsai-4B-Q1_0.gguf";
const MAX_DECODE_DURATION: Duration = Duration::from_mins(5);

/// Pinned Hugging Face source of the sentence embedding model.
pub(crate) const EMBEDDING_WEIGHTS_REPO: &str = "ggml-org/bge-small-en-v1.5-Q8_0-GGUF";
pub(crate) const EMBEDDING_WEIGHTS_FILE: &str = "bge-small-en-v1.5-q8_0.gguf";
pub(crate) const EMBEDDING_WEIGHTS_SHA256: &str =
    "f046db1dc724cf4f6f0a0c5917e922823b73eb1d27b8f9a9c2797f7866974804";
pub(crate) const EMBEDDING_MODEL_ID: &str = "bge-small-en-v1.5-q8_0-pipeline-v1@f046db1dc724cf4f6f0a0c5917e922823b73eb1d27b8f9a9c2797f7866974804";
const EMBEDDING_INTEGRITY_STAMP_FILE: &str = "bge-small-en-v1.5-q8_0.gguf.json";

const TEXTGEN_CONTEXT: u16 = 4_096;

const INFERENCE_THREADS_ENV: &str = "GOOSEDUMP_INFERENCE_THREADS";
const INFERENCE_PROFILE_ENV: &str = "GOOSEDUMP_PROFILE_INFERENCE";
const EMBEDDING_SINGLE_THREADS: usize = 2;

/// Fixed dense Qwen3 Bonsai-4B `Q1_0` model for claim extraction and compact filtering.
pub struct TextGen {
    model: TextModel,
    pool: ThreadPool,
}

#[cfg(feature = "bench")]
#[derive(serde::Serialize)]
pub(crate) struct BenchmarkGeneration {
    pub(crate) text: String,
    pub(crate) prompt_tokens: usize,
    pub(crate) generated_tokens: usize,
    pub(crate) decode_tokens: usize,
    pub(crate) token_ids: Vec<u32>,
    pub(crate) stop_reason: &'static str,
    pub(crate) prefill_allocations: AllocationProfile,
    pub(crate) decode_allocations: AllocationProfile,
    pub(crate) prefill_allocation_calls_per_token: f64,
    pub(crate) prefill_requested_bytes_per_token: f64,
    pub(crate) decode_allocation_calls_per_token: Option<f64>,
    pub(crate) decode_requested_bytes_per_token: Option<f64>,
}

/// Fixed BGE-small-en-v1.5 model for semantic claim retrieval.
pub struct Embedder {
    model: EmbeddingModel,
    single_pool: OnceLock<ThreadPool>,
    batch_pool: OnceLock<ThreadPool>,
    pool_init: Mutex<()>,
    max_threads: usize,
}

/// Validated, normalized embedding produced by [`Embedder`].
#[derive(Clone, Debug)]
pub struct Embedding(Vec<f32>);

impl TryFrom<Vec<f32>> for Embedding {
    type Error = anyhow::Error;

    fn try_from(values: Vec<f32>) -> Result<Self, Self::Error> {
        if values.len() != EMBEDDING_DIMENSIONS || !values.iter().all(|value| value.is_finite()) {
            bail!("embedding has an invalid shape or value");
        }
        let norm_squared = values.iter().map(|value| value * value).sum::<f32>();
        if (norm_squared - 1.0).abs() > 1e-3 {
            bail!("embedding is not normalized");
        }
        Ok(Self(values))
    }
}

impl From<&Embedding> for Vec<u8> {
    fn from(value: &Embedding) -> Self {
        value
            .0
            .iter()
            .flat_map(|component| component.to_ne_bytes())
            .collect()
    }
}

/// Finite cosine similarity between two normalized embeddings.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Similarity(f32);

impl Similarity {
    fn from_valid(value: f32) -> Self {
        Self(if value == 0.0 { 0.0 } else { value })
    }
}

impl TryFrom<f32> for Similarity {
    type Error = anyhow::Error;

    fn try_from(value: f32) -> Result<Self, Self::Error> {
        if !value.is_finite() || !(-1.0..=1.0).contains(&value) {
            bail!("embedding similarity is outside the cosine range");
        }
        Ok(Self::from_valid(value))
    }
}

impl Eq for Similarity {}

impl PartialOrd for Similarity {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Similarity {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.total_cmp(&other.0)
    }
}

impl Embedding {
    pub fn similarity(&self, other: &Self) -> Similarity {
        let value = self
            .0
            .iter()
            .zip(&other.0)
            .map(|(left, right)| left * right)
            .sum::<f32>();
        Similarity::from_valid(value.clamp(-1.0, 1.0))
    }
}

#[cfg(feature = "bench")]
#[expect(
    clippy::cast_precision_loss,
    reason = "bench rates only need f64 display precision"
)]
fn allocation_rate(value: u64, tokens: usize) -> Option<f64> {
    if tokens == 0 {
        return None;
    }
    Some(value as f64 / tokens as f64)
}

impl TextGen {
    /// Download and load the text model from the Hugging Face cache.
    ///
    /// # Errors
    /// Returns an error if downloading, loading, or configuring inference fails.
    pub fn load() -> anyhow::Result<Self> {
        let path = model_file(TEXTGEN_WEIGHTS_REPO, TEXTGEN_WEIGHTS_FILE)?;
        Self::load_path(&path)
    }

    #[cfg(feature = "bench")]
    pub(crate) fn load_for_benchmark(cold: bool) -> anyhow::Result<Self> {
        let path = model_file(TEXTGEN_WEIGHTS_REPO, TEXTGEN_WEIGHTS_FILE)?;
        if cold {
            evict_text_model_pages(&path)?;
        }
        Self::load_path(&path)
    }

    fn load_path(path: &Path) -> anyhow::Result<Self> {
        let model = TextModel::load(path)?;
        let threads = inference_threads()?;
        let thread_count = usize::try_from(threads).context("text inference thread count")?;
        let pool = ThreadPoolBuilder::new()
            .num_threads(thread_count)
            .build()
            .context("create text inference thread pool")?;
        Ok(Self { model, pool })
    }

    /// Return the formatted prompt's token count.
    ///
    /// # Errors
    /// Returns an error if tokenization fails.
    pub fn completion_prompt_tokens(&self, system: &str, user: &str) -> anyhow::Result<usize> {
        self.model.prompt_tokens(&completion_prompt(system, user))
    }

    /// Return whether the formatted prompt and requested generation fit the context.
    ///
    /// # Errors
    /// Returns an error if tokenization fails.
    pub fn completion_fits(
        &self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<bool> {
        let prompt_tokens = self.completion_prompt_tokens(system, user)?;
        Ok(prompt_tokens
            .checked_add(max_tokens)
            .is_some_and(|tokens| tokens <= usize::from(TEXTGEN_CONTEXT)))
    }

    /// Greedy instruction completion of `system` + `user`, capped at
    /// `max_tokens`.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails, or generation is stopped by a safety guard.
    pub fn complete(
        &mut self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<String> {
        let generation = self.generate(system, user, max_tokens, Some(MAX_DECODE_DURATION))?;
        print_generation_profile(&generation)?;
        finish_generation(generation)
    }

    /// Greedy instruction completion for background work without a wall-clock decode limit.
    /// Token and repetition guards still bound generation.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails, or generation is stopped by the
    /// repetition guard.
    pub fn complete_background(
        &mut self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<String> {
        let generation = self.generate(system, user, max_tokens, None)?;
        print_generation_profile(&generation)?;
        finish_generation(generation)
    }

    #[cfg(feature = "bench")]
    pub(crate) fn benchmark_complete(
        &self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<BenchmarkGeneration> {
        let generation = self.generate(system, user, max_tokens, Some(MAX_DECODE_DURATION))?;
        print_generation_profile(&generation)?;
        let prefill_allocation_calls_per_token = allocation_rate(
            generation.prefill_allocations.allocation_calls,
            generation.prompt_tokens,
        )
        .context("prefill has no tokens")?;
        let prefill_requested_bytes_per_token = allocation_rate(
            generation.prefill_allocations.requested_bytes,
            generation.prompt_tokens,
        )
        .context("prefill has no tokens")?;
        let decode_allocation_calls_per_token = allocation_rate(
            generation.decode_allocations.allocation_calls,
            generation.decode_tokens,
        );
        let decode_requested_bytes_per_token = allocation_rate(
            generation.decode_allocations.requested_bytes,
            generation.decode_tokens,
        );
        Ok(BenchmarkGeneration {
            text: generation.text,
            prompt_tokens: generation.prompt_tokens,
            generated_tokens: generation.generated_tokens,
            decode_tokens: generation.decode_tokens,
            token_ids: generation.generated_token_ids,
            stop_reason: stop_reason_name(generation.stop_reason),
            prefill_allocations: generation.prefill_allocations,
            decode_allocations: generation.decode_allocations,
            prefill_allocation_calls_per_token,
            prefill_requested_bytes_per_token,
            decode_allocation_calls_per_token,
            decode_requested_bytes_per_token,
        })
    }

    fn generate(
        &self,
        system: &str,
        user: &str,
        max_tokens: usize,
        max_decode_duration: Option<Duration>,
    ) -> anyhow::Result<Generation> {
        let prompt = completion_prompt(system, user);
        let profile_enabled = inference_profile_enabled();
        self.pool.install(|| {
            self.model.generate(
                &prompt,
                max_tokens,
                TEXTGEN_CONTEXT,
                profile_enabled,
                max_decode_duration,
            )
        })
    }
}

fn finish_generation(generation: Generation) -> anyhow::Result<String> {
    match generation.stop_reason {
        GenerationStopReason::EndOfSequence | GenerationStopReason::TokenLimit => {
            Ok(generation.text)
        }
        GenerationStopReason::TokenCycle => bail!(
            "text generation stopped after detecting a repeating token cycle ({} tokens)",
            generation.generated_tokens
        ),
        GenerationStopReason::TimeBudget => bail!(
            "text generation exceeded its decode time budget ({} tokens)",
            generation.generated_tokens
        ),
    }
}

fn inference_profile_enabled() -> bool {
    env::var_os(INFERENCE_PROFILE_ENV).is_some()
        || env::var_os("GOOSEDUMP_PROFILE_COMPACT").is_some()
}

fn print_generation_profile(generation: &Generation) -> anyhow::Result<()> {
    let Some(profile) = generation.profile.as_ref() else {
        return Ok(());
    };
    print_phase_profile(
        "prefill",
        generation.prompt_tokens,
        generation.prefill_duration,
        &profile.prefill,
    )?;
    print_phase_profile(
        "decode",
        generation.decode_tokens,
        generation.decode_duration,
        &profile.decode,
    )
}

fn print_phase_profile(
    name: &str,
    tokens: usize,
    duration: std::time::Duration,
    profile: &PhaseProfile,
) -> anyhow::Result<()> {
    let seconds = duration.as_secs_f64();
    let duration_ms = seconds * 1_000.0;
    let rate = token_rate(tokens, seconds)?;
    if let Some(process) = profile.process {
        eprintln!(
            "goosedump: profile: inference: {name}: tokens={tokens} duration_ms={duration_ms:.1} tokens_per_second={rate:.1} minor_faults={} major_faults={} read_bytes={} rss_kib_start={} rss_kib_end={} peak_rss_kib={}",
            process.minor_faults,
            process.major_faults,
            process.read_bytes,
            process.rss_kib_start,
            process.rss_kib_end,
            process.peak_rss_kib,
        );
    } else {
        eprintln!(
            "goosedump: profile: inference: {name}: tokens={tokens} duration_ms={duration_ms:.1} tokens_per_second={rate:.1} process_metrics=unavailable"
        );
    }

    Ok(())
}

#[cfg(feature = "bench")]
const fn stop_reason_name(reason: GenerationStopReason) -> &'static str {
    match reason {
        GenerationStopReason::EndOfSequence => "end_of_sequence",
        GenerationStopReason::TokenLimit => "token_limit",
        GenerationStopReason::TokenCycle => "token_cycle",
        GenerationStopReason::TimeBudget => "time_budget",
    }
}

#[cfg(all(feature = "bench", target_os = "linux"))]
fn evict_text_model_pages(path: &Path) -> anyhow::Result<()> {
    use std::os::fd::AsRawFd as _;
    let file =
        File::open(path).with_context(|| format!("open {} for cache eviction", path.display()))?;
    // SAFETY: posix_fadvise does not access Rust memory, and the descriptor remains
    // open for the duration of the call.
    let result = unsafe { libc::posix_fadvise(file.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED) };
    if result != 0 {
        return Err(std::io::Error::from_raw_os_error(result))
            .with_context(|| format!("evict cached pages for {}", path.display()));
    }
    Ok(())
}

#[cfg(all(feature = "bench", not(target_os = "linux")))]
fn evict_text_model_pages(_path: &Path) -> anyhow::Result<()> {
    bail!("cold inference benchmarking is supported only on Linux")
}

impl Embedder {
    /// Download and load the embedding model from the Hugging Face cache.
    pub fn load() -> anyhow::Result<Self> {
        let path = ensure_embedding_model()?;
        Self::load_path(&path)
    }

    /// Produce one normalized sentence embedding.
    pub fn embed(&self, text: &str) -> anyhow::Result<Embedding> {
        self.pool(false)?
            .install(|| self.model.embed(text))
            .and_then(Embedding::try_from)
    }

    /// Produce normalized embeddings for an independent batch of sentences.
    pub fn embed_batch(&self, texts: &[&str]) -> anyhow::Result<Vec<Embedding>> {
        if texts.len() <= 1 {
            return texts.first().map_or_else(
                || Ok(Vec::new()),
                |text| self.embed(text).map(|value| vec![value]),
            );
        }
        self.pool(true)?
            .install(|| self.model.embed_batch(texts))?
            .into_iter()
            .map(Embedding::try_from)
            .collect()
    }

    /// Score each candidate by its greatest similarity to any reference.
    pub fn relevance(
        &self,
        references: &[&str],
        candidates: &[&str],
    ) -> anyhow::Result<Vec<Similarity>> {
        if references.is_empty() {
            bail!("relevance has no references");
        }
        let inputs = references
            .iter()
            .chain(candidates)
            .copied()
            .collect::<Vec<_>>();
        let embeddings = self.embed_batch(&inputs)?;
        if embeddings.len() != inputs.len() {
            bail!("relevance embedding count differs");
        }
        let reference_embeddings = embeddings
            .get(..references.len())
            .context("relevance reference range is invalid")?;
        let candidate_embeddings = embeddings
            .get(references.len()..)
            .context("relevance candidate range is invalid")?;
        candidate_embeddings
            .iter()
            .map(|candidate| {
                reference_embeddings
                    .iter()
                    .map(|reference| candidate.similarity(reference))
                    .max()
                    .context("relevance has no reference embedding")
            })
            .collect()
    }

    fn pool(&self, batch: bool) -> anyhow::Result<&ThreadPool> {
        let single_threads = self.max_threads.min(EMBEDDING_SINGLE_THREADS);
        let (pool, threads) = if !batch || self.max_threads == single_threads {
            (&self.single_pool, single_threads)
        } else {
            (&self.batch_pool, self.max_threads)
        };
        if let Some(pool) = pool.get() {
            return Ok(pool);
        }
        let _guard = self
            .pool_init
            .lock()
            .map_err(|_| anyhow::anyhow!("embedding inference pool initialization is poisoned"))?;
        if let Some(pool) = pool.get() {
            return Ok(pool);
        }
        let candidate = ThreadPoolBuilder::new()
            .num_threads(threads)
            .build()
            .context("create embedding inference thread pool")?;
        pool.set(candidate)
            .map_err(|_| anyhow::anyhow!("embedding inference pool initialized twice"))?;
        pool.get()
            .context("initialize embedding inference thread pool")
    }

    fn load_path(path: &Path) -> anyhow::Result<Self> {
        let model = EmbeddingModel::load(path)?;
        let threads = inference_threads()?;
        let max_threads = usize::try_from(threads).context("embedding inference thread count")?;
        Ok(Self {
            model,
            single_pool: OnceLock::new(),
            batch_pool: OnceLock::new(),
            pool_init: Mutex::new(()),
            max_threads,
        })
    }
}

fn completion_prompt(system: &str, user: &str) -> String {
    format!(
        "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
    )
}

fn token_rate(tokens: usize, seconds: f64) -> anyhow::Result<f64> {
    if seconds == 0.0 {
        Ok(0.0)
    } else {
        let tokens = u32::try_from(tokens).context("token count exceeds u32")?;
        Ok(f64::from(tokens) / seconds)
    }
}

pub(crate) fn inference_threads() -> anyhow::Result<i32> {
    let available = std::thread::available_parallelism()
        .context("detect available parallelism")?
        .get();
    configured_threads(INFERENCE_THREADS_ENV, available)
}

fn configured_threads(name: &str, default: usize) -> anyhow::Result<i32> {
    let value = match env::var_os(name) {
        Some(value) => {
            let value = value
                .into_string()
                .map_err(|_| anyhow::anyhow!("{name} is not valid UTF-8"))?;
            let threads = value
                .parse::<usize>()
                .with_context(|| format!("parse {name}"))?;
            if threads == 0 {
                bail!("{name} must be greater than zero");
            }
            threads
        }
        None => default,
    };
    i32::try_from(value).with_context(|| format!("{name} exceeds i32"))
}

/// Fetch `filename` from a model repository into the standard Hugging Face cache.
fn model_file(repo_id: &str, filename: &str) -> anyhow::Result<PathBuf> {
    let api = hf_hub::api::sync::ApiBuilder::from_env()
        .with_cache_dir(hugging_face_cache_dir()?)
        .build()
        .context("configure Hugging Face Hub client")?;
    api.model(repo_id.to_owned())
        .get(filename)
        .with_context(|| format!("download {repo_id}/{filename}"))
}

/// Fetch and verify the pinned embedding model.
fn ensure_embedding_model() -> anyhow::Result<PathBuf> {
    let path = model_file(EMBEDDING_WEIGHTS_REPO, EMBEDDING_WEIGHTS_FILE)?;
    let stamp_path = embedding_integrity_stamp_path();
    verify_cached_sha256(&path, stamp_path.as_deref(), EMBEDDING_WEIGHTS_SHA256)?;
    Ok(path)
}

fn embedding_integrity_stamp_path() -> Option<PathBuf> {
    let state_dir = env::var_os("GOOSEDUMP_STATE_DIR")
        .filter(|path| !path.is_empty())
        .map(PathBuf::from)
        .or_else(|| {
            dirs::state_dir()
                .or_else(dirs::data_local_dir)
                .map(|path| path.join("goosedump"))
        })?;
    Some(
        state_dir
            .join("model-integrity")
            .join(EMBEDDING_INTEGRITY_STAMP_FILE),
    )
}

/// Standard Hugging Face Hub cache (honors `HF_HUB_CACHE`, `HF_HOME`, and XDG).
fn hugging_face_cache_dir() -> anyhow::Result<PathBuf> {
    match (env::var_os("HF_HUB_CACHE"), env::var_os("HF_HOME")) {
        (Some(path), _) => Ok(PathBuf::from(path)),
        (None, Some(path)) => Ok(PathBuf::from(path).join("hub")),
        (None, None) => Ok(dirs::cache_dir()
            .context("locate platform cache directory")?
            .join("huggingface")
            .join("hub")),
    }
}