def register_bert_model(bert_cls): """ This function wraps a BertModel inherited cls and automatically: 1. Creates an associated BertConfig 2. Creates an associated BertForMaskedLM 3. Creates an associated BertForSequenceClassification 4. Creates an associated BertForQuestionAnswering 5. Registers these classes with Transformers model mappings This last step ensures that the resulting config and models may be used by AutoConfig, AutoModelForMaskedLM, and AutoModelForSequenceClassification. Assumptions are made to auto-name these classes and the corresponding model type. For instance, SparseBertModel will have model_type="sparse_bert" and associated classes like SparseBertConfig. To customize the the inputs to the model's config, include the dataclass `bert_cls.ConfigKWargs`. This is, in fact, required. Upon initialization of the config, the fields of that dataclass will be used to extract extra keyword arguments and assign them as attributes to the config. Example ``` @register_bert_model class SparseBertModel(BertModel): @dataclass class ConfigKWargs: # Keyword arguments to configure sparsity. sparsity: float = 0.9 # Define __init__, ect. ... # Model is ready to auto load. config = AutoConfig.for_model("sparse_bert", sparsity=0.5) model = AutoModelForMaskedLM.from_config(model) config.sparsity >>> 0.5 type(model) >>> SparseBertModelForMaskedLM """ assert bert_cls.__name__.endswith("BertModel") # Get first part of name e.g. StaticSparseBertModel -> StaticSparse name_prefix = bert_cls.__name__.replace("BertModel", "") # Create new bert config and models based off of `bert_cls`. config_cls = create_config_class(bert_cls, name_prefix) masked_lm_cls = create_masked_lm_class(bert_cls, name_prefix) seq_classification_cls = create_sequence_classification_class(bert_cls, name_prefix) question_answering_cls = create_question_answering_class(bert_cls, name_prefix) # Specify the correct config class bert_cls.config_class = config_cls masked_lm_cls.config_class = config_cls seq_classification_cls.config_class = config_cls question_answering_cls.config_class = config_cls # Update Transformers mappings to auto-load these new models. CONFIG_MAPPING.update({ config_cls.model_type: config_cls }) TOKENIZER_MAPPING.update({ config_cls: (BertTokenizer, BertTokenizerFast), }) MODEL_FOR_MASKED_LM_MAPPING.update({ config_cls: masked_lm_cls, }) MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.update({ config_cls: seq_classification_cls }) MODEL_FOR_QUESTION_ANSWERING_MAPPING.update({ config_cls: question_answering_cls }) # Update the `models` modules so that these classes may be imported. __models_dict__.update({ config_cls.__name__: config_cls, masked_lm_cls.__name__: masked_lm_cls, seq_classification_cls.__name__: seq_classification_cls, question_answering_cls.__name__: question_answering_cls, })
from transformers import CONFIG_MAPPING, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, MODEL_NAMES_MAPPING, TOKENIZER_MAPPING from transformers.convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS, BertConverter from transformers.models.auto.modeling_auto import auto_class_factory from .modeling.layoutlmv2 import ( LayoutLMv2Config, LayoutLMv2ForTokenClassification, LayoutLMv2Tokenizer, LayoutLMv2TokenizerFast, ) CONFIG_MAPPING.update([("layoutlmv2", LayoutLMv2Config)]) MODEL_NAMES_MAPPING.update([("layoutlmv2", "LayoutLMv2")]) TOKENIZER_MAPPING.update([(LayoutLMv2Config, (LayoutLMv2Tokenizer, LayoutLMv2TokenizerFast))]) SLOW_TO_FAST_CONVERTERS.update({"LayoutLMv2Tokenizer": BertConverter}) MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.update([ (LayoutLMv2Config, LayoutLMv2ForTokenClassification) ]) AutoModelForTokenClassification = auto_class_factory( "AutoModelForTokenClassification", MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, head_doc="token classification")
MarkupLMTokenizer, MarkupLMForQuestionAnswering, MarkupLMForTokenClassification, MarkupLMTokenizerFast, ) CONFIG_MAPPING.update( [ ("markuplm", MarkupLMConfig), ] ) MODEL_NAMES_MAPPING.update([("markuplm", "MarkupLM")]) TOKENIZER_MAPPING.update( [ (MarkupLMConfig, (MarkupLMTokenizer, MarkupLMTokenizerFast)), ] ) SLOW_TO_FAST_CONVERTERS.update( {"MarkupLMTokenizer": RobertaConverter} ) MODEL_FOR_QUESTION_ANSWERING_MAPPING.update( [(MarkupLMConfig, MarkupLMForQuestionAnswering)] ) MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.update( [(MarkupLMConfig, MarkupLMForTokenClassification)] )