This pull request includes several changes to the train function and related methods in the oumi module to support additional model and trainer keyword arguments. The motivation was to allow passing non-serializable keyword arguments into the trainer class (and the model class). The most important changes are as follows:

Enhancements to train function:

• Modified the train function in src/oumi/__init__.py to accept additional_model_kwargs and additional_trainer_kwargs parameters and pass them to the oumi.train.train function.

Enhancements to train function in src/oumi/train.py:

• Updated the _create_optional_training_kwargs function to accept an additional_trainer_kwargs parameter and include it in the returned dictionary. [1] [2]
• Modified the train function to accept additional_model_kwargs and additional_trainer_kwargs parameters and use them when building the model and creating optional training kwargs. [1] [2]

Related issues

Fixes #1623 by allowing user to pass a preprocess_logits_for_metrics function to additional_trainer_kwargs in the train function.

Before submitting

• This PR only changes documentation. (You can ignore the following checks in that case)
• Did you read the contributor guideline Pull Request guidelines?
• Did you link the issue(s) related to this PR in the section above?
• Did you add / update tests where needed?

Reviewers

At least one review from a member of oumi-ai/oumi-staff is required.