mk / rp.py
pranavajay's picture
Create rp.py
5c67222 verified
raw
history blame
970 Bytes
import torch
from safetensors.torch import load_file, save_file
def reduce_key_size(input_file, output_file, reduction_factor=0.30):
# Load the model
model_data = load_file(input_file)
# Iterate through all the tensors and reduce their size
for key in model_data.keys():
original_tensor = model_data[key]
# Calculate the new size
new_size = int(original_tensor.size(0) * (1 - reduction_factor))
# Resize the tensor (this could vary depending on your requirements)
if new_size > 0: # Ensure new size is positive
reduced_tensor = original_tensor[:new_size]
model_data[key] = reduced_tensor
# Save the modified model
save_file(model_data, output_file)
# Usage example
input_file = 'merged_model.safetensors' # Replace with your input model file
output_file = 'merged_model_16.safetensors' # Desired output file name
reduce_key_size(input_file, output_file)