[Flax] Bugfixes in `run_mlm_flax.py`
# What does this PR do?

This PR fixes a few bugs I have observed when using `run_mlm_flax.py`:
- As discussed with @mfuntowicz , `jnp.split` is a lot slower than `np.split` on the first iteration, outright hanging in my tests on simplewiki (~20MB). As this operation doesn't need to be traced. we can use `np.split` instead.
- When using a HF `datasets`, the text column was also passed to the model as input, causing a bug. The PR removes the text column in `dataset.map` to avoid this.
- Finally, using `warmup_steps = 0` (as is default) causes the Flax optimizer to output NaNs. We use 1 as a minimum value for the same warmup-less behaviour.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.