Add TrainerArguments and adjust training loop

Add Trainer configuration as a separate class. This holds all the
training options as a separate dataclass; this can also be easily passed
in as a json file.

More code organization in the main training loop.