Integration with sklearn pipeline