Update vocab_resizer.py

vocab_resizer.py

@@ -1,117 +1,117 @@
-import torch
-from safetensors.torch import save_file, load_file
-import os
-import json
-import glob
-from collections import OrderedDict
-
-def fix_vocab_size(model_dir, output_dir, new_vocab_size=131072):
-    """
-    Resizes the vocabulary-dependent tensors of a sharded model to a new size.
-
-    Args:
-        model_dir (str): The directory of the model to fix.
-        output_dir (str): The directory where the fixed model will be saved.
-        new_vocab_size (int): The target vocabulary size.
-    """
-    try:
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
-            print(f"Created output directory: {output_dir}")
-
-        # --- Step 1: Find all safetensor shards ---
-        search_pattern = os.path.join(model_dir, '*.safetensors')
-        shard_paths = sorted(glob.glob(search_pattern))
-        if not shard_paths:
-            print(f"Error: No '.safetensors' files found in {model_dir}")
-            return
-        
-        print(f"Found {len(shard_paths)} shards to process.")
-
-        # --- Step 2: Identify which shards contain the vocab tensors ---
-        vocab_tensor_keys = ["model.embed_tokens.weight", "lm_head.weight"]
-        shards_to_modify = {} # {filename: {key: tensor}}
-
-        for shard_path in shard_paths:
-            with open(shard_path, "rb") as f:
-                data = f.read()
-            header_size = int.from_bytes(data[:8], 'little')
-            header_str = data[8:8+header_size].decode('utf-8')
-            header = json.loads(header_str)
-            
-            for key in header.keys():
-                if key in vocab_tensor_keys:
-                    filename = os.path.basename(shard_path)
-                    if filename not in shards_to_modify:
-                        shards_to_modify[filename] = {}
-                    # Load only the specific tensor
-                    shards_to_modify[filename][key] = load_file(shard_path)[key]
-                    print(f"Found '{key}' in shard: {filename}")
-
-        if not shards_to_modify:
-            print("Error: Could not find 'embed_tokens' or 'lm_head' tensors in any shard.")
-            return
-
-        # --- Step 3: Process all shards, modifying the ones with vocab tensors ---
-        for shard_path in shard_paths:
-            filename = os.path.basename(shard_path)
-            output_shard_path = os.path.join(output_dir, filename)
-
-            # Load all tensors from the current shard
-            tensors = load_file(shard_path)
-
-            if filename in shards_to_modify:
-                print(f"Resizing tensors in {filename}...")
-                for key, tensor in shards_to_modify[filename].items():
-                    original_size = tensor.shape[0]
-                    print(f"  - Resizing '{key}' from {original_size} to {new_vocab_size}")
-                    # Trim the tensor to the new vocabulary size
-                    resized_tensor = tensor[:new_vocab_size, :]
-                    tensors[key] = resized_tensor # Replace the tensor in the loaded dict
-            
-            # Save the (potentially modified) tensors to the new location
-            save_file(tensors, output_shard_path)
-            print(f"Saved new shard: {output_shard_path}")
-
-        # --- Step 4: Modify and save the config.json ---
-        config_path = os.path.join(model_dir, 'config.json')
-        new_config_path = os.path.join(output_dir, 'config.json')
-        if os.path.exists(config_path):
-            with open(config_path, 'r') as f:
-                config = json.load(f)
-            
-            print(f"\nUpdating config.json: 'vocab_size' from {config.get('vocab_size')} to {new_vocab_size}")
-            config['vocab_size'] = new_vocab_size
-            
-            with open(new_config_path, 'w') as f:
-                json.dump(config, f, indent=2)
-            print(f"Saved new config.json to {new_config_path}")
-        else:
-            print("Warning: config.json not found. Please create it manually.")
-
-        # --- Step 5: Copy other essential files ---
-        for filename in os.listdir(model_dir):
-            if filename.endswith(('.json', '.py', '.md', '.txt')) and filename != 'config.json':
-                 if not os.path.exists(os.path.join(output_dir, filename)):
-                    import shutil
-                    shutil.copy2(os.path.join(model_dir, filename), output_dir)
-                    print(f"Copied {filename} to output directory.")
-
-        print("\nVocabulary resizing complete. The model is now ready for merging.")
-
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-if __name__ == "__main__":
-    # --- Configuration ---
-    # Directory of the original DeepHermes-24B model
-    input_model_directory = r"path/to/your/DeepHermes-24B"
-
-    # Directory to save the fixed, merge-ready model
-    output_model_directory = r"path/to/your/DeepHermes-24B-fixed"
-    
-    # The standard vocab size you are targeting for the merge
-    target_vocab_size = 131072
-
-    # --- Run the script ---
+import torch
+from safetensors.torch import save_file, load_file
+import os
+import json
+import glob
+from collections import OrderedDict
+
+def fix_vocab_size(model_dir, output_dir, new_vocab_size=131072):
+    """
+    Resizes the vocabulary-dependent tensors of a sharded model to a new size.
+
+    Args:
+        model_dir (str): The directory of the model to fix.
+        output_dir (str): The directory where the fixed model will be saved.
+        new_vocab_size (int): The target vocabulary size.
+    """
+    try:
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+            print(f"Created output directory: {output_dir}")
+
+        # --- Step 1: Find all safetensor shards ---
+        search_pattern = os.path.join(model_dir, '*.safetensors')
+        shard_paths = sorted(glob.glob(search_pattern))
+        if not shard_paths:
+            print(f"Error: No '.safetensors' files found in {model_dir}")
+            return
+        
+        print(f"Found {len(shard_paths)} shards to process.")
+
+        # --- Step 2: Identify which shards contain the vocab tensors ---
+        vocab_tensor_keys = ["model.embed_tokens.weight", "lm_head.weight"]
+        shards_to_modify = {} # {filename: {key: tensor}}
+
+        for shard_path in shard_paths:
+            with open(shard_path, "rb") as f:
+                data = f.read()
+            header_size = int.from_bytes(data[:8], 'little')
+            header_str = data[8:8+header_size].decode('utf-8')
+            header = json.loads(header_str)
+            
+            for key in header.keys():
+                if key in vocab_tensor_keys:
+                    filename = os.path.basename(shard_path)
+                    if filename not in shards_to_modify:
+                        shards_to_modify[filename] = {}
+                    # Load only the specific tensor
+                    shards_to_modify[filename][key] = load_file(shard_path)[key]
+                    print(f"Found '{key}' in shard: {filename}")
+
+        if not shards_to_modify:
+            print("Error: Could not find 'embed_tokens' or 'lm_head' tensors in any shard.")
+            return
+
+        # --- Step 3: Process all shards, modifying the ones with vocab tensors ---
+        for shard_path in shard_paths:
+            filename = os.path.basename(shard_path)
+            output_shard_path = os.path.join(output_dir, filename)
+
+            # Load all tensors from the current shard
+            tensors = load_file(shard_path)
+
+            if filename in shards_to_modify:
+                print(f"Resizing tensors in {filename}...")
+                for key, tensor in shards_to_modify[filename].items():
+                    original_size = tensor.shape[0]
+                    print(f"  - Resizing '{key}' from {original_size} to {new_vocab_size}")
+                    # Trim the tensor to the new vocabulary size
+                    resized_tensor = tensor[:new_vocab_size, :]
+                    tensors[key] = resized_tensor # Replace the tensor in the loaded dict
+            
+            # Save the (potentially modified) tensors to the new location
+            save_file(tensors, output_shard_path)
+            print(f"Saved new shard: {output_shard_path}")
+
+        # --- Step 4: Modify and save the config.json ---
+        config_path = os.path.join(model_dir, 'config.json')
+        new_config_path = os.path.join(output_dir, 'config.json')
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+            
+            print(f"\nUpdating config.json: 'vocab_size' from {config.get('vocab_size')} to {new_vocab_size}")
+            config['vocab_size'] = new_vocab_size
+            
+            with open(new_config_path, 'w') as f:
+                json.dump(config, f, indent=2)
+            print(f"Saved new config.json to {new_config_path}")
+        else:
+            print("Warning: config.json not found. Please create it manually.")
+
+        # --- Step 5: Copy other essential files ---
+        for filename in os.listdir(model_dir):
+            if filename.endswith(('.json', '.py', '.md', '.txt')) and filename != 'config.json':
+                 if not os.path.exists(os.path.join(output_dir, filename)):
+                    import shutil
+                    shutil.copy2(os.path.join(model_dir, filename), output_dir)
+                    print(f"Copied {filename} to output directory.")
+
+        print("\nVocabulary resizing complete. The model is now ready for merging.")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+if __name__ == "__main__":
+    # --- Configuration ---
+    # Directory of the original DeepHermes-24B model
+    input_model_directory = r"path/to/your/DeepHermes-24B"
+
+    # Directory to save the fixed, merge-ready model
+    output_model_directory = r"path/to/your/DeepHermes-24B/fixed"
+    
+    # The standard vocab size you are targeting for the merge
+    target_vocab_size = 131072
+
+    # --- Run the script ---
     fix_vocab_size(input_model_directory, output_model_directory, target_vocab_size)