Small refactor & fixes for models