20 lines
670 B
Python
20 lines
670 B
Python
#!/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.")
|