Clean up model api

[SamTov]

I didn't realise through all of the PRs but the internal models classes are very convoluted now and it is hard to make changes as the steps are distributed over many modules. I would like to clean all of this up as currently making model training changes is difficult if it requires changing things like the apply function.