Modeling

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

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)

Abstract method of encode.

Parameters:

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

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

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