goosedump 0.8.2

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

//! Optional candle model layer. Every entry point degrades to `None`/empty when
//! weights are not present in the local cache, so `compact` and the rest of the
//! tool behave identically offline — the model only ever sharpens results, it
//! never gates them.

use std::path::PathBuf;

use anyhow::Context as _;
use candle_core::{Device, Tensor};
use candle_nn::{Module, VarBuilder};
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::jina_bert::{BertModel, Config, DTYPE};
use candle_transformers::models::quantized_mixformer::{
    Config as PhiConfig, MixFormerSequentialForCausalLM,
};
use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer};

/// Cache subdirectory and Hugging Face repo of the sentence-embedding model.
pub(crate) const EMBEDDER_NAME: &str = "jina-embeddings-v2-base-en";
pub(crate) const EMBEDDER_REPO: &str = "jinaai/jina-embeddings-v2-base-en";
/// The files `pull` fetches and `load` expects, in the flat cache layout.
pub(crate) const EMBEDDER_FILES: [&str; 3] = ["config.json", "tokenizer.json", "model.safetensors"];

/// Cache subdirectory of the Phi-1.5 judge model and its Hugging Face sources.
/// The GGUF weights and tokenizer both come from the candle helper repo.
pub(crate) const TEXTGEN_NAME: &str = "phi-1_5";
pub(crate) const TEXTGEN_WEIGHTS_REPO: &str = "lmz/candle-quantized-phi";
pub(crate) const TEXTGEN_WEIGHTS_FILE: &str = "model-q4k.gguf";
pub(crate) const TEXTGEN_TOKENIZER_REPO: &str = "lmz/candle-quantized-phi";
pub(crate) const TEXTGEN_TOKENIZER_FILE: &str = "tokenizer.json";
/// Phi-1.5 terminates with the standard `GPT-2` end-of-text token.
const TEXTGEN_EOS: &str = "<|endoftext|>";

/// Sentence-embedding model (`jina-embeddings-v2-base-en`, 768-dim, `ALiBi`)
/// and salience selection. Built from the local model cache; absent weights
/// yield `None` and callers fall back to deterministic behavior.
pub struct Embedder {
    model: BertModel,
    tokenizer: Tokenizer,
    device: Device,
}

impl Embedder {
    /// Load the embedder from the model cache, or `None` when its weights are
    /// not cached or fail to load. No network access is attempted here.
    #[must_use]
    pub fn load() -> Option<Self> {
        Self::from_cache().ok()
    }

    fn from_cache() -> anyhow::Result<Self> {
        let dir = model_cache_dir(EMBEDDER_NAME);
        let config_path = dir.join("config.json");
        let tokenizer_path = dir.join("tokenizer.json");
        let weights_path = dir.join("model.safetensors");
        anyhow::ensure!(
            config_path.exists() && tokenizer_path.exists() && weights_path.exists(),
            "embedder weights not cached in {}",
            dir.display()
        );

        let device = Device::Cpu;
        let config: Config = serde_json::from_str(&std::fs::read_to_string(config_path)?)?;
        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)?;
        // SAFETY: memory-mapping the weights is sound as long as the file is not
        // mutated while mapped; it lives in our read-only model cache.
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DTYPE, &device)? };
        let model = BertModel::new(vb, &config)?;
        Ok(Self {
            model,
            tokenizer,
            device,
        })
    }

    /// Mean-pooled, L2-normalized sentence embeddings — one row per input text,
    /// matching the jina-embeddings reference pooling.
    ///
    /// # Errors
    /// Returns an error if tokenization or the forward pass fails.
    pub fn embed(&self, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let mut tokenizer = self.tokenizer.clone();
        tokenizer.with_padding(Some(PaddingParams {
            strategy: PaddingStrategy::BatchLongest,
            ..PaddingParams::default()
        }));
        let encodings = tokenizer
            .encode_batch(texts.to_vec(), true)
            .map_err(anyhow::Error::msg)?;

        let mut ids = Vec::with_capacity(encodings.len());
        let mut masks = Vec::with_capacity(encodings.len());
        for encoding in &encodings {
            ids.push(Tensor::new(encoding.get_ids(), &self.device)?);
            masks.push(Tensor::new(encoding.get_attention_mask(), &self.device)?);
        }
        let token_ids = Tensor::stack(&ids, 0)?;
        let attention_mask = Tensor::stack(&masks, 0)?;

        // JinaBERT's encoder has no attention-mask input; padding tokens are
        // masked out here during mean-pooling so they do not skew the vector.
        let hidden = self.model.forward(&token_ids)?;

        // Mean-pool over the sequence with the attention mask, then L2-normalize.
        let mask = attention_mask.to_dtype(DTYPE)?.unsqueeze(2)?;
        let summed = hidden.broadcast_mul(&mask)?.sum(1)?;
        let counts = mask.sum(1)?;
        let pooled = summed.broadcast_div(&counts)?;
        let norms = pooled.sqr()?.sum_keepdim(1)?.sqrt()?;
        let normalized = pooled.broadcast_div(&norms)?;
        let rows = normalized.to_vec2::<f32>()?;
        Ok(rows)
    }
}

/// Small instruct model (`Qwen2.5-0.5B-Instruct`, 4-bit GGUF) used by `compact`
/// Small quantized `phi-1_5` model used by `compact` to judge which
/// deterministic goal/preference candidates are genuine, still-current
/// directives. Decoding is greedy, so its output is reproducible for a given
/// weights file. Absent weights yield `None` and callers keep the unjudged
/// candidate list.
pub struct TextGen {
    model: MixFormerSequentialForCausalLM,
    tokenizer: Tokenizer,
    device: Device,
    eos: u32,
}

impl TextGen {
    /// Load the text model from the model cache, or `None` when its weights are
    /// not cached or fail to load. No network access is attempted here.
    #[must_use]
    pub fn load() -> Option<Self> {
        Self::from_cache().ok()
    }

    fn from_cache() -> anyhow::Result<Self> {
        let dir = model_cache_dir(TEXTGEN_NAME);
        let weights_path = dir.join(TEXTGEN_WEIGHTS_FILE);
        let tokenizer_path = dir.join(TEXTGEN_TOKENIZER_FILE);
        anyhow::ensure!(
            weights_path.exists() && tokenizer_path.exists(),
            "text model weights not cached in {}",
            dir.display()
        );

        let device = Device::Cpu;
        let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(
            &weights_path,
            &device,
        )?;
        let model = MixFormerSequentialForCausalLM::new(&PhiConfig::v1_5(), vb)?;
        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)?;
        let eos = tokenizer
            .token_to_id(TEXTGEN_EOS)
            .with_context(|| format!("tokenizer has no {TEXTGEN_EOS} token"))?;
        Ok(Self {
            model,
            tokenizer,
            device,
            eos,
        })
    }

    /// Greedy instruction completion of `system` + `user`, capped at
    /// `max_tokens`.
    ///
    /// # Errors
    /// Returns an error if tokenization or a forward pass fails.
    pub fn complete(
        &mut self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<String> {
        let prompt = format!("Instruction:\n{system}\n\nInput:\n{user}\n\nAnswer:\n");
        let encoding = self
            .tokenizer
            .encode(prompt, true)
            .map_err(anyhow::Error::msg)?;

        self.model.clear_kv_cache();
        let mut sampler = LogitsProcessor::new(0, None, None);
        let mut generated: Vec<u32> = Vec::new();
        let mut input: Vec<u32> = encoding.get_ids().to_vec();
        for _ in 0..max_tokens {
            let tensor = Tensor::new(input.as_slice(), &self.device)?.unsqueeze(0)?;
            let logits = self.model.forward(&tensor)?;
            let next = sampler.sample(&logits)?;
            if next == self.eos {
                break;
            }
            generated.push(next);
            input = vec![next];
        }
        self.tokenizer
            .decode(&generated, true)
            .map_err(anyhow::Error::msg)
    }
}

/// Fetch `files` from the Hugging Face `repo_id` into `dest` (flat layout).
pub(crate) fn pull_files(
    repo_id: &str,
    files: &[&str],
    dest: &std::path::Path,
) -> anyhow::Result<()> {
    std::fs::create_dir_all(dest).with_context(|| format!("create {}", dest.display()))?;
    let api = hf_hub::api::sync::Api::new()?;
    let repo = api.model(repo_id.to_string());
    for file in files {
        let cached = repo
            .get(file)
            .with_context(|| format!("download {repo_id}/{file}"))?;
        let target = dest.join(file);
        std::fs::copy(&cached, &target).with_context(|| format!("write {}", target.display()))?;
    }
    Ok(())
}

/// The per-model cache directory (honors `XDG_CACHE_HOME`).
pub(crate) fn model_cache_dir(name: &str) -> PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("goosedump")
        .join("models")
        .join(name)
}