File size: 468 Bytes
ffdda5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
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") |