Update models.py

Finish the data model