Update dataset.py to be compatible with TF 2.x
Hi @LukeBoyer Can you please review this PR ? Thank you!
 looks good to me
The functions in the `dataset.py` are outdated and deprecated in Tensorflow 2.x. This commit updates the code to be compatible with the latest TF versions by replacing the deprecated functions with their TensorFlow 2.x equivalents. 

Thanks.