Spaces:
Running
Running
File size: 5,186 Bytes
aa8807d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 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) |