Cleaning up code for release