Modified code to give user option to save to a file