Updated to work with library