How to train a cross-encoder reranker with Sentence Transformers

First-stage retrievers trade some ranking precision for the speed needed to search a large corpus. A CrossEncoder can learn application-specific relevance by reading each query and candidate together, making it suitable for the second ranking stage when generic model weights miss the language used in support records or internal documentation.

Current Sentence Transformers training separates the one-label model, labeled pair dataset, loss, training arguments, and trainer. BinaryCrossEntropyLoss accepts numeric positive and negative labels, while CrossEncoderTrainer handles batching and optimizer updates.

Four optimizer steps over eight labeled pairs make this a CPU smoke run rather than a production evaluation. The two ranked documents also occur in the training data, so their ordering demonstrates only that the saved model reloads and handles that known pair; generalization requires a separate validation split with unseen queries and passages.

1. Install Sentence Transformers with trainer dependencies in the active Python virtual environment.

   $ python3 -m pip install --upgrade "sentence-transformers[train]"

   The train extra installs the Hugging Face dataset and trainer packages required by CrossEncoderTrainer.
   Related: How to install Sentence Transformers with pip

2. Create train_support_reranker.py with the imports and labeled query-document pairs.

   train_support_reranker.py

   from pathlib import Path
    
   import torch
   from datasets import Dataset
   from sentence_transformers import CrossEncoder
   from sentence_transformers.cross_encoder import (
       CrossEncoderTrainer,
       CrossEncoderTrainingArguments,
   )
   from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
    
    
   train_dataset = Dataset.from_dict(
       {
           "query": [
               "reset a user password",
               "reset a user password",
               "enable two-factor authentication",
               "enable two-factor authentication",
               "export audit logs",
               "export audit logs",
               "restore a deleted project",
               "restore a deleted project",
           ],
           "document": [
               "Open the user profile and send a password recovery email.",
               "Open billing settings and update the saved payment card.",
               "Open account security and enroll an authenticator app.",
               "Open the release dashboard and deploy the web service.",
               "Open compliance reports and export the audit log CSV file.",
               "Open team settings and invite a new member.",
               "Open deleted projects and restore the selected project.",
               "Open API tokens and create a replacement token.",
           ],
           "label": [1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0],
       }
   )

   Each row represents one query-document pair. A mined passage that also answers the query is a false negative and must not receive a 0.0 label.

3. Add the one-label model, binary loss, and CPU training arguments below the dataset definition.

   model_dir = Path("models/support-reranker")
   model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", num_labels=1)
   loss = BinaryCrossEntropyLoss(model)
   args = CrossEncoderTrainingArguments(
       output_dir="training-output/support-reranker",
       max_steps=4,
       per_device_train_batch_size=2,
       learning_rate=2e-5,
       warmup_steps=0,
       use_cpu=True,
       eval_strategy="no",
       save_strategy="no",
       logging_strategy="no",
       report_to=[],
       disable_tqdm=True,
       seed=7,
   )

   A new model_dir prevents files from an earlier run from being replaced. Longer training runs need a validation split and evaluator.

4. Append the training block after the training arguments.

   trainer = CrossEncoderTrainer(
       model=model,
       args=args,
       train_dataset=train_dataset,
       loss=loss,
   )
   train_result = trainer.train()
   model.save_pretrained(model_dir)

5. Append the saved-model reload and training-document ranking checks after the save operation.

   reranker = CrossEncoder(str(model_dir), activation_fn=torch.nn.Sigmoid())
   query = "How do I reset an account password?"
   documents = [
       "Open the user profile and send a password recovery email.",
       "Open the release dashboard and deploy the web service.",
   ]
   scores = reranker.predict(
       [(query, document) for document in documents],
       convert_to_numpy=True,
       show_progress_bar=False,
   )
    
   if train_result.global_step != 4:
       raise SystemExit(f"expected 4 training steps, got {train_result.global_step}")
   if not (model_dir / "config.json").is_file():
       raise SystemExit("saved model configuration is missing")
   if float(scores[0]) <= float(scores[1]):
       raise SystemExit("relevant passage did not outrank the unrelated passage")
    
   print(f"Training rows: {len(train_dataset)}")
   print(f"Training steps: {train_result.global_step}")
   print(f"Saved model: {model_dir}")
   print(f"Relevant score: {float(scores[0]):.3f}")
   print(f"Unrelated score: {float(scores[1]):.3f}")
   print("Reload check: relevant passage ranked first")

6. Run the completed reranker training program from the project directory.

   $ python3 train_support_reranker.py
   ##### snipped #####
   Training rows: 8
   Training steps: 4
   Saved model: models/support-reranker
   Relevant score: 0.557
   Unrelated score: 0.000
   Reload check: relevant passage ranked first

   The omitted trainer metrics contain machine-dependent runtime values. A missing model file, incomplete run, or failure to rank the matching training document above the unrelated training document exits with an error instead of printing the reload confirmation.