Skip to content

Commit

Permalink
add test fn for clearing import cache
Browse files Browse the repository at this point in the history
  • Loading branch information
sambhavnoobcoder committed Jan 23, 2025
1 parent 0fd0cf4 commit 05ca48f
Showing 1 changed file with 28 additions and 1 deletion.
29 changes: 28 additions & 1 deletion tests/utils/test_import_utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,30 @@
import sys
import pytest
from transformers.utils.import_utils import clear_import_cache, _LazyModule
from transformers.utils.import_utils import clear_import_cache, _LazyModule

def test_clear_import_cache():
# Import some transformers modules
from transformers import AutoModel, AutoTokenizer

# Get initial module count
initial_modules = {
name: mod for name, mod in sys.modules.items()
if name.startswith('transformers.')
}

# Verify we have some modules loaded
assert len(initial_modules) > 0

# Clear cache
clear_import_cache()

# Check modules were removed
remaining_modules = {
name: mod for name, mod in sys.modules.items()
if name.startswith('transformers.')
}
assert len(remaining_modules) < len(initial_modules)

# Verify we can reimport
from transformers import AutoModel, AutoTokenizer
assert 'transformers.models.auto.modeling_auto' in sys.modules

0 comments on commit 05ca48f

Please sign in to comment.