| import jax | |
| import jax.numpy as jnp | |
| from transformers import FlaxWav2Vec2ForPreTraining, Wav2Vec2ForPreTraining | |
| def to_f32(t): | |
| return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t) | |
| model_fx = FlaxWav2Vec2ForPreTraining.from_pretrained("../") | |
| model_fx.params = to_f32(model_fx.params) | |
| model_fx.save_pretrained("./fx") | |
| model_pt = Wav2Vec2ForPreTraining.from_pretrained("./fx", from_flax=True) | |
| model_pt.save_pretrained("./pt") |