When there is a trained model it is loaded directly, but if we need to train the code is written under the `else`. Refactor it to its own function.