add llamafy_qwen.py

Former-commit-id: 6cdc91543c022edcc98076488f06e809fde9bad7

tests/llamafy_baichuan2.py

@@ -1,6 +1,6 @@
 # coding=utf-8
 # Converts the Baichuan2-7B model in the same format as LLaMA2-7B.
-# Usage: python llamafy_baichuan2.py --llama2_json llama2.index.json --input_dir input --output_dir output
+# Usage: python llamafy_baichuan2.py --input_dir input --output_dir output --shard_size 10GB
 # Inspired by: https://huggingface.co/fireballoon/baichuan-llama-7b/blob/main/convert_baichuan_to_llama.py
 # Converted model: https://huggingface.co/hiyouga/Baichuan2-7B-Base-LLaMAfied
 
@@ -9,56 +9,77 @@ import fire
 import json
 import torch
 from collections import OrderedDict
+from transformers.modeling_utils import shard_checkpoint, WEIGHTS_NAME, WEIGHTS_INDEX_NAME
+from typing import Any, Dict
 
 
-SHARD_A = "pytorch_model-00001-of-00002.bin"
-SHARD_B = "pytorch_model-00002-of-00002.bin"
+CONFIG_NAME = "config.json"
 
 
-def llamafy_baichuan2(
-    llama2_json: str,
+def save_weight(
     input_dir: str,
-    output_dir: str
+    output_dir: str,
+    shard_size: str
 ):
-    baichuan2_state_dict = OrderedDict()
+    baichuan2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
     for filepath in os.listdir(input_dir):
         if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".bin"):
             shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu")
             baichuan2_state_dict.update(shard_weight)
 
-    llama2_state_dict = OrderedDict()
-    total_size = 0
+    llama2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
     for key, value in baichuan2_state_dict.items():
-        total_size += 2 * value.numel() # half precision
         if "W_pack" in key:
-            llama2_state_dict[key.replace("W_pack", "q_proj")] = value[:4096, :]
-            llama2_state_dict[key.replace("W_pack", "k_proj")] = value[4096:2*4096, :]
-            llama2_state_dict[key.replace("W_pack", "v_proj")] = value[2*4096:, :]
+            proj_size = value.size(0) // 3
+            llama2_state_dict[key.replace("W_pack", "q_proj")] = value[:proj_size, :]
+            llama2_state_dict[key.replace("W_pack", "k_proj")] = value[proj_size:2*proj_size, :]
+            llama2_state_dict[key.replace("W_pack", "v_proj")] = value[2*proj_size:, :]
         elif "lm_head" in key:
             llama2_state_dict[key] = torch.nn.functional.normalize(value)
         else:
             llama2_state_dict[key] = value
 
-    with open(os.path.join(input_dir, llama2_json), "r", encoding="utf-8") as f:
-        llama2_index = json.load(f)
+    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=WEIGHTS_NAME)
+    for shard_file, shard in shards.items():
+        torch.save(shard, os.path.join(output_dir, shard_file))
+    
+    if index is None:
+        print("Model weights saved in {}".format(os.path.join(output_dir, WEIGHTS_NAME)))
+    else:
+        with open(os.path.join(output_dir, WEIGHTS_INDEX_NAME), "w", encoding="utf-8") as f:
+            json.dump(index, f, indent=2, sort_keys=True)
+        print("Model weights saved in {}".format(output_dir))
 
-    merged_index = OrderedDict()
-    merged_index["metadata"] = {"total_size": total_size}
-    merged_index["weight_map"] = llama2_index["weight_map"]
 
-    state_dict_a, state_dict_b = OrderedDict(), OrderedDict()
-    for key, value in llama2_state_dict.items():
-        if merged_index["weight_map"][key] == SHARD_A:
-            state_dict_a[key] = value
-        else:
-            state_dict_b[key] = value
+def save_config(
+    input_dir: str,
+    output_dir: str
+):
+    with open(os.path.join(input_dir, CONFIG_NAME), "r", encoding="utf-8") as f:
+        llama2_config_dict: Dict[str, Any] = json.load(f)
 
-    os.makedirs(output_dir, exist_ok=True)
-    torch.save(state_dict_a, os.path.join(output_dir, SHARD_A))
-    torch.save(state_dict_b, os.path.join(output_dir, SHARD_B))
-    with open(os.path.join(output_dir, "pytorch_model.bin.index.json"), "w", encoding="utf-8") as f:
-        json.dump(merged_index, f, indent=2)
-    print("Completed!")
+    llama2_config_dict["architectures"] = ["LlamaForCausalLM"]
+    llama2_config_dict.pop("auto_map", None)
+    llama2_config_dict.pop("tokenizer_class", None)
+    llama2_config_dict["model_type"] = "llama"
+
+    with open(os.path.join(output_dir, CONFIG_NAME), "w", encoding="utf-8") as f:
+        json.dump(llama2_config_dict, f, indent=2)
+    print("Model config saved in {}".format(os.path.join(output_dir, CONFIG_NAME)))
+
+
+def llamafy_baichuan2(
+    input_dir: str,
+    output_dir: str,
+    shard_size: str
+):
+    try:
+        os.makedirs(output_dir, exist_ok=False)
+    except Exception as e:
+        raise print("Output dir already exists", e)
+
+    save_weight(input_dir, output_dir, shard_size)
+    save_config(input_dir, output_dir)    
 
 
 if __name__ == "__main__":