Updated to work with new file structure