| import torch | |
| from transformers import BertModel | |
| from .configuration import NewModelConfig | |
| # It doesn't inherit from a generate-compatible class, so it would throw an error when calling `generate` if | |
| # `transformers` wouldn't automatically patch inheritance | |
| class NewModelForCausalLM(BertModel): | |
| config_class = NewModelConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.last_layer = torch.nn.Linear(config.hidden_size, config.new_hidden_size) | |
| # prior to v4.45, overwriting `prepare_inputs_for_generation` was a hard requirement for a model to be | |
| # `generate`-compatible | |
| def prepare_inputs_for_generation(self): | |
| pass | |