Refactor model implementation