| import jax | |
| from jax import numpy as jnp | |
| from transformers import FlaxRobertaForMaskedLM, RobertaForMaskedLM | |
| def to_f32(t): | |
| return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t) | |
| flax_model = FlaxRobertaForMaskedLM.from_pretrained("./") | |
| flax_model.params = to_f32(flax_model.params) | |
| flax_model.save_pretrained("./") | |
| model = RobertaForMaskedLM.from_pretrained("./", from_flax=True) | |
| model.save_pretrained("./", save_config=False) | |