Modified the Entity model, cleaned up some code