Updated the code for the ui changes