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("../v0")
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")