MS MARCO Contrastive and/or Distillation sample training script

train_script.py

import argparse
import logging
import traceback
from collections import defaultdict
from collections.abc import Iterable
from enum import Enum, auto

import torch
from datasets import load_dataset
from torch import Tensor

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.evaluation import InformationRetrievalEvaluator, NanoBEIREvaluator, SequentialEvaluator
from sentence_transformers.losses import (
    CachedMultipleNegativesRankingLoss,
    DistillKLDivLoss,
    MarginMSELoss,
    MultipleNegativesRankingLoss,
)
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.util import pairwise_dot_score

# Set the log level to INFO to get more information
logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)


class LossType(Enum):
    MNRL = auto()
    CMNRL = auto()
    MARGIN_MSE = auto()
    KLDIV = auto()
    MARGIN_MSE_KLDIV = auto()

    def __str__(self):
        return self.name.lower()


class MarginMSEKLDivLoss(torch.nn.Module):
    def __init__(
        self,
        model: SentenceTransformer,
        similarity_fct=pairwise_dot_score,
        temperature=1.0,
        margin_mse_weight=1.0,
        kldiv_weight=1.0,
    ) -> None:
        super().__init__()
        self.model = model
        self.similarity_fct = similarity_fct
        self.temperature = temperature
        self.margin_mse_weight = margin_mse_weight
        self.kldiv_weight = kldiv_weight

        self.margin_mse_loss = MarginMSELoss(self.model, similarity_fct=self.similarity_fct)
        self.kl_div_loss = DistillKLDivLoss(
            self.model, similarity_fct=self.similarity_fct, temperature=self.temperature
        )

    def forward(self, sentence_features: Iterable[dict[str, Tensor]], labels: Tensor) -> Tensor:
        embeddings = [self.model(sentence_feature)["sentence_embedding"] for sentence_feature in sentence_features]

        return self.compute_loss_from_embeddings(embeddings, labels)

    def compute_loss_from_embeddings(self, embeddings: list[Tensor], labels: Tensor) -> Tensor:
        return {
            "margin_mse": self.margin_mse_loss.compute_loss_from_embeddings(embeddings, labels) * self.margin_mse_weight,
            "kl_div": self.kl_div_loss.compute_loss_from_embeddings(embeddings, labels) * self.kldiv_weight,
        }


def main(
    model_name_or_path: str,
    loss_type: LossType,
    kldiv_temperature: float,
    margin_mse_weight: float,
    kldiv_weight: float,
    mini_batch_size: int,
    mnrl_scale: float,
    num_train_epochs: int,
    per_device_batch_size: int,
    learning_rate: float,
    warmup_ratio: float,
    fp16: bool,
    bf16: bool,
    eval_save_steps: int,
    save_total_limit: int,
    logging_steps: int,
    evaluator_batch_size: int,
    quick: bool,
):
    # 1. Load a model with prompts to finetune with 2. (Optional) model card data
    model = SentenceTransformer(
        model_name_or_path,
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{model_name_or_path} trained on RLHN MS MARCO using {loss_type}",
        ),
        prompts={  # prompts with "query" and "document" keys are automatically used in evaluation via model.encode_query/model.encode_document
            "query": "query: ",
            "document": "document: ",
        },
    )

    # 3. Load a dataset to finetune on
    # TODO: Eventually we want this:
    """
    train_dataset = load_dataset("mixedbread-ai/rlhn-680k-msmarco-7negs", split="train")
    eval_dataset = load_dataset("mixedbread-ai/rlhn-680k-msmarco-7negs", split="eval")
    test_dataset = load_dataset("mixedbread-ai/rlhn-680k-msmarco-7negs", split="test")
    train_dataset = train_dataset.select_columns([column for column in train_dataset.column_names if column != 'logits'])
    eval_dataset = eval_dataset.select_columns([column for column in eval_dataset.column_names if column != 'logits'])
    test_dataset = test_dataset.select_columns([column for column in test_dataset.column_names if column != 'logits'])
    """
    # But for now we do it manually:
    dataset = load_dataset("mixedbread-ai/rlhn-680k-msmarco-7negs-scored", split="train")
    dataset = dataset.select_columns([column for column in dataset.column_names if column != "logits"])
    split_dataset = dataset.train_test_split(test_size=3_000)
    dataset = split_dataset["train"]
    eval_dataset = split_dataset["test"]
    split_dataset = dataset.train_test_split(test_size=10_000)
    train_dataset = split_dataset["train"]
    test_dataset = split_dataset["test"]

    # 4. Define a loss function
    batch_sampler = BatchSamplers.BATCH_SAMPLER
    gather_across_devices = torch.distributed.is_initialized() if torch.distributed.is_available() else False
    if loss_type == LossType.MNRL:
        loss = MultipleNegativesRankingLoss(model, scale=mnrl_scale, gather_across_devices=gather_across_devices)
        batch_sampler = BatchSamplers.NO_DUPLICATES
    elif loss_type == LossType.CMNRL:
        loss = CachedMultipleNegativesRankingLoss(
            model, scale=mnrl_scale, mini_batch_size=mini_batch_size, gather_across_devices=gather_across_devices
        )
        batch_sampler = BatchSamplers.NO_DUPLICATES
    elif loss_type == LossType.MARGIN_MSE:
        loss = MarginMSELoss(model)
    elif loss_type == LossType.KLDIV:
        loss = DistillKLDivLoss(model, temperature=kldiv_temperature)
    elif loss_type == LossType.MARGIN_MSE_KLDIV:
        loss = MarginMSEKLDivLoss(
            model, temperature=kldiv_temperature, margin_mse_weight=margin_mse_weight, kldiv_weight=kldiv_weight
        )

    # 5. (Optional) Specify training arguments
    short_model_name_or_path = model_name_or_path.split("/")[-1]
    run_name = f"{short_model_name_or_path}-{loss_type}-lr{learning_rate}-bs{per_device_batch_size}"
    column_names_to_prompts = {
        column_name: "query" if column_name == "query" else "document"
        for column_name in dataset.column_names
        if column_name != "scores"
    }
    args = SentenceTransformerTrainingArguments(
        # Required parameter:
        output_dir=f"models/{run_name}",
        # Optional training parameters:
        num_train_epochs=0.05 if quick else num_train_epochs,
        per_device_train_batch_size=per_device_batch_size,
        per_device_eval_batch_size=per_device_batch_size,
        learning_rate=learning_rate,
        warmup_ratio=warmup_ratio,
        fp16=fp16,  # Set to False if you get an error that your GPU can't run on FP16
        bf16=bf16,  # Set to True if you have a GPU that supports BF16
        batch_sampler=batch_sampler,  # (C)MNRL benefits from no duplicate samples in a batch
        prompts=column_names_to_prompts,  # Let's incorporate prompts for a ~1% improvement
        # Optional tracking/debugging parameters:
        eval_strategy="steps",
        eval_steps=eval_save_steps,
        save_strategy="steps",
        save_steps=eval_save_steps,
        save_total_limit=save_total_limit,
        logging_steps=logging_steps,
        run_name=run_name,
    )

    # 6. (Optional) Create evaluator & evaluate the base model
    nano_beir_evaluator = NanoBEIREvaluator(
        dataset_names=["msmarco", "nfcorpus", "nq"],
        batch_size=evaluator_batch_size,
        query_prompts=model.prompts["query"],  # This will be done automatically starting from the next version
        corpus_prompts=model.prompts["document"],  # This will be done automatically starting from the next version
    )
    eval_queries = {}
    eval_documents = {}
    eval_relevant_docs = defaultdict(set)
    for query, positive in zip(eval_dataset["query"], eval_dataset["positive"]):
        query_id = len(eval_queries)
        eval_queries[query_id] = query
        document_id = len(eval_documents)
        eval_documents[document_id] = positive
        eval_relevant_docs[query_id].add(document_id)
    for column_name in test_dataset.column_names:
        if column_name.startswith("negative"):
            for negative in test_dataset[column_name]:
                document_id = len(eval_documents)
                eval_documents[document_id] = negative
    eval_ir_evaluator = InformationRetrievalEvaluator(
        queries=eval_queries,
        corpus=eval_documents,
        relevant_docs=eval_relevant_docs,
        name="rlhn-msmarco-eval",
        batch_size=evaluator_batch_size,
        query_prompt_name="query",  # This will be done automatically starting from the next version
        corpus_prompt_name="document",  # This will be done automatically starting from the next version
    )
    eval_evaluator = SequentialEvaluator([nano_beir_evaluator, eval_ir_evaluator])
    if not quick:
        eval_evaluator(model)

    test_queries = {}
    test_documents = {}
    test_relevant_docs = defaultdict(set)
    for query, positive in zip(test_dataset["query"], test_dataset["positive"]):
        query_id = len(test_queries)
        test_queries[query_id] = query
        document_id = len(test_documents)
        test_documents[document_id] = positive
        test_relevant_docs[query_id].add(document_id)
    for column_name in test_dataset.column_names:
        if column_name.startswith("negative"):
            for negative in test_dataset[column_name]:
                document_id = len(test_documents)
                test_documents[document_id] = negative
    test_ir_evaluator = InformationRetrievalEvaluator(
        queries=test_queries,
        corpus=test_documents,
        relevant_docs=test_relevant_docs,
        name="rlhn-msmarco-test",
        batch_size=evaluator_batch_size,
        query_prompt_name="query",  # This will be done automatically starting from the next version
        corpus_prompt_name="document",  # This will be done automatically starting from the next version
    )
    test_evaluator = SequentialEvaluator([test_ir_evaluator])
    if not quick:
        test_evaluator(model)

    # 7. Create a trainer & train
    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=eval_evaluator,
    )
    trainer.train()

    # (Optional) Evaluate the trained model on the eval & test sets again
    eval_evaluator(model)
    test_evaluator(model)

    # 8. Save the final model
    final_output_dir = f"models/{run_name}/final"
    model.save_pretrained(final_output_dir)

    # 9. (Optional) save the model to the Hugging Face Hub!
    # It is recommended to run `huggingface-cli login` to log into your Hugging Face account first
    try:
        model.push_to_hub(run_name, private=True)
    except Exception:
        logging.error(
            f"Error uploading model to the Hugging Face Hub:\n{traceback.format_exc()}To upload it manually, you can run "
            f"`huggingface-cli login`, followed by loading the model using `model = CrossEncoder({final_output_dir!r})` "
            f"and saving it using `model.push_to_hub('{run_name}')`."
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a sentence transformer model on RLHN MS MARCO dataset")
    parser.add_argument(
        "--model_name_or_path", type=str, default="jhu-clsp/ettin-encoder-17m", help="Model name or path to load"
    )
    parser.add_argument(
        "--loss_type",
        type=lambda x: LossType[x.upper()],
        default=LossType.CMNRL,
        choices=list(LossType),
        help="Loss function to use",
    )
    parser.add_argument("--kldiv_temperature", type=float, default=1.0, help="Temperature for KL divergence loss")
    parser.add_argument("--margin_mse_weight", type=float, default=1.0, help="Weight for margin MSE in combined loss")
    parser.add_argument("--kldiv_weight", type=float, default=1.0, help="Weight for KL divergence in combined loss")
    parser.add_argument("--mini_batch_size", type=int, default=16, help="Mini-batch size for cached MNRL")
    parser.add_argument("--mnrl_scale", type=float, default=20.0, help="Scale factor for MNRL loss")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--per_device_batch_size", type=int, default=128, help="Batch size per device")
    parser.add_argument("--evaluator_batch_size", type=int, default=32, help="Batch size for the evaluators")
    parser.add_argument("--learning_rate", type=float, default=2e-5, help="Learning rate")
    parser.add_argument("--warmup_ratio", type=float, default=0.1, help="Ratio of warmup steps")
    parser.add_argument("--fp16", action="store_true", help="Use FP16 precision")
    parser.add_argument("--bf16", action="store_true", default=True, help="Use BF16 precision")
    parser.add_argument(
        "--eval_save_steps",
        type=float,
        default=0.2,
        help="Steps between evaluations and checkpoint saves. If less than 1, "
        "it will be treated as a fraction of the total steps.",
    )
    parser.add_argument("--save_total_limit", type=int, default=3, help="Maximum number of checkpoints to keep")
    parser.add_argument("--logging_steps", type=int, default=100, help="Steps between logging")
    parser.add_argument("--quick", action="store_true", help="Run with only 5% of training data for quick testing")

    args = parser.parse_args()

    main(
        model_name_or_path=args.model_name_or_path,
        loss_type=args.loss_type,
        kldiv_temperature=args.kldiv_temperature,
        margin_mse_weight=args.margin_mse_weight,
        kldiv_weight=args.kldiv_weight,
        mini_batch_size=args.mini_batch_size,
        mnrl_scale=args.mnrl_scale,
        num_train_epochs=args.num_train_epochs,
        per_device_batch_size=args.per_device_batch_size,
        learning_rate=args.learning_rate,
        warmup_ratio=args.warmup_ratio,
        fp16=args.fp16,
        bf16=args.bf16,
        eval_save_steps=args.eval_save_steps,
        save_total_limit=args.save_total_limit,
        logging_steps=args.logging_steps,
        evaluator_batch_size=args.evaluator_batch_size,
        quick=args.quick,
    )