Modeling#

class FlagEmbedding.finetune.reranker.decoder_only.layerwise.CrossDecoderModel(base_model: PreTrainedModel, tokenizer: AutoTokenizer = None, train_batch_size: int = 4, start_layer: int = 8)[source]#

Model class for decoder only reranker.

Parameters:

  • base_model (PreTrainedModel) – The underlying pre-trained model used for encoding and scoring input pairs.

  • tokenizer (AutoTokenizer, optional) – The tokenizer for encoding input text. Defaults to None.

  • train_batch_size (int, optional) – The batch size to use. Defaults to 4.

  • start_layer (int, optional) – Starting layer for layerwise. Defaults to 8.

Methods#

CrossDecoderModel.encode(features)[source]#

Abstract method of encode.

Parameters:

features (dict) – Teatures to pass to the model.

CrossDecoderModel.forward(pair: Dict[str, Tensor] | List[Dict[str, Tensor]] = None, teacher_scores: Tensor | None = None)[source]#

The computation performed at every call.

Parameters:

  • pair (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional) – The query-document pair. Defaults to None.

  • teacher_scores (Optional[Tensor], optional) – Teacher scores of knowledge distillation. Defaults to None.

Returns:

Output of reranker model.

Return type:

RerankerOutput