Modeling

class FlagEmbedding.finetune.reranker.encoder_only.base.CrossEncoderModel(base_model: PreTrainedModel, tokenizer: AutoTokenizer | None = None, train_batch_size: int = 4)

Model class for 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.

Methods

CrossEncoderModel.encode(features)

Encodes input features to logits.

Parameters:

features (dict) – Dictionary with input features.

Returns:

The logits output from the model.

Return type:

torch.Tensor