added some code to models.py