model: update merge script config

merge.py

@@ -22,8 +22,8 @@ def create_merge_plan(tensor_locations, layer_config):
     # Special handling for specific weights
     special_weights = {
         "model.embed_tokens.weight": 1,
-        "lm_head.weight": 48,
-        "model.norm.weight": 48
+        "lm_head.weight": 156,
+        "model.norm.weight": 156
     }
 
     for slice_config in layer_config:
@@ -65,13 +65,18 @@ def create_merge_plan(tensor_locations, layer_config):
     
     return merge_plan
 
-def merge_layers(input_dir, output_dir, merge_plan):
+def merge_layers(input_dir, output_dir, merge_plan, start_file_index=1):
     output_tensors = {}
-    current_new_file_index = 1
     max_file_index = max(item['new_file_index'] for item in merge_plan)
 
     with tqdm(total=len(merge_plan), desc="Merging layers") as pbar:
-        for file_index in range(1, max_file_index + 1):
+        for file_index in range(start_file_index, max_file_index + 1):
+            output_file = os.path.join(output_dir, f"model-{file_index:05d}-of-{max_file_index:05d}.safetensors")
+            
+            if os.path.exists(output_file):
+                pbar.update(sum(1 for item in merge_plan if item['new_file_index'] == file_index))
+                continue
+
             for item in merge_plan:
                 if item['new_file_index'] == file_index:
                     input_file = os.path.join(input_dir, f"model-{item['original_file_index']:05d}-of-00051.safetensors")
@@ -81,7 +86,6 @@ def merge_layers(input_dir, output_dir, merge_plan):
                     pbar.update(1)
 
             if output_tensors:
-                output_file = os.path.join(output_dir, f"model-{file_index:05d}-of-{max_file_index:05d}.safetensors")
                 save_file(output_tensors, output_file)
                 output_tensors = {}
 
@@ -92,6 +96,7 @@ def main():
     parser.add_argument("input_dir", help="Directory containing input safetensors files")
     parser.add_argument("output_dir", help="Directory for output safetensors files")
     parser.add_argument("--dry-run", action="store_true", help="Perform a dry run and output merge plan")
+    parser.add_argument("--continue-from", type=int, default=1, help="Continue merging from this file index")
     args = parser.parse_args()
 
     layer_config = [
@@ -116,8 +121,8 @@ def main():
         print("Merge plan saved to merge_plan.json")
     else:
         os.makedirs(args.output_dir, exist_ok=True)
-        merge_layers(args.input_dir, args.output_dir, merge_plan)
+        merge_layers(args.input_dir, args.output_dir, merge_plan, start_file_index=args.continue_from)
         print(f"Merged model saved to {args.output_dir}")
 
 if __name__ == "__main__":
-    main()
+    main()