File size: 698 Bytes
7582c1e 32c02cf 7582c1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from transformers import PreTrainedModel
from .configuration_test import TestConfig
import torch.nn as nn
from transformers import AutoModelForMaskedLM, AutoConfig
from transformers import AutoModelForSequenceClassification
class TestModel(PreTrainedModel):
config_class = TestConfig
def __init__(self, config: TestConfig):
super().__init__(config)
self.input_dim = config.input_dim
self.model1 = nn.Linear(config.input_dim, config.output_dim)
# self.model2 = AutoModelForMaskedLM.from_config(
# AutoConfig.from_pretrained("bert-base-uncased")
# )
def forward(self, tensor):
return self.model1(tensor)
|