little changes in main.py