Updated main code to work with current library