#!/usr/bin/env python3
from safetensors.torch import save_file, load_file
import torch

# 1. Clean data (safetensors strictly accepts ONLY a flat dictionary of tensors)
clean_state_dict = {
    "layer1.weight": torch.randn(5, 5),
    "layer1.bias": torch.randn(5)
}

# 2. Saving to safetensors format
safetensors_path = "model.safetensors"
save_file(clean_state_dict, safetensors_path)
print(f"Successfully migrated weights to {safetensors_path}")

# 3. Loading back safely
# No flags needed; it is fundamentally impossible to embed a pickle exploit here
safe_loaded_data = load_file(safetensors_path)
print("Safetensors file loaded with absolute mathematical safety.")