candding 0.1.0

//! `BAAI/bge-m3`'s sparse head: `sparse_linear.pt` gives one logit per position, scattered host-side into the
//! column its input id names, keeping the larger logit on a shared id (candle has no `scatter_reduce`).

use std::path::Path;

use candle_core::DType;
use candle_core::Module;
use candle_core::Tensor;
use candle_core::pickle::PthTensors;
use candle_nn::Linear;
use candle_nn::VarBuilder;

use crate::error::CanddingError;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::SparseModel;
use crate::models::xlm_roberta::XlmRobertaConfig;
use crate::models::xlm_roberta::XlmRobertaModel;
use crate::sparse::SparseEmbedding;

/// `<s>`, `</s>`, `<pad>`, `<unk>`: the four ids the reference zeros after the scatter (`modeling.py:152-156`).
const DROPPED: [u32; 4] = [0, 2, 1, 3];

/// `BAAI/bge-m3`'s sparse head, built on `xlm_roberta::XlmRobertaModel`.
#[derive(Debug)]
pub struct M3Sparse {
  encoder: XlmRobertaModel,
  /// `(1, 1024)`, already `(out, in)` as `Linear` wants; not an embedding table, so no transpose.
  weight: Tensor,
  /// `sparse_linear.pt`'s `bias`, `(1,)`.
  bias: Tensor,
  vocab_size: usize,
  /// Always `DROPPED`.
  dropped: [u32; 4],
}

/// Like `PthTensors::get`, but a missing tensor names the file too.
fn read_tensor(pth: &PthTensors, name: &str, file: &Path) -> Result<Tensor> {
  pth.get(name)?.ok_or_else(|| CanddingError::MissingTensor {
    file: file.to_path_buf(),
    name: name.to_string(),
  })
}

impl M3Sparse {
  /// Reads `sparse_linear.pt` (f16 on disk) through candle's pickle reader. A missing file is an error, where
  /// FlagEmbedding silently builds a random head.
  pub fn load(config: XlmRobertaConfig, vb: VarBuilder, sparse_linear: &Path) -> Result<Self> {
    if !sparse_linear.is_file() {
      return Err(CanddingError::MissingFile(sparse_linear.to_path_buf()));
    }
    let pth = PthTensors::new(sparse_linear, None)?;
    let weight = read_tensor(&pth, "weight", sparse_linear)?;
    let bias = read_tensor(&pth, "bias", sparse_linear)?;
    Self::from_raw_head(config, vb, weight, bias)
  }

  /// Casts the head to the run's dtype and builds the encoder; split out so tests can pass hand-built tensors.
  fn from_raw_head(
    config: XlmRobertaConfig,
    vb: VarBuilder,
    weight: Tensor,
    bias: Tensor,
  ) -> Result<Self> {
    let vocab_size = config.bert.vocab_size;
    let dtype = vb.dtype();
    let weight = weight.to_dtype(dtype)?.to_device(vb.device())?;
    let bias = bias.to_dtype(dtype)?.to_device(vb.device())?;
    let encoder = XlmRobertaModel::load(config, vb)?;
    Ok(Self {
      encoder,
      weight,
      bias,
      vocab_size,
      dropped: DROPPED,
    })
  }

  /// `relu(h · weightᵀ + bias)` per position, truncated to each row's length and minus `dropped`, folded by
  /// `SparseEmbedding::from_pairs` (max per id); nothing here scales with `vocab_size`.
  fn scatter_rows(
    &self,
    hidden: &Tensor,
    input_ids: &Tensor,
    lengths: &[usize],
  ) -> Result<Vec<SparseEmbedding>> {
    let linear = Linear::new(self.weight.clone(), Some(self.bias.clone()));
    let logits = linear
      .forward(hidden)?
      .relu()?
      .squeeze(2)?
      .to_dtype(DType::F32)?;
    let weights = logits.to_vec2::<f32>()?;
    let ids = input_ids.to_vec2::<u32>()?;
    Ok(
      weights
        .into_iter()
        .zip(ids)
        .zip(lengths)
        .map(|((row_weights, row_ids), &len)| {
          SparseEmbedding::from_pairs(
            row_ids
              .into_iter()
              .zip(row_weights)
              .take(len)
              .filter(|(id, _)| !self.dropped.contains(id)),
          )
        })
        .collect(),
    )
  }
}

impl SparseModel for M3Sparse {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    self.scatter_rows(&hidden, input.input_ids, input.lengths)
  }

  fn vocab_size(&self) -> usize {
    self.vocab_size
  }

  fn max_position_embeddings(&self) -> usize {
    self.encoder.max_position_embeddings()
  }

  fn weights_dtype(&self) -> DType {
    self.encoder.weights_dtype()
  }
}

#[cfg(test)]
mod tests;