Minor style improvements for the Flax BERT and RoBERTa examples (#8178)
* Minor style improvements:
1. Use `@nn.compact` rather than `@compact` (as to not make it seem
like compact is a standard Python decorator.
2. Move attribute docstrings from two `__call__` methods to comments
on the attributes themselves. (This was probably a remnant from
the pre-Linen version where the attributes were arguments to
`call`.)
* Use black on the Flax modeling code