Clean-up

* Moves module demo code to main block

* Improves code style