mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-04 05:10:06 +06:00
36 lines
838 B
Python
36 lines
838 B
Python
import torch
|
|
|
|
from transformers import PreTrainedModel
|
|
|
|
from .custom_configuration import CustomConfig, NoSuperInitConfig
|
|
|
|
|
|
class CustomModel(PreTrainedModel):
|
|
config_class = CustomConfig
|
|
base_model_prefix = "custom"
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.linear = torch.nn.Linear(config.hidden_size, config.hidden_size)
|
|
|
|
def forward(self, x):
|
|
return self.linear(x)
|
|
|
|
def _init_weights(self, module):
|
|
pass
|
|
|
|
|
|
class NoSuperInitModel(PreTrainedModel):
|
|
config_class = NoSuperInitConfig
|
|
base_model_prefix = "custom"
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.linear = torch.nn.Linear(config.attribute, config.attribute)
|
|
|
|
def forward(self, x):
|
|
return self.linear(x)
|
|
|
|
def _init_weights(self, module):
|
|
pass
|