Add an API to register objects to Auto classes (#13989)
* Add API to register a new object in auto classes * Fix test * Documentation * Add to tokenizers and test * Add cleanup after tests * Be more careful * Move import * Move import * Cleanup in TF test too * Add consistency check * Add documentation * Style * Update docs/source/model_doc/auto.rst Co-authored-by: Lysandre Debut <lysandre@huggingface.co> * Update src/transformers/models/auto/auto_factory.py Co-authored-by: Lysandre Debut <lysandre@huggingface.co> Co-authored-by: Lysandre Debut <lysandre@huggingface.co>
This commit is contained in:
@@ -27,7 +27,32 @@ Instantiating one of :class:`~transformers.AutoConfig`, :class:`~transformers.Au
|
||||
|
||||
will create a model that is an instance of :class:`~transformers.BertModel`.
|
||||
|
||||
There is one class of :obj:`AutoModel` for each task, and for each backend (PyTorch or TensorFlow).
|
||||
There is one class of :obj:`AutoModel` for each task, and for each backend (PyTorch, TensorFlow, or Flax).
|
||||
|
||||
Extending the Auto Classes
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Each of the auto classes has a method to be extended with your custom classes. For instance, if you have defined a
|
||||
custom class of model :obj:`NewModel`, make sure you have a :obj:`NewModelConfig` then you can add those to the auto
|
||||
classes like this:
|
||||
|
||||
.. code-block::
|
||||
|
||||
from transformers import AutoConfig, AutoModel
|
||||
|
||||
AutoConfig.register("new-model", NewModelConfig)
|
||||
AutoModel.register(NewModelConfig, NewModel)
|
||||
|
||||
You will then be able to use the auto classes like you would usually do!
|
||||
|
||||
.. warning::
|
||||
|
||||
If your :obj:`NewModelConfig` is a subclass of :class:`~transformer.PretrainedConfig`, make sure its
|
||||
:obj:`model_type` attribute is set to the same key you use when registering the config (here :obj:`"new-model"`).
|
||||
|
||||
Likewise, if your :obj:`NewModel` is a subclass of :class:`~transformers.PreTrainedModel`, make sure its
|
||||
:obj:`config_class` attribute is set to the same class you use when registering the model (here
|
||||
:obj:`NewModelConfig`).
|
||||
|
||||
|
||||
AutoConfig
|
||||
|
||||
Reference in New Issue
Block a user