goosedump 0.9.13

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

// Based on BERT model code from https://github.com/huggingface/candle,
// which is licensed with MIT OR Apache-2.0.

#![allow(dead_code, clippy::all, clippy::pedantic)]

//! Quantized BERT encoder for sentence embeddings.

use anyhow::{Context as _, Result};
use candle_core::quantized::gguf_file;
use candle_core::quantized::{GgmlDType, QTensor};
use candle_core::safetensors::MmapedSafetensors;
use candle_core::{DType, Device, Module, Result as CResult, Tensor};
use candle_nn::LayerNorm;
use candle_transformers::models::bert::{Config, HiddenAct};
use candle_transformers::quantized_nn::{Embedding, Linear, layer_norm, linear};
use candle_transformers::quantized_var_builder::VarBuilder;
use std::io::Write as _;
use std::path::Path;

pub const Q8_0_GGUF_NAME: &str = "all-MiniLM-L6-v2-q8_0.gguf";

struct HiddenActLayer {
    act: HiddenAct,
}

impl HiddenActLayer {
    fn new(act: HiddenAct) -> Self {
        Self { act }
    }
}

impl Module for HiddenActLayer {
    fn forward(&self, xs: &Tensor) -> CResult<Tensor> {
        match self.act {
            HiddenAct::Gelu => xs.gelu_erf(),
            HiddenAct::GeluApproximate => xs.gelu(),
            HiddenAct::Relu => xs.relu(),
        }
    }
}

struct Dropout;

impl Module for Dropout {
    fn forward(&self, x: &Tensor) -> CResult<Tensor> {
        Ok(x.clone())
    }
}

struct BertEmbeddings {
    word_embeddings: Embedding,
    position_embeddings: Embedding,
    token_type_embeddings: Embedding,
    layer_norm: LayerNorm,
}

impl BertEmbeddings {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let word_embeddings = Embedding::new(
            config.vocab_size,
            config.hidden_size,
            vb.pp("word_embeddings"),
        )?;
        let position_embeddings = Embedding::new(
            config.max_position_embeddings,
            config.hidden_size,
            vb.pp("position_embeddings"),
        )?;
        let token_type_embeddings = Embedding::new(
            config.type_vocab_size,
            config.hidden_size,
            vb.pp("token_type_embeddings"),
        )?;
        let layer_norm = layer_norm(
            config.hidden_size,
            config.layer_norm_eps,
            vb.pp("LayerNorm"),
        )?;
        Ok(Self {
            word_embeddings,
            position_embeddings,
            token_type_embeddings,
            layer_norm,
        })
    }

    fn forward(&self, input_ids: &Tensor, token_type_ids: &Tensor) -> CResult<Tensor> {
        let (_bsize, seq_len) = input_ids.dims2()?;
        let input_embeddings = self.word_embeddings.forward(input_ids)?;
        let token_type_embeddings = self.token_type_embeddings.forward(token_type_ids)?;
        let mut embeddings = (&input_embeddings + token_type_embeddings)?;
        let position_ids = (0..seq_len as u32).collect::<Vec<_>>();
        let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
        embeddings = embeddings.broadcast_add(&self.position_embeddings.forward(&position_ids)?)?;
        let embeddings = self.layer_norm.forward(&embeddings)?;
        Ok(embeddings)
    }
}

struct BertSelfAttention {
    query: Linear,
    key: Linear,
    value: Linear,
    num_attention_heads: usize,
    attention_head_size: usize,
}

impl BertSelfAttention {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let attention_head_size = config.hidden_size / config.num_attention_heads;
        let all_head_size = config.num_attention_heads * attention_head_size;
        let hidden_size = config.hidden_size;
        let query = linear(hidden_size, all_head_size, vb.pp("query"))?;
        let value = linear(hidden_size, all_head_size, vb.pp("value"))?;
        let key = linear(hidden_size, all_head_size, vb.pp("key"))?;
        Ok(Self {
            query,
            key,
            value,
            num_attention_heads: config.num_attention_heads,
            attention_head_size,
        })
    }

    fn transpose_for_scores(&self, xs: &Tensor) -> CResult<Tensor> {
        let mut new_x_shape = xs.dims().to_vec();
        new_x_shape.pop();
        new_x_shape.push(self.num_attention_heads);
        new_x_shape.push(self.attention_head_size);
        let xs = xs.reshape(new_x_shape.as_slice())?.transpose(1, 2)?;
        xs.contiguous()
    }

    fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> CResult<Tensor> {
        let query_layer = self.transpose_for_scores(&self.query.forward(hidden_states)?)?;
        let key_layer = self.transpose_for_scores(&self.key.forward(hidden_states)?)?;
        let value_layer = self.transpose_for_scores(&self.value.forward(hidden_states)?)?;

        let attention_scores = query_layer.matmul(&key_layer.t()?)?;
        let attention_scores = (attention_scores / (self.attention_head_size as f64).sqrt())?;
        let attention_scores = attention_scores.broadcast_add(attention_mask)?;
        let attention_probs = candle_nn::ops::softmax(&attention_scores, candle_core::D::Minus1)?;

        let context_layer = attention_probs.matmul(&value_layer)?;
        let context_layer = context_layer.transpose(1, 2)?.contiguous()?;
        let context_layer = context_layer.flatten_from(candle_core::D::Minus2)?;
        Ok(context_layer)
    }
}

struct BertSelfOutput {
    dense: Linear,
    layer_norm: LayerNorm,
    dropout: Dropout,
}

impl BertSelfOutput {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let dense = linear(config.hidden_size, config.hidden_size, vb.pp("dense"))?;
        let layer_norm = layer_norm(
            config.hidden_size,
            config.layer_norm_eps,
            vb.pp("LayerNorm"),
        )?;
        Ok(Self {
            dense,
            layer_norm,
            dropout: Dropout,
        })
    }

    fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> CResult<Tensor> {
        let hidden_states = self.dense.forward(hidden_states)?;
        self.layer_norm.forward(&(hidden_states + input_tensor)?)
    }
}

struct BertAttention {
    self_attention: BertSelfAttention,
    self_output: BertSelfOutput,
}

impl BertAttention {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let self_attention = BertSelfAttention::load(vb.pp("self"), config)?;
        let self_output = BertSelfOutput::load(vb.pp("output"), config)?;
        Ok(Self {
            self_attention,
            self_output,
        })
    }

    fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> CResult<Tensor> {
        let self_outputs = self.self_attention.forward(hidden_states, attention_mask)?;
        self.self_output.forward(&self_outputs, hidden_states)
    }
}

struct BertIntermediate {
    dense: Linear,
    intermediate_act: HiddenActLayer,
}

impl BertIntermediate {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let dense = linear(config.hidden_size, config.intermediate_size, vb.pp("dense"))?;
        Ok(Self {
            dense,
            intermediate_act: HiddenActLayer::new(config.hidden_act),
        })
    }
}

impl Module for BertIntermediate {
    fn forward(&self, hidden_states: &Tensor) -> CResult<Tensor> {
        let hidden_states = self.dense.forward(hidden_states)?;
        self.intermediate_act.forward(&hidden_states)
    }
}

struct BertOutput {
    dense: Linear,
    layer_norm: LayerNorm,
}

impl BertOutput {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let dense = linear(config.intermediate_size, config.hidden_size, vb.pp("dense"))?;
        let layer_norm = layer_norm(
            config.hidden_size,
            config.layer_norm_eps,
            vb.pp("LayerNorm"),
        )?;
        Ok(Self { dense, layer_norm })
    }

    fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> CResult<Tensor> {
        let hidden_states = self.dense.forward(hidden_states)?;
        self.layer_norm.forward(&(hidden_states + input_tensor)?)
    }
}

pub struct BertLayer {
    attention: BertAttention,
    intermediate: BertIntermediate,
    output: BertOutput,
}

impl BertLayer {
    fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let attention = BertAttention::load(vb.pp("attention"), config)?;
        let intermediate = BertIntermediate::load(vb.pp("intermediate"), config)?;
        let output = BertOutput::load(vb.pp("output"), config)?;
        Ok(Self {
            attention,
            intermediate,
            output,
        })
    }

    fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> CResult<Tensor> {
        let attention_output = self.attention.forward(hidden_states, attention_mask)?;
        let intermediate_output = self.intermediate.forward(&attention_output)?;
        self.output.forward(&intermediate_output, &attention_output)
    }
}

pub struct BertEncoder {
    pub layers: Vec<BertLayer>,
}

impl BertEncoder {
    pub fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let layers = (0..config.num_hidden_layers)
            .map(|index| BertLayer::load(vb.pp(format!("layer.{index}")), config))
            .collect::<CResult<Vec<_>>>()?;
        Ok(BertEncoder { layers })
    }

    pub fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> CResult<Tensor> {
        let mut hidden_states = hidden_states.clone();
        for layer in self.layers.iter() {
            hidden_states = layer.forward(&hidden_states, attention_mask)?;
        }
        Ok(hidden_states)
    }
}

pub struct BertModel {
    embeddings: BertEmbeddings,
    encoder: BertEncoder,
    pub device: Device,
}

impl BertModel {
    pub fn load(vb: VarBuilder, config: &Config) -> CResult<Self> {
        let embeddings = BertEmbeddings::load(vb.pp("embeddings"), config)?;
        let encoder = BertEncoder::load(vb.pp("encoder"), config)?;
        Ok(Self {
            embeddings,
            encoder,
            device: vb.device().clone(),
        })
    }

    pub fn forward(
        &self,
        input_ids: &Tensor,
        token_type_ids: &Tensor,
        attention_mask: Option<&Tensor>,
    ) -> CResult<Tensor> {
        let embedding_output = self.embeddings.forward(input_ids, token_type_ids)?;
        let attention_mask = match attention_mask {
            Some(attention_mask) => attention_mask.clone(),
            None => input_ids.ones_like()?,
        };
        let dtype = embedding_output.dtype();
        let attention_mask = get_extended_attention_mask(&attention_mask, dtype)?;
        self.encoder.forward(&embedding_output, &attention_mask)
    }
}

fn get_extended_attention_mask(attention_mask: &Tensor, dtype: DType) -> CResult<Tensor> {
    let attention_mask = match attention_mask.rank() {
        3 => attention_mask.unsqueeze(1)?,
        2 => attention_mask.unsqueeze(1)?.unsqueeze(1)?,
        _ => candle_core::bail!("Wrong shape for input_ids or attention_mask"),
    };
    let attention_mask = attention_mask.to_dtype(dtype)?;
    (attention_mask.ones_like()? - &attention_mask)?.broadcast_mul(
        &Tensor::try_from(f32::MIN)?
            .to_device(attention_mask.device())?
            .to_dtype(dtype)?,
    )
}

fn quantize_dtype_for(name: &str, rank: usize) -> GgmlDType {
    if rank == 2 && name.ends_with(".weight") {
        GgmlDType::Q8_0
    } else {
        GgmlDType::F32
    }
}

pub fn build_q8_0_gguf(f32_safetensors: &Path, out_gguf: &Path) -> Result<()> {
    let device = Device::Cpu;
    let mst = unsafe { MmapedSafetensors::multi(&[f32_safetensors]) }
        .with_context(|| format!("mmap {}", f32_safetensors.display()))?;
    let names: Vec<String> = mst.tensors().into_iter().map(|(name, _)| name).collect();
    eprintln!("building q8_0 GGUF ({} tensors)...", names.len());
    let mut qtensors: Vec<(String, QTensor)> = Vec::with_capacity(names.len());
    for name in &names {
        let tensor = mst.load(name, &device)?;
        // Skip non-float tensors (e.g. the I64 `position_ids`); the model
        // recomputes positions via arange, so the stored copy is unused.
        if tensor.dtype() != DType::F32 && tensor.dtype() != DType::F16 {
            continue;
        }
        let dtype = quantize_dtype_for(name, tensor.rank());
        let q = QTensor::quantize(&tensor, dtype)
            .with_context(|| format!("quantize {name} to {dtype:?}"))?;
        qtensors.push((name.clone(), q));
    }
    let file = std::fs::File::create(out_gguf)
        .with_context(|| format!("create {}", out_gguf.display()))?;
    let mut buf = std::io::BufWriter::new(file);
    let refs: Vec<(&str, &QTensor)> = qtensors.iter().map(|(n, q)| (n.as_str(), q)).collect();
    gguf_file::write(&mut buf, &[], &refs)
        .with_context(|| format!("write GGUF {}", out_gguf.display()))?;
    buf.flush().context("flush GGUF")?;
    eprintln!("q8_0 GGUF done: {}", out_gguf.display());
    Ok(())
}