candding 0.1.0

//! `naver/splade-code-06B`'s sparse head: SPLADE's max-pooled `relu(log1p(logits))`, decoding off Qwen3's tied
//! embedding table with no transform or bias. Its tokenizer adds no special tokens, so rows can be empty.

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

use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::SparseModel;
use crate::models::qwen3::Qwen3Config;
use crate::models::qwen3::Qwen3Model;
use crate::sparse::SparseEmbedding;

/// `splade::CHUNK` over a 5x larger vocabulary: about 1.19 GiB per chunk at batch 8, whatever the max length.
const CHUNK: usize = 64;

/// `naver/splade-code-06B`'s sparse head, built on `qwen3::Qwen3Model`.
pub struct SpladeTiedSparse {
  qwen3: Qwen3Model,
  /// `(vocab, hidden)`, shared via `Qwen3Model::embed_tokens`; the checkpoint's `lm_head` has no bias.
  embed_tokens: Tensor,
  vocab_size: usize,
  chunk: usize,
}

impl SpladeTiedSparse {
  /// Builds the Qwen3 backbone and reuses its tied embedding table as the decoder.
  pub fn load(config: Qwen3Config, vb: VarBuilder) -> Result<Self> {
    let qwen3 = Qwen3Model::load(config, vb)?;
    let embed_tokens = qwen3.embed_tokens().clone();
    let vocab_size = embed_tokens.dims()[0];
    Ok(Self {
      qwen3,
      embed_tokens,
      vocab_size,
      chunk: CHUNK,
    })
  }

  /// `hidden · embed_tokensᵀ`: no transform, no bias, and no transpose since `Linear` already wants `(out, in)`.
  fn chunk_logits(&self, h_chunk: &Tensor) -> Result<Tensor> {
    let decoder = Linear::new(self.embed_tokens.clone(), None);
    Ok(decoder.forward(h_chunk)?)
  }

  /// Folds a running `(batch, vocab)` maximum over sequence chunks like `SpladeSparse::pool_logits`;
  /// a zero-length `hidden` gives every row an empty embedding instead of a panic.
  fn pool_logits(&self, hidden: &Tensor, attention_mask: &Tensor) -> Result<Vec<SparseEmbedding>> {
    let (batch, seq, _) = hidden.dims3()?;
    let mut pooled: Option<Tensor> = None;
    let mut start = 0;
    while start < seq {
      let len = self.chunk.min(seq - start);
      let h_chunk = hidden.narrow(1, start, len)?;
      let mask_chunk = attention_mask
        .narrow(1, start, len)?
        .to_dtype(hidden.dtype())?
        .unsqueeze(2)?;
      let masked = self.chunk_logits(&h_chunk)?.broadcast_mul(&mask_chunk)?;
      let weight = masked.relu()?.affine(1.0, 1.0)?.log()?;
      let chunk_max = weight.max(1)?;
      pooled = Some(match pooled {
        Some(prev) => prev.broadcast_maximum(&chunk_max)?,
        None => chunk_max,
      });
      start += len;
    }
    let Some(pooled) = pooled else {
      return Ok(vec![SparseEmbedding::from_pairs([]); batch]);
    };
    let rows = pooled.to_dtype(DType::F32)?.to_vec2::<f32>()?;
    Ok(
      rows
        .into_iter()
        .map(|row| {
          SparseEmbedding::from_pairs(row.into_iter().enumerate().map(|(i, v)| (i as u32, v)))
        })
        .collect(),
    )
  }

  /// A zero-length row beside real ones would leave its queries no visible key (NaN on Metal), so the real rows
  /// are encoded on their own at the same width and empty embeddings are spliced back in.
  fn forward_with_empty_rows(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let mut rows = vec![SparseEmbedding::from_pairs([]); input.lengths.len()];
    let real_rows: Vec<usize> = input
      .lengths
      .iter()
      .enumerate()
      .filter(|&(_, &len)| len > 0)
      .map(|(i, _)| i)
      .collect();
    if real_rows.is_empty() {
      return Ok(rows);
    }
    let device = input.input_ids.device();
    let index: Vec<u32> = real_rows.iter().map(|&i| i as u32).collect();
    let index = Tensor::from_vec(index, real_rows.len(), device)?;
    let input_ids = input.input_ids.index_select(&index, 0)?;
    let attention_mask = input.attention_mask.index_select(&index, 0)?;
    let token_type_ids = input.token_type_ids.index_select(&index, 0)?;
    let lengths: Vec<usize> = real_rows.iter().map(|&i| input.lengths[i]).collect();
    let real_input = EncoderInput {
      input_ids: &input_ids,
      attention_mask: &attention_mask,
      token_type_ids: &token_type_ids,
      lengths: &lengths,
    };
    let hidden = self.qwen3.forward(&real_input)?;
    let real_rows_out = self.pool_logits(&hidden, &attention_mask)?;
    for (row, embedding) in real_rows.into_iter().zip(real_rows_out) {
      rows[row] = embedding;
    }
    Ok(rows)
  }
}

impl SparseModel for SpladeTiedSparse {
  /// All-empty batches return empty rows before the backbone runs; mixed ones go to `forward_with_empty_rows`.
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let (batch, seq) = input.input_ids.dims2()?;
    if seq == 0 {
      return Ok(vec![SparseEmbedding::from_pairs([]); batch]);
    }
    if input.lengths.contains(&0) {
      return self.forward_with_empty_rows(input);
    }
    let hidden = self.qwen3.forward(input)?;
    self.pool_logits(&hidden, input.attention_mask)
  }

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

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

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

#[cfg(test)]
mod tests;