Refactor Model class
Cleaned up imports