[FLAX] Minor fixes in CLM example
Hi,

this PR fixes some minor issues that I've seen when training a new GPT-2 model from scratch:

* It uses the correct method for retrieving the vocab size from tokenizer instance
* Fixes train and validation assignment of dataset instance when using train or validation files