Still refine the model class