Made improvements to pull request as per comments