change apply_rotary_pos_emb of Glmmodel for GLM-Edge Series model

src/transformers/models/glm/configuration_glm.py

@@ -93,6 +93,7 @@ def __init__(
         num_hidden_layers=40,
         num_attention_heads=32,
         num_key_value_heads=2,
+        partial_rotary_factor=0.5,
         head_dim=128,
         hidden_act="silu",
         attention_dropout=0.0,
@@ -114,6 +115,7 @@ def __init__(
         self.intermediate_size = intermediate_size
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
+        self.partial_rotary_factor = partial_rotary_factor
         self.head_dim = head_dim
         self.num_key_value_heads = num_key_value_heads
         self.hidden_act = hidden_act

src/transformers/models/glm/convert_glm_weights_to_hf.py

@@ -2,7 +2,6 @@
 import json
 import os
 import re
-
 import torch
 from safetensors.torch import load_file
 from tokenizers import processors
@@ -37,16 +36,28 @@
 # fmt: on
 
 
-def merge_safetensors(input_dir: str):
-    all_files = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith(".safetensors")]
-    all_files = sorted(all_files, key=lambda x: int(x.rsplit("-", 3)[1]))
+def load_weights(input_dir: str):
+    safetensor_files = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith(".safetensors")]
+    bin_files = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith(".bin")]
 
     all_weights = {}
-    for file in all_files:
-        tensors = load_file(file)
-        all_weights.update(tensors)
 
-    return all_weights
+    if safetensor_files:
+        safetensor_files = sorted(safetensor_files, key=lambda x: int(x.rsplit("-", 3)[1]))
+        for file in safetensor_files:
+            tensors = load_file(file)
+            all_weights.update(tensors)
+        return all_weights
+
+    elif bin_files:
+        bin_files = sorted(bin_files, key=lambda x: int(x.rsplit("-", 3)[1]))
+        for file in bin_files:
+            tensors = torch.load(file, map_location="cpu")
+            all_weights.update(tensors)
+        return all_weights
+
+    else:
+        raise ValueError("No .safetensors or .bin files found in the specified directory.")
 
 
 def map_old_key_to_new(old_key):
@@ -100,7 +111,8 @@ def convert_config(original_config: dict):
         "attention_bias": "add_qkv_bias",
     }
     similar_keys_to_keep = [
-        "num_attention_heads" "hidden_size",
+        "num_attention_heads",
+        "hidden_size",
         "attention_dropout",
         "use_cache",
         "eos_token_id",
@@ -120,40 +132,43 @@ def convert_config(original_config: dict):
     return new_config
 
 
-def convert_glm_tokenizer(input_dir):
+def convert_glm_tokenizer(input_dir, use_post_processor=False):
     fast_tok = PreTrainedTokenizerFast.from_pretrained(input_dir, model_input_names=["input_ids", "attention_mask"])
-    # Add the two tokens automatically with post processor
-    fast_tok._tokenizer.post_processor = processors.Sequence(
-        [
-            processors.ByteLevel(trim_offsets=False),
-            processors.TemplateProcessing(
-                single="[gMASK]:0 <sop>:0 $A:0",
-                pair="[gMASK]:0 <sop>:0 $A:0 $B:1",
-                special_tokens=[("[gMASK]", 151331), ("<sop>", 151333)],
-            ),
-        ],
-    )
-
+    if use_post_processor:
+        fast_tok._tokenizer.post_processor = processors.Sequence(
+            [
+                processors.ByteLevel(trim_offsets=False),
+                processors.TemplateProcessing(
+                    single="[gMASK]:0 <sop>:0 $A:0",
+                    pair="[gMASK]:0 <sop>:0 $A:0 $B:1",
+                    special_tokens=[("[gMASK]", 151331), ("<sop>", 151333)],
+                ),
+            ],
+        )
+    else:
+        fast_tok._tokenizer.post_processor = processors.Sequence(
+            [processors.ByteLevel(trim_offsets=False)],
+        )
     return fast_tok
 
 
-def convert_glm_model(input_dir, output_dir):
+def convert_glm_model(input_dir, output_dir, use_post_processor=False):
     # Load and convert config
     with open(os.path.join(input_dir, "config.json")) as f:
         original_config = json.load(f)
     config = convert_config(original_config)
     config.save_pretrained(output_dir)
 
     # Load and convert weights
-    original_state_dict = merge_safetensors(input_dir)
+    original_state_dict = load_weights(input_dir)
     new_dict = convert_state_dict(original_state_dict, config)
     with torch.device("meta"):
         model = GlmForCausalLM(config)
     model.load_state_dict(new_dict, strict=True, assign=True)
     model.save_pretrained(output_dir)
 
     # Load and convert tokenizer
-    tokenizer = convert_glm_tokenizer(input_dir)
+    tokenizer = convert_glm_tokenizer(input_dir, use_post_processor)
     tokenizer.save_pretrained(output_dir)
 
 
@@ -169,6 +184,11 @@ def convert_glm_model(input_dir, output_dir):
         type=str,
         help="Location to write HF model and tokenizer",
     )
+    parser.add_argument(
+        "--use_post_processor",
+        action="store_true",
+        help="Whether to apply post processor with special tokens",
+    )
 
     args = parser.parse_args()
-    convert_glm_model(args.input_dir, args.output_dir)
+    convert_glm_model(args.input_dir, args.output_dir, args.use_post_processor)