goosedump 0.9.1

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

//! Candle model layer. Models are downloaded to the local cache on first use,
//! then loaded locally for inference.

use std::path::PathBuf;

use anyhow::Context as _;
use candle_core::quantized::gguf_file;
use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::bert::{BertModel, Config, DTYPE};
use candle_transformers::models::quantized_qwen2::ModelWeights as Qwen2;
use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer};

/// Cache subdirectory and Hugging Face repo of the sentence-embedding model.
pub(crate) const EMBEDDER_NAME: &str = "all-MiniLM-L6-v2";
pub(crate) const EMBEDDER_REPO: &str = "sentence-transformers/all-MiniLM-L6-v2";
/// 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 Qwen2 judge model and its Hugging Face sources.
/// The GGUF weights come from the official Qwen2 instruct GGUF repo; the
/// tokenizer ships in the base instruct repo.
pub(crate) const TEXTGEN_NAME: &str = "qwen2-0_5b-instruct";
pub(crate) const TEXTGEN_WEIGHTS_REPO: &str = "Qwen/Qwen2-0.5B-Instruct-GGUF";
pub(crate) const TEXTGEN_WEIGHTS_FILE: &str = "qwen2-0_5b-instruct-q4_0.gguf";
pub(crate) const TEXTGEN_TOKENIZER_REPO: &str = "Qwen/Qwen2-0.5B-Instruct";
pub(crate) const TEXTGEN_TOKENIZER_FILE: &str = "tokenizer.json";
/// Qwen2 instruct turns end with the `<|im_end|>` marker.
const TEXTGEN_EOS: &str = "<|im_end|>";

/// Sentence-embedding model (`all-MiniLM-L6-v2`, 384-dim) used for semantic
/// deduplication during compaction.
pub struct Embedder {
    model: BertModel,
    tokenizer: Tokenizer,
    device: Device,
}

impl Embedder {
    /// Download and load the embedder from the local model cache.
    ///
    /// # Errors
    /// Returns an error if downloading or loading the model fails.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(EMBEDDER_NAME);
        ensure_files(EMBEDDER_REPO, &EMBEDDER_FILES, &dir)?;
        Self::from_cache(&dir)
    }

    fn from_cache(dir: &std::path::Path) -> anyhow::Result<Self> {
        let config_path = dir.join("config.json");
        let tokenizer_path = dir.join("tokenizer.json");
        let weights_path = dir.join("model.safetensors");

        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::load(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 token_type_ids = token_ids.zeros_like()?;
        let attention_mask = Tensor::stack(&masks, 0)?;

        let hidden = self
            .model
            .forward(&token_ids, &token_type_ids, Some(&attention_mask))?;

        // 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)
    }
}

/// Quantized `Qwen2-0.5B-Instruct` 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.
pub struct TextGen {
    model: Qwen2,
    tokenizer: Tokenizer,
    device: Device,
    eos: u32,
}

impl TextGen {
    /// Download and load the text model from the local model cache.
    ///
    /// # Errors
    /// Returns an error if downloading or loading the model fails.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(TEXTGEN_NAME);
        ensure_files(TEXTGEN_WEIGHTS_REPO, &[TEXTGEN_WEIGHTS_FILE], &dir)?;
        ensure_files(TEXTGEN_TOKENIZER_REPO, &[TEXTGEN_TOKENIZER_FILE], &dir)?;
        Self::from_cache(&dir)
    }

    fn from_cache(dir: &std::path::Path) -> anyhow::Result<Self> {
        let weights_path = dir.join(TEXTGEN_WEIGHTS_FILE);
        let tokenizer_path = dir.join(TEXTGEN_TOKENIZER_FILE);

        let device = Device::Cpu;
        let mut file = std::fs::File::open(&weights_path)
            .with_context(|| format!("open {}", weights_path.display()))?;
        let content =
            gguf_file::Content::read(&mut file).map_err(|e| e.with_path(&weights_path))?;
        let model = Qwen2::from_gguf(content, &mut file, &device)?;
        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!(
            "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\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 prompt_ids = encoding.get_ids().to_vec();
        let mut generated: Vec<u32> = Vec::new();
        let mut index_pos = 0usize;

        let input = Tensor::new(prompt_ids.as_slice(), &self.device)?.unsqueeze(0)?;
        let logits = self.model.forward(&input, index_pos)?;
        let logits = logits.squeeze(0)?;
        let mut next = sampler.sample(&logits)?;
        index_pos += prompt_ids.len();

        for _ in 0..max_tokens {
            if next == self.eos {
                break;
            }
            generated.push(next);
            let input = Tensor::new(&[next], &self.device)?.unsqueeze(0)?;
            let logits = self.model.forward(&input, index_pos)?;
            let logits = logits.squeeze(0)?;
            next = sampler.sample(&logits)?;
            index_pos += 1;
        }

        self.tokenizer
            .decode(&generated, true)
            .map_err(anyhow::Error::msg)
    }
}

fn ensure_files(repo_id: &str, files: &[&str], dest: &std::path::Path) -> anyhow::Result<()> {
    if files.iter().all(|file| dest.join(file).is_file()) {
        return Ok(());
    }
    pull_files(repo_id, files, dest)
}

/// 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)
}