Cleaning up model.py