AbsTrainer

AbsRerankerTrainer

class FlagEmbedding.abc.finetune.reranker.AbsRerankerTrainer(model: PreTrainedModel | Module | None = None, args: TrainingArguments | None = None, data_collator: Callable[[list[Any]], dict[str, Any]] | None = None, train_dataset: Dataset | IterableDataset | Dataset | None = None, eval_dataset: Dataset | dict[str, Dataset] | Dataset | None = None, processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin | None = None, model_init: Callable[[...], PreTrainedModel] | None = None, compute_loss_func: Callable | None = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, optimizers: tuple[Optimizer | None, LambdaLR | None] = (None, None), optimizer_cls_and_kwargs: tuple[type[Optimizer], dict[str, Any]] | None = None, preprocess_logits_for_metrics: Callable[[Tensor, Tensor], Tensor] | None = None)

Abstract class for the trainer of reranker.

Methods

AbsRerankerTrainer.compute_loss(model, inputs, return_outputs=False, **kwargs)

How the loss is computed by Trainer. By default, all models return the loss in the first element.

Subclass and override for custom behavior.

Parameters:

• model (AbsRerankerModel) – The model being trained.

• inputs (dict) – A dictionary of input tensors to be passed to the model.

• return_outputs (bool, optional) – If True, returns both the loss and the model’s outputs. Otherwise, returns only the loss. Defaults to False.

Returns:

The computed loss. If return_outputs is True,

also returns the model’s outputs in a tuple (loss, outputs).

Return type:

Union[torch.Tensor, tuple(torch.Tensor, RerankerOutput)]