AbsModeling#

AbsEmbedderModel#

class FlagEmbedding.abc.finetune.embedder.AbsEmbedderModel(base_model, tokenizer: AutoTokenizer | None = None, negatives_cross_device: bool = False, temperature: float = 1.0, sub_batch_size: int = -1, kd_loss_type: str = 'kl_div')[source]#

Abstract class of embedding model for training.

Parameters:

base_model – 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 "kl_div".

Methods#

abstract AbsEmbedderModel.encode(features)[source]#

Abstract method encode and get the embedding.

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

abstract AbsEmbedderModel.compute_loss(scores, target)[source]#

Abstract method compute the loss.

Parameters:

scores (torch.Tensor) – Computed score.
target (torch.Tensor) – The target value.

abstract AbsEmbedderModel.compute_score(q_reps, p_reps)[source]#

Abstract method to compute the score.

Parameters:

q_reps (torch.Tensor) – Queries representations.
p_reps (torch.Tensor) – Passages rerpresentations.

abstract AbsEmbedderModel.save(output_dir: str)[source]#

Abstract method to save the model.

Parameters:: output_dir (str) – Directory for saving the model.

AbsEmbedderModel.get_local_score(q_reps, p_reps, all_scores)[source]#

Get the local score of queries and passages.

Parameters:

q_reps (torch.Tensor) – Queries representations.
p_reps (torch.Tensor) – Passages rerpresentations.
all_scores (torch.Tensor) – All the query-passage scores computed.

Returns:

Local scores to compute loss.

Return type:

torch.Tensor

AbsEmbedderModel.compute_local_score(q_reps, p_reps, compute_score_func=None, **kwargs)[source]#

Compute the local score of queries and passages.

Parameters:

q_reps (torch.Tensor) – Queries representations.
p_reps (torch.Tensor) – Passages rerpresentations.
compute_score_func (function, optional) – Function to compute score. Defaults to None, which will use the self.compute_score().

Returns:

Local scores to compute loss.

Return type:

torch.Tensor

AbsEmbedderModel.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)[source]#

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

static AbsEmbedderModel.distill_loss(kd_loss_type, teacher_targets, student_scores, group_size=None)[source]#

Compute the distillation loss.

Parameters:

kd_loss_type (str) – Type of knowledge distillation loss, supports “kl_div” and “m3_kd_loss”.
teacher_targets (torch.Tensor) – Targets from the teacher model.
student_scores (torch.Tensor) – Score of student model.
group_size (int, optional) – Number of groups for . Defaults to None.

Raises:

ValueError – Invalid kd_loss_type

Returns:

A scalar of computed distillation loss.

Return type:

torch.Tensor

AbsEmbedderModel._compute_no_in_batch_neg_loss(q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs)[source]#: Compute loss when using no in-batch negatives and no cross-device negatives

AbsEmbedderModel._compute_in_batch_neg_loss(q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs)[source]#: Compute loss when only using in-batch negatives

AbsEmbedderModel._compute_cross_device_neg_loss(q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs)[source]#: Compute loss when using both in-batch negatives and cross-device negatives

AbsEmbedderModel._dist_gather_tensor(t: Tensor | None)[source]#

Gather a tensor from all processes in a distributed setting.

Parameters:

t (Optional[torch.Tensor]) – The input tensor to be gathered. If None, no gathering is performed.

Returns:

A concatenated tensor from all processes if t is not None,: otherwise returns None.

Return type:

Union[torch.Tensor, None]

EmbedderOutput#

class FlagEmbedding.abc.finetune.embedder.EmbedderOutput(q_reps: Tensor | None = None, p_reps: Tensor | None = None, loss: Tensor | None = None, scores: Tensor | None = None)[source]#: Output information returned by the model.