Updating the code snippet in the readme and reformatting.