Refactored code so the can add different models more easily.