Load dynamic module (remote code) only once if code isn't change (#33162)

* Load remote code only once

* Use hash as load indicator

* Add a new option `force_reload` for old behavior (i.e. always reload)

* Add test for dynamic module is cached

* Add more type annotations to improve code readability

* Address comments from code review
This commit is contained in:
Xuehai Pan
2024-09-06 19:49:35 +08:00
committed by GitHub
parent 1bd9d1c899
commit e1c2b69c34
6 changed files with 139 additions and 12 deletions

View File

@@ -122,12 +122,27 @@ class AutoConfigTest(unittest.TestCase):
config = AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model", trust_remote_code=True)
self.assertEqual(config.__class__.__name__, "NewModelConfig")
# Test the dynamic module is loaded only once.
reloaded_config = AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model", trust_remote_code=True)
self.assertIs(config.__class__, reloaded_config.__class__)
# Test config can be reloaded.
with tempfile.TemporaryDirectory() as tmp_dir:
config.save_pretrained(tmp_dir)
reloaded_config = AutoConfig.from_pretrained(tmp_dir, trust_remote_code=True)
self.assertEqual(reloaded_config.__class__.__name__, "NewModelConfig")
# The configuration file is cached in the snapshot directory. So the module file is not changed after dumping
# to a temp dir. Because the revision of the configuration file is not changed.
# Test the dynamic module is loaded only once if the configuration file is not changed.
self.assertIs(config.__class__, reloaded_config.__class__)
# Test the dynamic module is reloaded if we force it.
reloaded_config = AutoConfig.from_pretrained(
"hf-internal-testing/test_dynamic_model", trust_remote_code=True, force_download=True
)
self.assertIsNot(config.__class__, reloaded_config.__class__)
def test_from_pretrained_dynamic_config_conflict(self):
class NewModelConfigLocal(BertConfig):
model_type = "new-model"