Spaces:
Runtime error
Runtime error
File size: 804 Bytes
4121bec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import logging
from transformers import AutoConfig
from transformers import AutoModel
from .registry import register_lang_encoder
logger = logging.getLogger(__name__)
@register_lang_encoder
def lang_encoder(config_encoder, tokenizer, verbose, **kwargs):
hf_model = None
if config_encoder['LOAD_PRETRAINED']:
hf_model = AutoModel.from_pretrained(config_encoder['HF_MODEL'])
else:
hf_config = AutoConfig.from_pretrained(config_encoder['HF_MODEL'])
if 'CONFIG_OVERRIDE' in config_encoder:
logger.warning(f'Override config: {config_encoder["CONFIG_OVERRIDE"]}')
hf_config.update(config_encoder['CONFIG_OVERRIDE'])
logger.info(f'HF model config: {hf_config}')
hf_model = AutoModel.from_config(hf_config)
return hf_model
|