File size: 677 Bytes
fbcbb1a
9e037f6
fbcbb1a
 
 
9e037f6
 
 
fbcbb1a
 
 
 
29ac323
60de8dd
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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