Skip to content

Commit

Permalink
Moved Model Card Callback init in Trainer to a separate function (#3047)
Browse files Browse the repository at this point in the history
* Moved Model Card Callback init in Trainer to a separate function for subclassing

* revert formatting

* Rename slightly; add docstring which'll go in docs

* Initialize default_args_dict before the super().__init__()

---------

Co-authored-by: Tom Aarsen <Cubiegamedev@gmail.com>
  • Loading branch information
tRosenflanz and tomaarsen authored Nov 11, 2024
1 parent 7ede83b commit 6347b75
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion sentence_transformers/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,8 +277,24 @@ def __init__(
self.eval_dataset = self.maybe_add_prompts_or_dataset_name_column(
eval_dataset, args.prompts, dataset_name="eval"
)
self.add_model_card_callback(default_args_dict)

def add_model_card_callback(self, default_args_dict: dict[str, Any]) -> None:
"""
Add a callback responsible for automatically tracking data required for the automatic model card generation
This method is called in the ``__init__`` method of the
:class:`~sentence_transformers.trainer.SentenceTransformerTrainer` class.
Args:
default_args_dict (Dict[str, Any]): A dictionary of the default training arguments, so we can determine
which arguments have been changed for the model card.
.. note::
This method can be overriden by subclassing the trainer to remove/customize this callback in custom uses cases
"""

# Add a callback responsible for automatically tracking data required for the automatic model card generation
model_card_callback = ModelCardCallback(self, default_args_dict)
self.add_callback(model_card_callback)
model_card_callback.on_init_end(self.args, self.state, self.control, self.model)
Expand Down

0 comments on commit 6347b75

Please sign in to comment.