Fickle: Firewalled Pickle

People abuse pickle. Especially researchers. Pickle is not secure. Published datasets and ML training weights are often distributed as pickle files (or formats which use pickle files, such as PyTorch checkpoint.ckpt files). Sometimes it is the only format that they are available in.

Examples

Loading basic types is easy:

>>> from fickle import DefaultFirewall
>>> import pickle
>>>
>>> my_picked_data = pickle.dumps({"list": [1, 2, "three", b"four"]})
>>>
>>> firewall = DefaultFirewall()
>>> firewall.loads(my_picked_data)
{'list': [1, 2, 'three', b'four']}

Safely loading PyTorch checkpoint files into numpy arrays is just as easy:

>>> from fickle.ext.pytorch import fake_torch_load_zipped
>>> from zipfile import ZipFile
>>>
>>> zf = ZipFile("/path/to/sd-v1-4.ckpt")
>>> ckpt = fake_torch_load_zipped(zf)
>>> tensor = ckpt["state_dict"]["model.diffusion_model.output_blocks.3.1.norm.weight"]
>>> tensor.array
array([0.39097363, 0.3898967 , 0.35191917, ..., 0.41924757, 0.4031702 ,
       0.37156993], dtype=float32)

You can, optionally, even use marshmallow for validation!

Alternatives

• picklemagic
• pikara