update

dbiir · Oct 12, 2023 · 3e82267 · 3e82267
1 parent 05432c0
commit 3e82267
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/scripts/convert_bert_from_uer_to_huggingface.py b/scripts/convert_bert_from_uer_to_huggingface.py
@@ -63,6 +63,18 @@ def main():
 
     convert_bert_transformer_encoder_from_uer_to_huggingface(input_model, output_model, args.layers_num)
 
+    if args.type == "bert":
+        output_model["bert.pooler.dense.weight"] = input_model["target.sp.linear_1.weight"]
+        output_model["bert.pooler.dense.bias"] = input_model["target.sp.linear_1.bias"]
+        output_model["cls.seq_relationship.weight"] = input_model["target.sp.linear_2.weight"]
+        output_model["cls.seq_relationship.bias"] = input_model["target.sp.linear_2.bias"]
+    output_model["cls.predictions.transform.dense.weight"] = input_model["target.mlm.linear_1.weight"]
+    output_model["cls.predictions.transform.dense.bias"] = input_model["target.mlm.linear_1.bias"]
+    output_model["cls.predictions.transform.LayerNorm.weight"] = input_model["target.layer_norm.gamma"]
+    output_model["cls.predictions.transform.LayerNorm.bias"] = input_model["target.layer_norm.beta"]
+    output_model["cls.predictions.decoder.weight"] = input_model["target.mlm.linear_2.weight"]
+    output_model["cls.predictions.bias"] = input_model["target.mlm.linear_2.bias"]
+
     torch.save(output_model, args.output_model_path)