[Flax] Fix eval and data_args usage in streaming example

[falcaopetri]

What does this PR do?

This PR fixes the evaluation loop in run_mlm_flax_stream.py. Current behavior didn't update the correct variable, which leads to data leakage during evaluation.

It also takes the opportunity to improve some DataTrainingArguments usages.

---

It's a draft PR because there is an open improvement that could be made: the script splits train-eval based solely in data_args.{dataset_name,num_eval_samples}, but also accepts unused args train_file, validation_file, train_ref_file, validation_ref_file, validation_split_percentage. Other data args that are unused: pad_to_max_length, line_by_line.

My suggestion would be to remove all these unused args. May I proceed with that?

Before submitting

• Did you read the contributor guideline, Request section?
• Did you make sure to update the documentation with your changes?
  The script is mentioned in jax-projects/dataset-streaming/README, but no changes are required.

Who can review?

@patrickvonplaten

[github-actions]

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 are likely to be ignored.

[patrickvonplaten]

@patil-suraj could you take a look here? :-)