candding 0.1.0

//! `Qdrant/minicoil-v1`'s sparse head over the jina-embeddings-v2-small-en encoder's hidden states.
//! Query and passage differ only in the term weight, so this head overrides `forward_query`.

use candle_core::DType;
use candle_core::Tensor;

use super::index;
use super::pooling;
use super::vocab::Vocab;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::SparseModel;
use crate::models::sparse::stopwords::Stopwords;
use crate::sparse::SparseEmbedding;

/// `Qdrant/minicoil-v1`'s sparse head; `new` takes the side files already loaded. `encoder` is a `dyn Encoder`
/// so the registry builds it through the shared `load_encoder`, which is why there is no `Debug`.
pub struct MiniCoilSparse {
  encoder: Box<dyn Encoder>,
  /// `(vocab_size_including_unrecognized, hidden_size, embedding_size)` in F32, whatever the encoder's dtype.
  matrix: Tensor,
  vocab: Vocab,
  stopwords: Stopwords,
  /// The encoder's id-indexed WordPiece vocabulary (`##` marks a continuation), not minicoil's word vocabulary.
  token_text: Vec<String>,
}

impl MiniCoilSparse {
  /// Checks `matrix` against `vocab` (`dims[0] == vocab.len() + 1`) and against `index::SHIFT`, since a grown
  /// vocabulary can stay self-consistent while its dense indices collide with the hashed range.
  pub fn new(
    encoder: Box<dyn Encoder>,
    matrix: Tensor,
    vocab: Vocab,
    stopwords: Stopwords,
    token_text: Vec<String>,
  ) -> Result<Self> {
    let dims = matrix.dims().to_vec();
    if dims.len() != 3 || dims[2] != index::EMBEDDING_SIZE as usize {
      return Err(
        candle_core::Error::Msg(format!(
          "minicoil's per-word matrix must be 3-dimensional with a last dimension of {}, got {dims:?}",
          index::EMBEDDING_SIZE
        ))
        .into(),
      );
    }
    if dims[0] != vocab.len() + 1 {
      return Err(
        candle_core::Error::Msg(format!(
          "minicoil's per-word matrix declares {} rows but the vocabulary's {} words plus the reserved \
           unrecognized id 0 is {}",
          dims[0],
          vocab.len(),
          vocab.len() + 1
        ))
        .into(),
      );
    }
    let max_dense_index = dims[0] as u64 * index::EMBEDDING_SIZE as u64;
    if max_dense_index > index::SHIFT as u64 {
      return Err(
        candle_core::Error::Msg(format!(
          "minicoil's per-word matrix has {} rows, {} more than the {} this crate's compiled-in index \
           arithmetic assumes (plus the reserved unrecognized id 0); a dense index built from a word this far \
           into the vocabulary would reach {}, past the {} boundary where the hashed range for unrecognized \
           words begins, silently colliding the two ranges — index.rs's own constants need updating for a \
           vocabulary this large before it can be loaded safely",
          dims[0],
          dims[0].saturating_sub(index::VOCAB_SIZE_INCLUDING_UNRECOGNIZED as usize),
          index::VOCAB_SIZE_INCLUDING_UNRECOGNIZED,
          max_dense_index - 1,
          index::SHIFT
        ))
        .into(),
      );
    }
    let matrix = matrix.to_dtype(DType::F32)?;
    Ok(Self {
      encoder,
      matrix,
      vocab,
      stopwords,
      token_text,
    })
  }

  /// Pools each row through `pooling::pool_row` in F32. The encoder's hidden width is checked first: `zip` in
  /// the pooling would silently truncate a mismatch.
  fn embed_rows(
    &self,
    hidden: &Tensor,
    input_ids: &Tensor,
    lengths: &[usize],
    uniform_weight: bool,
  ) -> Result<Vec<SparseEmbedding>> {
    let matrix_hidden_size = self.matrix.dims()[1];
    if hidden.dims()[2] != matrix_hidden_size {
      return Err(
        candle_core::Error::Msg(format!(
          "the encoder's own hidden width ({}) does not match minicoil's per-word matrix ({}); pairing this \
           head with a mismatched encoder checkpoint would otherwise silently truncate every hidden state to \
           the narrower of the two instead of erroring",
          hidden.dims()[2],
          matrix_hidden_size
        ))
        .into(),
      );
    }
    let hidden = hidden.to_dtype(DType::F32)?;
    let hidden_rows = hidden.to_vec3::<f32>()?;
    let id_rows = input_ids.to_vec2::<u32>()?;
    hidden_rows
      .iter()
      .zip(id_rows.iter())
      .zip(lengths.iter())
      .map(|((hidden_row, id_row), &len)| {
        let len = len.min(id_row.len());
        pooling::pool_row(
          hidden_row,
          &id_row[..len],
          &self.token_text,
          &self.vocab,
          &self.stopwords,
          &self.matrix,
          uniform_weight,
        )
      })
      .collect()
  }
}

impl SparseModel for MiniCoilSparse {
  /// The passage path: term-frequency weights over reconstructed words.
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    self.embed_rows(&hidden, input.input_ids, input.lengths, false)
  }

  /// The query path: every term weight is 1.0, as the reference's `is_query=True` does.
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    self.embed_rows(&hidden, input.input_ids, input.lengths, true)
  }

  /// The per-word matrix's first dimension, counting the reserved unrecognized id 0 as the reference does.
  fn vocab_size(&self) -> usize {
    self.matrix.dims()[0]
  }

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

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

#[cfg(test)]
mod tests;