Struct ColBERT 

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

The main ColBERT model structure.

This struct encapsulates the language model, the chain of Dense projection layers declared in modules.json, 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_modules: Vec<DenseModuleData>, tokenizer_bytes: Vec<u8>, config_bytes: Vec<u8>, query_prefix: String, document_prefix: String, query_prompt: String, document_prompt: String, mask_token: String, do_query_expansion: bool, 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.

dense_modules carries the ordered list of Dense projection layers declared in modules.json. They are applied left-to-right after the transformer, with the last layer’s out_features as the final stored embedding dimension.

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

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

pub fn document_token_lengths( &mut self, sentences: &[String], ) -> Result<Vec<u32>, ColbertError>

Compute the post-truncation token count for each document sentence.

Mirrors the prefix + prompt + tokenization pipeline used inside [encode] for documents so the returned lengths match the per-row valid-token count of the encoded tensor exactly. Used by callers that need to slice the padded output tensor without scanning for all-zero rows on the host.

pub fn encode_documents_with_lengths( &mut self, sentences: &[String], ) -> Result<(Tensor, Vec<u32>), ColbertError>

Encode documents and return (Tensor, per_doc_valid_token_counts).

The lengths are returned in the caller-supplied order and index directly into the returned 3D tensor’s axis=1: row i has the first lengths[i] rows populated with L2-normalized embeddings, and the remaining rows (up to the batch’s padded length) zeroed by [finalize_embeddings].

Intended for the docbert indexing path, which slices the tensor per-doc before serializing embeddings. Returning real counts lets callers skip the previous O(padded_tokens · dim) per-doc all-zero scan.

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, this method leverages Rayon for parallel batch processing to accelerate encoding. On accelerators (GPU), 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.