move cleaned data tokenizer to main

cleaned_data/config.json

@@ -1,25 +0,0 @@
-{
-  "architectures": [
-    "RobertaForMaskedLM"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "bos_token_id": 0,
-  "eos_token_id": 2,
-  "gradient_checkpointing": false,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 1024,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-05,
-  "max_position_embeddings": 514,
-  "model_type": "roberta",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 24,
-  "pad_token_id": 1,
-  "position_embedding_type": "absolute",
-  "transformers_version": "4.9.0.dev0",
-  "type_vocab_size": 1,
-  "use_cache": true,
-  "vocab_size": 50265
-}

flax_model_to_pytorch.py

@@ -0,0 +1,30 @@
+from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM, AutoTokenizer
+import torch
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+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)
+
+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)
+
+# also save tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+tokenizer.save_pretrained(MODEL_PATH)