Refactor and clean the newly introduced code

* Use better method names, clean up the internal API.
* Call save_model from the parent when finished.

It is a good practice to always call the original method from
the overriden one, unless you have a very good reason not to.
This way it allows others to do their extensions as well.