Refactor model code to be extensible