Refactoring of the codebase

[crangelsmith]

The codebase has gotten too big and some refactoring is needed. Some ideas to address:

• Data loading: The data loader has been extracted out into caked, in order to be able to load many more types of datasets, and from different pipelines. This simplifies what happens in data.py. Refactoring happening on Integrating Caked to Affinity #293

• Training and evaluation: Too many things are happening in these scripts that can be moved in to its own files, so these only show the pipeline for training (Training refactoring #303):
  - Load data
  - Define model
  - Define loss function
  - Run loop for training
  - Visualise and save results.

• Functions with too many parameters: train, eval, pass_batch, etc have too many parameters and it can get confusing and introduce bugs, we already have a a pedantic data model that can be passed around and in some cases, a class seems better than a function. Some functions can be broken into smaller parts as well.

• Many minor bug fixes that have been opened for a while.

Please add any more ideas for refactoring, I'll keep this as a meta issue for discussion.

Before this refactoring begins I think we should merge the current version of the code to main as it that has been used to produce results for the paper.