Refer to commit https://github.com/dsfsi/textaugment/commit/f2c05ec32fc79a5dbcc81519ebb291cfaf6c0ce2 @idzingirai please review