goosedump 0.8.1

// 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 base64::Engine as _;
use candle_core::quantized::gguf_file;
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::bert::{BertModel, Config, DTYPE};
use candle_transformers::models::paddleocr_vl::{Config as OcrConfig, PaddleOCRVLModel};
use candle_transformers::models::quantized_qwen2::ModelWeights;
use sha2::{Digest as _, Sha256};
use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer};

use crate::message::{ConversationMessage, Part};

/// 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 instruct text model and its Hugging Face sources.
/// The quantized weights and the tokenizer live in different repos upstream.
pub(crate) const TEXTGEN_NAME: &str = "Qwen2.5-0.5B-Instruct";
pub(crate) const TEXTGEN_WEIGHTS_REPO: &str = "Qwen/Qwen2.5-0.5B-Instruct-GGUF";
pub(crate) const TEXTGEN_WEIGHTS_FILE: &str = "qwen2.5-0.5b-instruct-q4_k_m.gguf";
pub(crate) const TEXTGEN_TOKENIZER_REPO: &str = "Qwen/Qwen2.5-0.5B-Instruct";
pub(crate) const TEXTGEN_TOKENIZER_FILE: &str = "tokenizer.json";
/// `ChatML` end-of-turn token that terminates generation.
const TEXTGEN_EOS: &str = "<|im_end|>";

/// Cache subdirectory and Hugging Face repo of the OCR vision-language model.
pub(crate) const OCR_NAME: &str = "PaddleOCR-VL";
pub(crate) const OCR_REPO: &str = "PaddlePaddle/PaddleOCR-VL";
pub(crate) const OCR_FILES: [&str; 3] = ["config.json", "tokenizer.json", "model.safetensors"];
/// Vision patch geometry and pixel bounds from the upstream
/// `preprocessor_config.json`.
const OCR_PATCH_SIZE: usize = 14;
const OCR_MIN_PIXELS: usize = 147_384;
const OCR_MAX_PIXELS: usize = 2_822_400;
const OCR_MAX_TOKENS: usize = 1024;

/// Sentence-embedding model (`all-MiniLM-L6-v2`) used for paraphrase-aware dedup
/// 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::load(vb, &config)?;
        Ok(Self {
            model,
            tokenizer,
            device,
        })
    }

    /// Mean-pooled, L2-normalized sentence embeddings — one row per input text,
    /// matching the `sentence-transformers` 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)?;
        let token_type_ids = token_ids.zeros_like()?;

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

/// Small instruct model (`Qwen2.5-0.5B-Instruct`, 4-bit GGUF) 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: ModelWeights,
    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 mut file = std::fs::File::open(&weights_path)?;
        let content = gguf_file::Content::read(&mut file)?;
        let model = ModelWeights::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 `ChatML` 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 mut generated: Vec<u32> = Vec::new();
        let mut input: Vec<u32> = encoding.get_ids().to_vec();
        let mut index_pos = 0;
        for _ in 0..max_tokens {
            let tensor = Tensor::new(input.as_slice(), &self.device)?.unsqueeze(0)?;
            let logits = self.model.forward(&tensor, index_pos)?.squeeze(0)?;
            index_pos += input.len();
            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)
    }
}

/// OCR vision-language model (`PaddleOCR-VL`, 0.9B) that turns document images
/// into text. Greedy decode; absent weights yield `None` and images simply stay
/// text-free.
pub struct Ocr {
    model: PaddleOCRVLModel,
    tokenizer: Tokenizer,
    device: Device,
    image_token_id: u32,
    vision_start_token_id: u32,
    vision_end_token_id: u32,
    spatial_merge: usize,
    bos: u32,
    eos: u32,
}

impl Ocr {
    /// Load the OCR 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(OCR_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(),
            "OCR weights not cached in {}",
            dir.display()
        );

        let device = Device::Cpu;
        let config: OcrConfig = 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::F32, &device)? };
        let bos = tokenizer.token_to_id("<|begin_of_sentence|>").unwrap_or(1);
        let eos = tokenizer
            .token_to_id("</s>")
            .or_else(|| tokenizer.token_to_id("<|end_of_sentence|>"))
            .or_else(|| tokenizer.token_to_id("<|endoftext|>"))
            .unwrap_or(2);
        let model = PaddleOCRVLModel::new(&config, vb)?;
        Ok(Self {
            model,
            tokenizer,
            device,
            image_token_id: config.image_token_id,
            vision_start_token_id: config.vision_start_token_id,
            vision_end_token_id: config.vision_end_token_id,
            spatial_merge: config.vision_config.spatial_merge_size,
            bos,
            eos,
        })
    }

    /// Recognize the text in an encoded image (PNG/JPEG/WebP/GIF bytes).
    ///
    /// # Errors
    /// Returns an error if the image cannot be decoded or a forward pass fails.
    pub fn recognize(&mut self, image_bytes: &[u8]) -> anyhow::Result<String> {
        let (pixel_values, grid_thw, h_patches, w_patches) = self.preprocess(image_bytes)?;
        let num_image_tokens = (h_patches / self.spatial_merge) * (w_patches / self.spatial_merge);

        // `<BOS>User: <IMAGE_START><IMAGE>...<IMAGE_END>OCR:\nAssistant: `,
        // matching the upstream chat format.
        let mut input_ids: Vec<u32> = vec![self.bos];
        input_ids.extend(self.encode_plain("User: ")?);
        input_ids.push(self.vision_start_token_id);
        input_ids.extend(std::iter::repeat_n(self.image_token_id, num_image_tokens));
        input_ids.push(self.vision_end_token_id);
        input_ids.extend(self.encode_plain("OCR:")?);
        input_ids.extend(self.encode_plain("\nAssistant: ")?);
        let input_ids = Tensor::new(input_ids.as_slice(), &self.device)?.unsqueeze(0)?;

        self.model.clear_kv_cache();
        let generated = self.model.generate(
            &input_ids,
            &pixel_values,
            &grid_thw,
            OCR_MAX_TOKENS,
            self.eos,
        )?;
        let text = self
            .tokenizer
            .decode(&generated, true)
            .map_err(anyhow::Error::msg)?;
        Ok(text.trim().to_string())
    }

    fn encode_plain(&self, text: &str) -> anyhow::Result<Vec<u32>> {
        let encoding = self
            .tokenizer
            .encode(text, false)
            .map_err(anyhow::Error::msg)?;
        Ok(encoding.get_ids().to_vec())
    }

    /// Decode, smart-resize, and normalize an image into the `(1, 3, H, W)`
    /// pixel tensor and `(1, 3)` grid the vision encoder expects.
    fn preprocess(&self, image_bytes: &[u8]) -> anyhow::Result<(Tensor, Tensor, usize, usize)> {
        let img = image::load_from_memory(image_bytes)?.to_rgb8();
        let (width, height) = (img.width() as usize, img.height() as usize);
        let factor = OCR_PATCH_SIZE * self.spatial_merge;
        let (new_height, new_width) =
            smart_resize(height, width, factor, OCR_MIN_PIXELS, OCR_MAX_PIXELS)?;

        #[allow(clippy::cast_possible_truncation)]
        let resized = image::imageops::resize(
            &img,
            new_width as u32,
            new_height as u32,
            image::imageops::FilterType::CatmullRom,
        );

        // Normalize to [-1, 1], matching the upstream processor output.
        let mut normalized = vec![0f32; 3 * new_height * new_width];
        for (idx, value) in normalized.iter_mut().enumerate() {
            let c = idx / (new_height * new_width);
            let y = (idx / new_width) % new_height;
            let x = idx % new_width;
            #[allow(clippy::cast_possible_truncation)]
            let pixel = resized.get_pixel(x as u32, y as u32);
            *value = f32::from(pixel[c]) / 255.0 * 2.0 - 1.0;
        }

        let pixel_values =
            Tensor::from_vec(normalized, (1, 3, new_height, new_width), &self.device)?;
        let h_patches = new_height / OCR_PATCH_SIZE;
        let w_patches = new_width / OCR_PATCH_SIZE;
        #[allow(clippy::cast_possible_truncation)]
        let grid_thw = Tensor::new(&[[1u32, h_patches as u32, w_patches as u32]], &self.device)?;
        Ok((pixel_values, grid_thw, h_patches, w_patches))
    }
}

/// Rescale image dimensions so both are divisible by `factor` and the pixel
/// count lands within `[min_pixels, max_pixels]`, preserving aspect ratio —
/// the upstream `smart_resize` algorithm.
#[allow(
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss
)]
fn smart_resize(
    height: usize,
    width: usize,
    factor: usize,
    min_pixels: usize,
    max_pixels: usize,
) -> anyhow::Result<(usize, usize)> {
    let mut h = height;
    let mut w = width;
    anyhow::ensure!(h > 0 && w > 0, "empty image");
    if h < factor {
        w = (w * factor + h / 2) / h;
        h = factor;
    }
    if w < factor {
        h = (h * factor + w / 2) / w;
        w = factor;
    }

    let aspect = if h > w {
        h as f64 / w as f64
    } else {
        w as f64 / h as f64
    };
    anyhow::ensure!(aspect <= 200.0, "aspect ratio {aspect:.1} exceeds 200");

    let mut h_bar = ((h + factor / 2) / factor) * factor;
    let mut w_bar = ((w + factor / 2) / factor) * factor;
    let total_pixels = h_bar * w_bar;
    if total_pixels > max_pixels {
        let beta = ((h * w) as f64 / max_pixels as f64).sqrt();
        h_bar = ((h as f64 / beta / factor as f64).floor() as usize) * factor;
        w_bar = ((w as f64 / beta / factor as f64).floor() as usize) * factor;
    } else if total_pixels < min_pixels {
        let beta = (min_pixels as f64 / (h * w) as f64).sqrt();
        h_bar = ((h as f64 * beta / factor as f64).ceil() as usize) * factor;
        w_bar = ((w as f64 * beta / factor as f64).ceil() as usize) * factor;
    }
    Ok((h_bar, w_bar))
}

/// Fill `ocr_text` on every image part from the process-once OCR cache,
/// running the OCR model only for uncached images and only when its weights
/// are present. Absent weights or undecodable payloads leave images text-free.
pub fn ocr_enrich(messages: &mut [ConversationMessage]) {
    // Lazy two-level option: outer `None` until the first cache miss forces a
    // load attempt; inner `None` when the weights are not cached.
    let mut ocr: Option<Option<Ocr>> = None;
    for message in messages {
        for part in &mut message.parts {
            let Part::Image(image) = part else { continue };
            if !image.ocr_text.is_empty() {
                continue;
            }
            let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(image.data.trim())
            else {
                continue;
            };
            let cache_path = {
                let digest = Sha256::digest(&bytes);
                let mut name = String::with_capacity(68);
                for byte in digest {
                    use std::fmt::Write as _;
                    let _ = write!(name, "{byte:02x}");
                }
                name.push_str(".txt");
                dirs::cache_dir()
                    .unwrap_or_else(|| PathBuf::from("."))
                    .join("goosedump")
                    .join("ocr")
                    .join(name)
            };
            if let Ok(text) = std::fs::read_to_string(&cache_path) {
                image.ocr_text = text;
                continue;
            }
            let Some(model) = ocr.get_or_insert_with(Ocr::load).as_mut() else {
                continue;
            };
            let Ok(text) = model.recognize(&bytes) else {
                continue;
            };
            if let Some(parent) = cache_path.parent() {
                let _ = std::fs::create_dir_all(parent);
            }
            let _ = std::fs::write(&cache_path, &text);
            image.ocr_text = text;
        }
    }
}

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