modified code to run both on AWS and local machine