Refactored code and modified interface.py to visualize training