import torch
from safetensors.torch import load_file, save_file

def reduce_key_size(input_file, output_file, reduction_factor=0.50):
    # 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 = 'model-00002-of-00002.safetensors'  # Replace with your input model file
output_file = 'model-00002-of-00002.safetensors'  # Desired output file name
reduce_key_size(input_file, output_file)