candding 0.1.0

use candle_core::DType;
use candle_core::Device;
use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Embedding;
use candle_nn::VarBuilder;
use candle_nn::embedding;

use super::config::MPNetConfig;
use crate::error::Result;

/// The reference always buckets with this default, whatever the config says; `ensure_supported` refuses others.
pub(super) const NUM_BUCKETS: usize = 32;

/// Scales the logarithm rather than the cut: the `min` against `half - 1` bites first, at distance 91.
const MAX_DISTANCE: i64 = 128;

/// `MPNetEncoder.relative_position_bucket` for `rel = j - i`: 16 past and 16 future buckets, exact for the
/// eight nearest distances, logarithmic beyond, saturating at 91.
fn relative_position_bucket(rel: i64) -> u32 {
  let n = -rel;
  let half = (NUM_BUCKETS / 2) as i64;
  let past_or_future = if n < 0 { half } else { 0 };
  let n = n.abs();
  let max_exact = half / 2;
  let offset = if n < max_exact {
    n
  } else {
    let scaled = (n as f64 / max_exact as f64).ln() / (MAX_DISTANCE as f64 / max_exact as f64).ln();
    (max_exact + (scaled * (half - max_exact) as f64) as i64).min(half - 1)
  };
  (past_or_future + offset) as u32
}

/// `encoder.relative_attention_bias.weight`, looked up by bucket into a per-head additive term.
pub struct RelativeAttentionBias {
  table: Embedding,
}

impl RelativeAttentionBias {
  /// Loaded in F32: the bias composes with the padding floor in F32, with one cast at the end.
  pub fn load(config: &MPNetConfig, vb: VarBuilder) -> Result<Self> {
    Ok(Self {
      table: embedding(
        config.relative_attention_num_buckets,
        config.num_attention_heads,
        vb.set_dtype(DType::F32),
      )?,
    })
  }

  /// `(1, heads, seq, seq)` in F32, bucketed from plain zero-based positions as the reference does (not the
  /// embeddings' offset ids); computed once per forward and shared by every layer and row.
  pub fn forward(&self, seq: usize, device: &Device) -> Result<Tensor> {
    let mut buckets = Vec::with_capacity(seq * seq);
    for i in 0..seq {
      for j in 0..seq {
        buckets.push(relative_position_bucket(j as i64 - i as i64));
      }
    }
    let indices = Tensor::from_vec(buckets, (seq, seq), device)?;
    Ok(
      self
        .table
        .forward(&indices)?
        .permute((2, 0, 1))?
        .unsqueeze(0)?
        .contiguous()?,
    )
  }
}

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

  use super::*;

  /// `(first_distance, offset)` pairs read off `relative_position_bucket` for distances 0 to 400, past the catalog's 384 tokens.
  const OFFSETS: &[(i64, i64)] = &[
    (0, 0),
    (1, 1),
    (2, 2),
    (3, 3),
    (4, 4),
    (5, 5),
    (6, 6),
    (7, 7),
    (8, 8),
    (12, 9),
    (16, 10),
    (23, 11),
    (32, 12),
    (46, 13),
    (64, 14),
    (91, 15),
  ];

  fn expected_offset(distance: i64) -> i64 {
    OFFSETS
      .iter()
      .rev()
      .find(|(first, _)| distance >= *first)
      .expect("a non-negative distance")
      .1
  }

  /// The past half starts at bucket 0, the future half at 16, and a key right of its query lands in the future half.
  #[test]
  fn every_distance_lands_in_the_reference_bucket() {
    for rel in -400..=400i64 {
      let expected = if rel > 0 {
        (NUM_BUCKETS as i64 / 2) + expected_offset(rel)
      } else {
        expected_offset(-rel)
      };
      assert_eq!(relative_position_bucket(rel) as i64, expected, "rel {rel}");
    }
  }

  /// Bucket 16 is unreachable, and the checkpoint agrees: its row is the only one within 0.006 of zero.
  /// A backwards sign convention would put the diagonal in that untrained row.
  #[test]
  fn bucket_sixteen_is_never_reached() {
    let used: BTreeSet<u32> = (-400..=400).map(relative_position_bucket).collect();
    assert!(!used.contains(&16), "the future half starts at 17");
    assert_eq!(used.len(), NUM_BUCKETS - 1);
    assert_eq!(used.first().copied(), Some(0));
    assert_eq!(used.last().copied(), Some(NUM_BUCKETS as u32 - 1));
  }

  /// A transposed bias or a `context - memory` subtraction inverts the sign, costing cosine 0.5371.
  #[test]
  fn the_diagonal_and_its_neighbours_pin_the_sign_convention() {
    assert_eq!(relative_position_bucket(0), 0);
    assert_eq!(relative_position_bucket(-1), 1);
    assert_eq!(relative_position_bucket(1), 17);
    assert_eq!(relative_position_bucket(-8), 8);
    assert_eq!(relative_position_bucket(8), 24);
  }

  /// The `min` against `half - 1` bites at distance 91, before `MAX_DISTANCE`, as in the reference.
  #[test]
  fn distant_keys_share_one_bucket_per_direction() {
    for rel in [91, 128, 200, 384, 400] {
      assert_eq!(relative_position_bucket(rel), 31, "rel {rel}");
      assert_eq!(relative_position_bucket(-rel), 15, "rel -{rel}");
    }
    assert_ne!(relative_position_bucket(90), relative_position_bucket(91));
  }
}