Improve dataset generation