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