candding 0.1.0

//! The sentence-transformers `Dense` module: a linear layer applied to pooled vectors, always in F32.

use std::fs;
use std::path::Path;

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::Result;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DenseActivation {
  Identity,
  Tanh,
}

#[derive(Deserialize)]
struct DenseConfig {
  in_features: usize,
  out_features: usize,
  #[serde(default = "default_bias")]
  bias: bool,
  #[serde(default = "default_activation")]
  activation_function: String,
}

fn default_bias() -> bool {
  true
}

fn default_activation() -> String {
  "torch.nn.modules.linear.Identity".to_string()
}

#[derive(Debug)]
pub struct Dense {
  /// Stored transposed, `(in_features, out_features)`, so the forward pass is one matmul.
  weight_t: Tensor,
  bias: Option<Tensor>,
  activation: DenseActivation,
  in_features: usize,
  out_features: usize,
}

impl Dense {
  /// Reads a `{n}_Dense` directory's `config.json` and `model.safetensors` (`linear.weight`, optional `linear.bias`).
  pub fn load(dir: &Path, device: &Device) -> Result<Self> {
    let config: DenseConfig = serde_json::from_str(&fs::read_to_string(dir.join("config.json"))?)?;
    let activation = match config.activation_function.as_str() {
      "torch.nn.modules.linear.Identity" => DenseActivation::Identity,
      "torch.nn.modules.activation.Tanh" => DenseActivation::Tanh,
      other => return Err(CanddingError::UnsupportedDenseActivation(other.to_string())),
    };
    let file = dir.join("model.safetensors");
    if !file.is_file() {
      return Err(CanddingError::MissingFile(file));
    }
    let mut tensors = candle_core::safetensors::load(&file, device)?;
    let mut take = |name: &str| -> Result<Tensor> {
      Ok(
        tensors
          .remove(name)
          .ok_or_else(|| CanddingError::MissingTensor {
            file: file.clone(),
            name: name.to_string(),
          })?
          .to_dtype(DType::F32)?,
      )
    };
    let weight = take("linear.weight")?;
    if weight.dims() != [config.out_features, config.in_features] {
      return Err(CanddingError::Candle(candle_core::Error::msg(format!(
        "dense weight has shape {:?}, config says ({}, {})",
        weight.dims(),
        config.out_features,
        config.in_features
      ))));
    }
    let bias = if config.bias {
      Some(take("linear.bias")?)
    } else {
      None
    };
    Ok(Self {
      weight_t: weight.t()?.contiguous()?,
      bias,
      activation,
      in_features: config.in_features,
      out_features: config.out_features,
    })
  }

  pub fn in_features(&self) -> usize {
    self.in_features
  }

  pub fn out_features(&self) -> usize {
    self.out_features
  }

  /// `xs` is `(batch, in_features)` in F32; returns `(batch, out_features)`.
  pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    let ys = xs.matmul(&self.weight_t)?;
    let ys = match &self.bias {
      Some(bias) => ys.broadcast_add(bias)?,
      None => ys,
    };
    Ok(match self.activation {
      DenseActivation::Identity => ys,
      DenseActivation::Tanh => ys.tanh()?,
    })
  }
}

#[cfg(test)]
mod tests {
  use std::collections::HashMap;
  use std::fs;

  use candle_core::Device;
  use candle_core::Tensor;
  use tempfile::tempdir;

  use super::*;

  fn write_dense(dir: &std::path::Path, config: &str, tensors: HashMap<String, Tensor>) {
    fs::write(dir.join("config.json"), config).unwrap();
    candle_core::safetensors::save(&tensors, dir.join("model.safetensors")).unwrap();
  }

  #[test]
  fn identity_dense_applies_weight_and_bias() {
    let dir = tempdir().unwrap();
    let weight = Tensor::new(&[[1f32, 0.0], [0.0, 1.0], [1.0, 1.0]], &Device::Cpu).unwrap();
    let bias = Tensor::new(&[0.5f32, 0.0, -1.0], &Device::Cpu).unwrap();
    write_dense(
      dir.path(),
      r#"{"in_features":2,"out_features":3,"bias":true,"activation_function":"torch.nn.modules.linear.Identity"}"#,
      HashMap::from([
        ("linear.weight".to_string(), weight),
        ("linear.bias".to_string(), bias),
      ]),
    );
    let dense = Dense::load(dir.path(), &Device::Cpu).unwrap();
    assert_eq!((dense.in_features(), dense.out_features()), (2, 3));
    let out = dense
      .forward(&Tensor::new(&[[2f32, 3.0]], &Device::Cpu).unwrap())
      .unwrap();
    assert_eq!(out.to_vec2::<f32>().unwrap(), vec![vec![2.5, 3.0, 4.0]]);
  }

  #[test]
  fn tanh_dense_without_bias() {
    let dir = tempdir().unwrap();
    let weight = Tensor::new(&[[0.5f32, 0.0], [0.0, 0.5]], &Device::Cpu).unwrap();
    write_dense(
      dir.path(),
      r#"{"in_features":2,"out_features":2,"bias":false,"activation_function":"torch.nn.modules.activation.Tanh"}"#,
      HashMap::from([("linear.weight".to_string(), weight)]),
    );
    let dense = Dense::load(dir.path(), &Device::Cpu).unwrap();
    let out = dense
      .forward(&Tensor::new(&[[2f32, -2.0]], &Device::Cpu).unwrap())
      .unwrap()
      .to_vec2::<f32>()
      .unwrap();
    assert!((out[0][0] - 1f32.tanh()).abs() < 1e-6 && (out[0][1] + 1f32.tanh()).abs() < 1e-6);
  }

  #[test]
  fn unknown_activation_and_missing_tensor_are_reported() {
    let dir = tempdir().unwrap();
    let weight = Tensor::zeros((2, 2), candle_core::DType::F32, &Device::Cpu).unwrap();
    write_dense(
      dir.path(),
      r#"{"in_features":2,"out_features":2,"bias":false,"activation_function":"torch.nn.modules.activation.GELU"}"#,
      HashMap::from([("linear.weight".to_string(), weight.clone())]),
    );
    assert!(matches!(
      Dense::load(dir.path(), &Device::Cpu).unwrap_err(),
      CanddingError::UnsupportedDenseActivation(name) if name.ends_with("GELU")
    ));
    write_dense(
      dir.path(),
      r#"{"in_features":2,"out_features":2,"bias":true,"activation_function":"torch.nn.modules.linear.Identity"}"#,
      HashMap::from([("linear.weight".to_string(), weight)]),
    );
    assert!(matches!(
      Dense::load(dir.path(), &Device::Cpu).unwrap_err(),
      CanddingError::MissingTensor { name, .. } if name == "linear.bias"
    ));
  }
}