Add pytorch model

config.json

@@ -1,4 +1,5 @@
 {
+  "_name_or_path": "./",
   "architectures": [
     "RobertaForMaskedLM"
   ],
@@ -19,6 +20,7 @@
   "num_hidden_layers": 24,
   "pad_token_id": 1,
   "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
   "transformers_version": "4.10.0.dev0",
   "type_vocab_size": 1,
   "use_cache": true,

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b0ad769a9250a2ed0509dbf2adfe5395758fe975acf4491ff56028ad151e62f9
-size 711588089
+oid sha256:7cf30f16ac72e048be2b0ad47ce76fdf2efcb13b5346dcf8a7d20d633848f7ac
+size 1421662309

flax_model_to_pytorch.py

@@ -0,0 +1,20 @@
+from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM, AutoTokenizer
+import torch
+import numpy as np
+import jax
+import jax.numpy as jnp
+jax.config.update('jax_platform_name', 'cpu')
+MODEL_PATH = "./"
+model = FlaxRobertaForMaskedLM.from_pretrained(MODEL_PATH)
+def to_f32(t):
+    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+model.params = to_f32(model.params)
+model.save_pretrained(MODEL_PATH)
+pt_model = RobertaForMaskedLM.from_pretrained(MODEL_PATH, from_flax=True).to('cpu')
+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 = model(input_ids).logits
+print(logits_fx)
+pt_model.save_pretrained(MODEL_PATH)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:190578a5782162752e1c2aaf1b04ef8b3db300a245a57746ffaea3f22db44963
+size 1421780139