Struct ColBERT

pub struct ColBERT {
    pub device: Device,
    /* private fields */
}

The main ColBERT model structure.

This struct encapsulates the language model, a linear projection layer, the tokenizer, and all necessary configuration for performing encoding and similarity calculations based on the ColBERT architecture.

Fields

device: Device

The device (CPU or GPU) on which the model is loaded.

Implementations

impl ColBERT

pub fn new( weights: Vec<u8>, dense_weights: Vec<u8>, tokenizer_bytes: Vec<u8>, config_bytes: Vec<u8>, dense_config_bytes: Vec<u8>, query_prefix: String, document_prefix: String, mask_token: String, attend_to_expansion_tokens: bool, query_length: Option<usize>, document_length: Option<usize>, batch_size: Option<usize>, device: &Device, ) -> Result<Self, ColbertError>

Creates a new instance of the ColBERT model from byte buffers.

pub fn from(repo_id: &str) -> ColbertBuilder

Creates a ColbertBuilder to construct a ColBERT model from a Hugging Face repository.

pub fn encode( &mut self, sentences: &[String], is_query: bool, ) -> Result<Tensor, ColbertError>

Encodes a batch of sentences (queries or documents) into embeddings.

On CPU and non-WASM targets, this method leverages Rayon for parallel batch processing to accelerate encoding. On other targets (like GPU or WASM), it processes batches sequentially.

pub fn similarity( &self, queries_embeddings: &Tensor, documents_embeddings: &Tensor, ) -> Result<Similarities, ColbertError>

Calculates the similarity scores between query and document embeddings.

pub fn raw_similarity( &self, queries_embeddings: &Tensor, documents_embeddings: &Tensor, ) -> Result<Tensor, ColbertError>

Computes the raw, un-reduced similarity matrix between query and document embeddings.

Trait Implementations

impl TryFrom<ColbertBuilder> for ColBERT

fn try_from(builder: ColbertBuilder) -> Result<Self, Self::Error>

Builds the ColBERT model by downloading files from the hub and initializing the model.

type Error = ColbertError

The type returned in the event of a conversion error.