docbert-pylate 0.8.0

Rust library for late interaction (ColBERT) models, vendored into the docbert workspace.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12

use candle_core::Tensor;

use crate::error::ColbertError;

/// Normalizes a tensor using L2 normalization along the last dimension.
///
/// A small epsilon is applied to avoid NaNs when a row is entirely zero.
pub fn normalize_l2(v: &Tensor) -> Result<Tensor, ColbertError> {
    let inv_norm_l2 =
        ((v.sqr()?.sum_keepdim(v.rank() - 1)?.sqrt()? + 1e-12f64)?).recip()?;
    v.broadcast_mul(&inv_norm_l2).map_err(ColbertError::from)
}