from safetensors import safe_open
from safetensors.torch import save_file
import torch
import hashlib
import random

filename = "model-00001-of-00001.safetensors"
modified_model_name = "modified_model.safetensors"


with safe_open(filename, framework="pt") as f:
    tensors = {key: f.get_tensor(key) for key in f.keys()}

def introduce_noise(tensor, noise_level=1e-8):
    noise = torch.randn(tensor.size()) * noise_level
    return (tensor + noise).to(tensor.dtype)
modified = 0;
for key in tensors:
    modified +=1
    tensors[key] = introduce_noise(tensors[key])
print(modified)

# # modified = 0
# # modification_scale=1e-7
# # for key, tensor in tensors.items():
# #     if not modified and tensor.numel() > 0:
# #             # Find the first non-zero element
# #             non_zero_indices = torch.nonzero(tensor, as_tuple=True)
# #             if len(non_zero_indices[0]) > 0:
# #                 idx = tuple(index[0] for index in non_zero_indices)
# #                 original_value = tensor[idx].item()
# #                 # Add a tiny value
# #                 tensor[idx] += modification_scale * abs(original_value)
# #                 modified += 1
# #                 if modified == 20:
# #                     break
# #                 if not modified:
# #                     print("Could not find a suitable tensor to modify.")

# tensor_keys = list(tensors.keys())
# num_changes=1000
# changes_made = 0
# modification_scale=1e-7
# while changes_made < num_changes and tensor_keys:
#         # Randomly select a tensor
#     key = random.choice(tensor_keys)
#     tensor = tensors[key]
#     if tensor.numel() > 0:
#             # Find non-zero elements
#         non_zero_indices = torch.nonzero(tensor, as_tuple=False)
#         if len(non_zero_indices) > 0:
#                 # Randomly select a non-zero element
#             idx = tuple(random.choice(non_zero_indices).tolist())
#             original_value = tensor[idx].item()
                
#                 # Add a small random value
#             modification = random.uniform(-modification_scale, modification_scale) * abs(original_value)
#             tensor[idx] += modification
                
#             changes_made += 1
            
#         # Remove this tensor from the list to avoid repeated selection
#     tensor_keys.remove(key)

# if changes_made == 0:
#     print("Could not find suitable tensors to modify.")

# print(f"Made {changes_made} changes to the SafeTensors file.")


metadata = {
"format": "pt"  # Adjust based on actual format needed
}

save_file(tensors, modified_model_name, metadata=metadata)


def compute_hash(filename):
    hasher = hashlib.sha256()
    with open(filename, "rb") as f:
        buf = f.read()
        hasher.update(buf)
        return hasher.hexdigest()

original_hash = compute_hash(filename)
modified_hash = compute_hash(modified_model_name)

print(f"Original Hash: {original_hash}")
print(f"Modified Hash: {modified_hash}")