wehub-resource-sync
35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from contextlib import contextmanager
|
|
from functools import wraps
|
|
from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase
|
|
|
|
|
|
@contextmanager
|
|
def patch_auto_tokenizer(tokenizer: PreTrainedTokenizerBase):
|
|
_old_from_pretrained = AutoTokenizer.from_pretrained
|
|
|
|
@wraps(_old_from_pretrained)
|
|
def _from_pretrained(*args, **kwargs):
|
|
return tokenizer
|
|
|
|
AutoTokenizer.from_pretrained = _from_pretrained
|
|
try:
|
|
yield
|
|
finally:
|
|
AutoTokenizer.from_pretrained = _old_from_pretrained
|
|
|
|
|
|
@contextmanager
|
|
def patch_auto_config(config: PretrainedConfig):
|
|
_old_from_pretrained = AutoConfig.from_pretrained
|
|
|
|
@wraps(_old_from_pretrained)
|
|
def _from_pretrained(*args, **kwargs):
|
|
return (config, {}) if 'return_unused_kwargs' in kwargs else config
|
|
|
|
AutoConfig.from_pretrained = _from_pretrained
|
|
try:
|
|
yield
|
|
finally:
|
|
AutoConfig.from_pretrained = _old_from_pretrained
|