Cleaning up model code