[Flax] Refactor gpt2 & bert example docs
# What does this PR do?

This PR mainly refactors the docs of the official Flax MLM, CLM examples. The CLM training script is also slightly changed for consistency with the MLM script.  