| # https://huggingface.co/docs/transformers/custom_models | |
| from transformers import PreTrainedModel, AutoModelForCausalLM, AutoTokenizer, AutoModel, AutoConfig | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| from torch.nn import CrossEntropyLoss | |
| from torch.nn.functional import log_softmax | |
| from torch.nn.modules.container import ModuleList | |
| from .configuration_custom5 import CustomConfig5 | |
| class CustomModel5(PreTrainedModel): | |
| config_class = CustomConfig5 | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.model = ModuleList([AutoModelForCausalLM.from_pretrained(m) for m in config.models]) | |
| def forward(self, *args, labels=None, **kwargs): | |
| loss = None | |
| logits = None | |
| for model, coeff in zip(self.model, self.config.coeffs): | |
| logp = log_softmax(model.forward(*args, **kwargs).logits, dim=-1) | |
| logits = coeff * logp if logits is None else logits + coeff * logp | |
| # The rest copied from modeling_llama.py: | |
| if labels is not None: | |
| # Shift so that tokens < n predict n | |
| shift_logits = logits[..., :-1, :].contiguous() | |
| shift_labels = labels[..., 1:].contiguous() | |
| # Flatten the tokens | |
| loss_fct = CrossEntropyLoss() | |
| shift_logits = shift_logits.view(-1, self.config.vocab_size) | |
| shift_labels = shift_labels.view(-1) | |
| # Enable model parallelism | |
| shift_labels = shift_labels.to(shift_logits.device) | |
| loss = loss_fct(shift_logits, shift_labels) | |
| return CausalLMOutputWithPast(loss=loss, logits=logits) | |
| ## Which one do we use? | |
| ## You have to tell the library you want to copy the code files of those objects when using the save_pretrained method and properly register them with a given Auto class (especially for models), just run: | |
| # CustomConfig5.register_for_auto_class() | |
| # CustomModel5.register_for_auto_class('AutoModelForCausalLM') | |
| # CustomModel5.register_for_auto_class('AutoModel') | |
| ## If you are writing a library that extends 🤗 Transformers, you may want to extend the auto classes to include your own model. This is different from pushing the code to the Hub in the sense that users will need to import your library to get the custom models (contrarily to automatically downloading the model code from the Hub). | |
| ## As long as your config has a model_type attribute that is different from existing model types, and that your model classes have the right config_class attributes, you can just add them to the auto classes like this: | |
| # AutoConfig.register("custom5", CustomConfig5) | |
| # AutoModel.register(CustomConfig5, CustomModel5) | |
| # AutoModelForCausalLM.register(CustomConfig5, CustomModel5) | |