examples/research_projects/rag-end2end-retriever/callbacks_rag.py/get_checkpoint_callback
def get_checkpoint_callback(output_dir, metric):
    """Saves the best model by validation EM score."""
    if metric == "rouge2":
        exp = "{val_avg_rouge2:.4f}-{step_count}"
    elif metric == "bleu":
        exp = "{val_avg_bleu:.4f}-{step_count}"
    elif metric == "em":
        exp = "{val_avg_em:.4f}-{step_count}"
    elif metric == "loss":
        exp = "{val_avg_loss:.4f}-{step_count}"
    else:
        raise NotImplementedError(
            f"seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this function."
        )

    checkpoint_callback = ModelCheckpoint(
        dirpath=output_dir,
        filename=exp,
        monitor=f"val_{metric}",
        mode="max",
        save_top_k=1,
        every_n_val_epochs=1,  # works only with PL > 1.3
    )

    return checkpoint_callback