Improve the main function
Clean a bit
mv benchmark.py and improve it