Remote code improvements (#23959)
* Fix model load when it has both code on the Hub and locally * Add input check with timeout * Add tests * Apply suggestions from code review Co-authored-by: Lysandre Debut <lysandre.debut@reseau.eseo.fr> * Some non-saved stuff * Add feature extractors * Add image processor * Add model * Add processor and tokenizer * Reduce timeout --------- Co-authored-by: Lysandre Debut <lysandre.debut@reseau.eseo.fr>
This commit is contained in:
@@ -19,6 +19,7 @@ import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import transformers
|
||||
from transformers import (
|
||||
CONFIG_MAPPING,
|
||||
IMAGE_PROCESSOR_MAPPING,
|
||||
@@ -37,6 +38,9 @@ from test_module.custom_image_processing import CustomImageProcessor # noqa E40
|
||||
|
||||
|
||||
class AutoImageProcessorTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
transformers.dynamic_module_utils.TIME_OUT_REMOTE_CODE = 0
|
||||
|
||||
def test_image_processor_from_model_shortcut(self):
|
||||
config = AutoImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
|
||||
self.assertIsInstance(config, CLIPImageProcessor)
|
||||
@@ -130,6 +134,15 @@ class AutoImageProcessorTest(unittest.TestCase):
|
||||
_ = AutoImageProcessor.from_pretrained("hf-internal-testing/config-no-model")
|
||||
|
||||
def test_from_pretrained_dynamic_image_processor(self):
|
||||
# If remote code is not set, we will time out when asking whether to load the model.
|
||||
with self.assertRaises(ValueError):
|
||||
image_processor = AutoImageProcessor.from_pretrained("hf-internal-testing/test_dynamic_image_processor")
|
||||
# If remote code is disabled, we can't load this config.
|
||||
with self.assertRaises(ValueError):
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
"hf-internal-testing/test_dynamic_image_processor", trust_remote_code=False
|
||||
)
|
||||
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
"hf-internal-testing/test_dynamic_image_processor", trust_remote_code=True
|
||||
)
|
||||
@@ -171,3 +184,35 @@ class AutoImageProcessorTest(unittest.TestCase):
|
||||
del CONFIG_MAPPING._extra_content["custom"]
|
||||
if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
|
||||
del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]
|
||||
|
||||
def test_from_pretrained_dynamic_image_processor_conflict(self):
|
||||
class NewImageProcessor(CLIPImageProcessor):
|
||||
is_local = True
|
||||
|
||||
try:
|
||||
AutoConfig.register("custom", CustomConfig)
|
||||
AutoImageProcessor.register(CustomConfig, NewImageProcessor)
|
||||
# If remote code is not set, the default is to use local
|
||||
image_processor = AutoImageProcessor.from_pretrained("hf-internal-testing/test_dynamic_image_processor")
|
||||
self.assertEqual(image_processor.__class__.__name__, "NewImageProcessor")
|
||||
self.assertTrue(image_processor.is_local)
|
||||
|
||||
# If remote code is disabled, we load the local one.
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
"hf-internal-testing/test_dynamic_image_processor", trust_remote_code=False
|
||||
)
|
||||
self.assertEqual(image_processor.__class__.__name__, "NewImageProcessor")
|
||||
self.assertTrue(image_processor.is_local)
|
||||
|
||||
# If remote is enabled, we load from the Hub
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
"hf-internal-testing/test_dynamic_image_processor", trust_remote_code=True
|
||||
)
|
||||
self.assertEqual(image_processor.__class__.__name__, "NewImageProcessor")
|
||||
self.assertTrue(not hasattr(image_processor, "is_local"))
|
||||
|
||||
finally:
|
||||
if "custom" in CONFIG_MAPPING._extra_content:
|
||||
del CONFIG_MAPPING._extra_content["custom"]
|
||||
if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
|
||||
del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]
|
||||
|
||||
Reference in New Issue
Block a user