File size: 537 Bytes
b860da0 70a6103 eb8e175 70a6103 b860da0 70a6103 b860da0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from torch import nn
from transformers.modeling_utils import PreTrainedModel
from .configuration_my_model import MyModelConfig
class MyModelPretrainedModel(PreTrainedModel):
config_class = MyModelConfig
class MyModel(MyModelPretrainedModel):
def __init__(self, config: MyModelConfig):
super().__init__(config)
self.config = config
self.n_layers = config.n_layers
self.hidden_dim = config.hidden_dim
self.linear = nn.Linear(config.hidden_dim, config.hidden_dim)
|