Slight model refactor