Modeling

EncoderOnlyEmbedderM3Model

class FlagEmbedding.finetune.embedder.encoder_only.m3.EncoderOnlyEmbedderM3Model(base_model: Dict[str, Any], tokenizer: AutoTokenizer | None = None, negatives_cross_device: bool = False, temperature: float = 1, sub_batch_size: int = -1, kd_loss_type: str = 'm3_kd_loss', sentence_pooling_method: str = 'cls', normalize_embeddings: bool = False, unified_finetuning: bool = True, use_self_distill: bool = False, self_distill_start_step: int = -1)

Embedder class for M3 model.

Parameters:

• base_model (AutoModel) – The base model to train on.

• tokenizer (AutoTokenizer, optional) – The tokenizer to use. Defaults to None.

• negatives_cross_device (bool, optional) – If True, will compute cross devices negative loss. Defaults to False.

• temperature (float, optional) – Temperature to control the scale of scores. Defaults to 1.0.

• sub_batch_size (int, optional) – Sub-batch size during encoding. If negative, will not split to sub-batch. Defaults to -1.

• kd_loss_type (str, optional) – Type of knowledge distillation loss. Defaults to 'm3_kd_loss'.

• sentence_pooling_method (str, optional) – Pooling method to get sentence embedding. Defaults to 'cls'.

• normalize_embeddings (bool, optional) – If True, normalize the embedding vector. Defaults to False.

• unified_finetuning (bool, optional) – If True, will finetune colbert vector and sparce embedding. Defaults to True.

• use_self_distill (bool, optional) – If True, will do self distillation. Defaults to False.

• self_distill_start_step (int, optional) – Step num to start self distillation. Defaults to -1.

Methods

EncoderOnlyEmbedderM3Model.encode(features)

Encode and get the embedding.

Parameters:

features (Union[list, dict]) – Features feed to the model.

Returns:

Dense embeddings. torch.Tensor: Sparce embeddings. torch.Tensor: Colbert vectors.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model.compute_score(q_reps, p_reps, q_mask: Tensor, dense_weight: float = 1.0, sparse_weight: float = 0.3, colbert_weight: float = 1.0)

_summary_

Parameters:

• q_reps (_type_) – Query representations.

• p_reps (_type_) – Passage representations.

• q_mask (torch.Tensor) – _description_

• dense_weight (float, optional) – _description_. Defaults to 1.0.

• sparse_weight (float, optional) – _description_. Defaults to 0.3.

• colbert_weight (float, optional) – _description_. Defaults to 1.0.

Returns:

_description_

Return type:

_type_

EncoderOnlyEmbedderM3Model.compute_dense_score(q_reps, p_reps)

Compute the dense score.

Parameters:

• q_reps (torch.Tensor) – Query representations.

• p_reps (torch.Tensor) – Passage representations.

Returns:

The computed dense scores, adjusted by temperature.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model.compute_sparse_score(q_reps, p_reps)

Compute the sparse score.

Parameters:

• q_reps (torch.Tensor) – Query representations.

• p_reps (torch.Tensor) – Passage representations.

Returns:

The computed sparse scores, adjusted by temperature.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model.compute_colbert_score(q_reps, p_reps, q_mask: Tensor | None = None)

Compute the colbert score.

Parameters:

• q_reps (torch.Tensor) – Query representations.

• p_reps (torch.Tensor) – Passage representations.

Returns:

The computed colber scores, adjusted by temperature.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model.ensemble_score(q_reps, p_reps, dense_scores=None, sparse_scores=None, colbert_scores=None)

Compute the ensemble score of the three methods.

Parameters:

• q_reps (torch.Tensor) – Query representations.

• p_reps (torch.Tensor) – Passage representations.

• dense_scores (torch.Tensor, optional) – The dense scores. Defaults to None.

• sparse_scores (torch.Tensor, optional) – The sparse scores. Defaults to None.

• colbert_scores (torch.Tensor, optional) – The colbert scores. Defaults to None.

Raises:

ValueError – dense_scores, sparse_scores, colbert_scores must be provided

Returns:

The ensemble score of the three methods.

Return type:

_type_

EncoderOnlyEmbedderM3Model.forward(queries: Dict[str, Tensor] | List[Dict[str, Tensor]] | None = None, passages: Dict[str, Tensor] | List[Dict[str, Tensor]] | None = None, teacher_scores: None | List[float] = None, no_in_batch_neg_flag: bool = False)

The computation performed at every call.

Parameters:

• queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional) – Input queries. Defaults to None.

• passages (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional) – Input passages. Defaults to None.

• teacher_scores (Union[None, List[float]], optional) – Teacher scores for distillation. Defaults to None.

• no_in_batch_neg_flag (bool, optional) – If True, use no in-batch negatives and no cross-device negatives. Defaults to False.

Returns:

Output of the forward call of model.

Return type:

EmbedderOutput

EncoderOnlyEmbedderM3Model.compute_loss(scores, target)

Compute the loss using cross entropy.

Parameters:

• scores (torch.Tensor) – Computed score.

• target (torch.Tensor) – The target value.

Returns:

The computed cross entropy loss.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model.gradient_checkpointing_enable(**kwargs)

Activates gradient checkpointing for the current model.

EncoderOnlyEmbedderM3Model.enable_input_require_grads(**kwargs)

Enables the gradients for the input embeddings.

EncoderOnlyEmbedderM3Model.save(output_dir: str)

Save the model to the directory.

Parameters:

output_dir (str) – Directory for saving the model.

EncoderOnlyEmbedderM3Model._dense_embedding(last_hidden_state, attention_mask)

Use the pooling method to get the dense embedding.

Parameters:

• last_hidden_state (torch.Tensor) – The model output’s last hidden state.

• attention_mask (torch.Tensor) – Mask out padding tokens during pooling.

Raises:

NotImplementedError – Specified pooling method not implemented.

Returns:

The dense embeddings.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model._sparse_embedding(hidden_state, input_ids, return_embedding: bool = True)

Compute and return the sparse embedding.

Parameters:

• hidden_state (torch.Tensor) – The model output’s last hidden state.

• input_ids (_type_) – Ids from input features.

• return_embedding (bool, optional) – If True, return the computed embedding, otherwise just return the token weights. Defaults to True.

Returns:

The sparse embedding or just the token weights.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model._colbert_embedding(last_hidden_state, mask)

Get the colbert vectors.

Parameters:

• last_hidden_state (torch.Tensor) – The model output’s last hidden state.

• attention_mask (torch.Tensor) – Mask out padding tokens during pooling.

Returns:

The colbert vectors.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model._encode(features)

Helper function to encode using input features.

Parameters:

features (Union[list, dict]) – Features feed to the model.

Returns:

Dense embedding. torch.Tensor: Sparce embedding. torch.Tensor: Colbert vector.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model._compute_similarity(q_reps, p_reps)

Computes the similarity between query and passage representations using inner product.

Parameters:

• q_reps (torch.Tensor) – Query representations.

• p_reps (torch.Tensor) – Passage representations.

Returns:

The computed similarity matrix.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3Model._get_queries_attention_mask(queries: Dict[str, Tensor] | List[Dict[str, Tensor]])

padding attention mask for colbert

Parameters:

queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]]) – Input queries.

Returns:

The query attention mask.

Return type:

torch.Tensor

EncoderOnlyEmbedderM3ModelForInference

class FlagEmbedding.finetune.embedder.encoder_only.m3.EncoderOnlyEmbedderM3ModelForInference(base_model: Dict[str, Any], tokenizer: AutoTokenizer | None = None, negatives_cross_device: bool = False, temperature: float = 1, sub_batch_size: int = -1, kd_loss_type: str = 'm3_kd_loss', sentence_pooling_method: str = 'cls', normalize_embeddings: bool = False, unified_finetuning: bool = True, use_self_distill: bool = False, self_distill_start_step: int = -1)

Inference class of M3 model.

forward(text_input: Dict[str, Tensor] | None = None, return_dense: bool = True, return_sparse: bool = False, return_colbert_vecs: bool = False, return_sparse_embedding: bool = False)

Encode the text input using the selected way.

Parameters:

• text_input (Dict[str, Tensor], optional) – Text inputs. Defaults to None.

• return_dense (bool, optional) – If True, return the dense embedding. Defaults to True.

• return_sparse (bool, optional) – If True, return the sparse embedding. Defaults to False.

• return_colbert_vecs (bool, optional) – If True, return the colbert vectors. Defaults to False.

• return_sparse_embedding (bool, optional) – Parameter for _sparse_embedding(). If True, will return sparse embedding. Otherwise, return the token weights. Defaults to False.

Returns:

A dictionary containing the three types of embeddings.

Return type:

dict