File size: 1,937 Bytes
16e93f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from transformers import T5ForConditionalGeneration, TFT5ForConditionalGeneration
pt_model = T5ForConditionalGeneration.from_pretrained(".", from_flax=True)
pt_model.save_pretrained(".")
# tf_model = TFT5ForConditionalGeneration.from_pretrained(".", from_pt=True)
# tf_model.save_pretrained(".")
exit()
# from transformers import T5ForConditionalGeneration, TFT5ForConditionalGeneration, FlaxT5ForConditionalGeneration
# import numpy as np
# import torch
#
# fx_model = FlaxT5ForConditionalGeneration.from_pretrained(".")
#
# pt_model = T5ForConditionalGeneration.from_pretrained(".", from_flax=True)
# pt_model.save_pretrained(".")
#
#
# # tf_model = TFT5ForConditionalGeneration.from_pretrained(".", from_pt=True)
# # tf_model.save_pretrained(".")
#
#!/usr/bin/env python
import tempfile
import jax
import numpy as np
import torch
from jax import numpy as jnp
from transformers import AutoTokenizer, FlaxT5ForConditionalGeneration, T5ForConditionalGeneration
def to_f32(t):
return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
def main():
# Saving extra files from config.json and tokenizer.json files
tokenizer = AutoTokenizer.from_pretrained("./")
tokenizer.save_pretrained("./")
# Temporary saving bfloat16 Flax model into float32
tmp = tempfile.mkdtemp()
flax_model = FlaxT5ForConditionalGeneration.from_pretrained("./")
flax_model.params = to_f32(flax_model.params)
flax_model.save_pretrained(tmp)
# Converting float32 Flax to PyTorch
pt_model = T5ForConditionalGeneration.from_pretrained(tmp, from_flax=True)
pt_model.save_pretrained("./", save_config=False)
input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
input_ids_pt = torch.tensor(input_ids)
logits_pt = pt_model(input_ids_pt).logits
print(logits_pt)
logits_fx = flax_model(input_ids).logits
print(logits_fx)
if __name__ == "__main__":
main() |