Skip to content

Commit

Permalink
Fix failing test with race condition (#32140)
Browse files Browse the repository at this point in the history
* Fix failing test with race condition

* make fixup

* monotonic_ns instead of randint

* uuid4 instead of monotonic_ns

* Add a finally cleanup step
  • Loading branch information
Rocketknight1 authored Jul 22, 2024
1 parent 5a649ff commit 7ba028f
Showing 1 changed file with 39 additions and 34 deletions.
73 changes: 39 additions & 34 deletions tests/models/auto/test_processor_auto.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import unittest
from pathlib import Path
from shutil import copyfile
from uuid import uuid4

from huggingface_hub import HfFolder, Repository, create_repo, delete_repo
from requests.exceptions import HTTPError
Expand Down Expand Up @@ -431,38 +432,42 @@ def test_push_to_hub_dynamic_processor(self):

processor = CustomProcessor(feature_extractor, tokenizer)

with tempfile.TemporaryDirectory() as tmp_dir:
create_repo(f"{USER}/test-dynamic-processor", token=self._token)
repo = Repository(tmp_dir, clone_from=f"{USER}/test-dynamic-processor", token=self._token)
processor.save_pretrained(tmp_dir)

# This has added the proper auto_map field to the feature extractor config
self.assertDictEqual(
processor.feature_extractor.auto_map,
{
"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor",
"AutoProcessor": "custom_processing.CustomProcessor",
},
)

# This has added the proper auto_map field to the tokenizer config
with open(os.path.join(tmp_dir, "tokenizer_config.json")) as f:
tokenizer_config = json.load(f)
self.assertDictEqual(
tokenizer_config["auto_map"],
{
"AutoTokenizer": ["custom_tokenization.CustomTokenizer", None],
"AutoProcessor": "custom_processing.CustomProcessor",
},
)

# The code has been copied from fixtures
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_feature_extraction.py")))
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_tokenization.py")))
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_processing.py")))

repo.push_to_hub()
random_repo_id = f"{USER}/test-dynamic-processor-{uuid4()}"
try:
with tempfile.TemporaryDirectory() as tmp_dir:
create_repo(random_repo_id, token=self._token)
repo = Repository(tmp_dir, clone_from=random_repo_id, token=self._token)
processor.save_pretrained(tmp_dir)

new_processor = AutoProcessor.from_pretrained(f"{USER}/test-dynamic-processor", trust_remote_code=True)
# Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
self.assertEqual(new_processor.__class__.__name__, "CustomProcessor")
# This has added the proper auto_map field to the feature extractor config
self.assertDictEqual(
processor.feature_extractor.auto_map,
{
"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor",
"AutoProcessor": "custom_processing.CustomProcessor",
},
)

# This has added the proper auto_map field to the tokenizer config
with open(os.path.join(tmp_dir, "tokenizer_config.json")) as f:
tokenizer_config = json.load(f)
self.assertDictEqual(
tokenizer_config["auto_map"],
{
"AutoTokenizer": ["custom_tokenization.CustomTokenizer", None],
"AutoProcessor": "custom_processing.CustomProcessor",
},
)

# The code has been copied from fixtures
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_feature_extraction.py")))
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_tokenization.py")))
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, "custom_processing.py")))

repo.push_to_hub()

new_processor = AutoProcessor.from_pretrained(random_repo_id, trust_remote_code=True)
# Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
self.assertEqual(new_processor.__class__.__name__, "CustomProcessor")
finally:
delete_repo(repo_id=random_repo_id)

0 comments on commit 7ba028f

Please sign in to comment.