Upload vocab_resizer_mistral24B.py

vocab_resizer_mistral24B.py

@@ -0,0 +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 ---
+    fix_vocab_size(input_model_directory, output_model_directory, target_vocab_size)