Update models.py

Refactoring models