File size: 855 Bytes
fe668e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_gpjtgpt2 import GPJTGPT2Config
from .gpt import GPTModel
class GPJTGPT2Model(PreTrainedModel):
config_class = GPJTGPT2Config
def __init__(self, config):
super().__init__(config)
self.model = GPTModel(config.cfg)
def forward(self, input_ids, **kwargs):
return self.model.forward(input_ids)
class GPJTGPT2ModelForCausalLM(PreTrainedModel, GenerationMixin):
config_class = GPJTGPT2Config
def __init__(self, config):
super().__init__(config)
self.model = GPTModel(config.cfg)
def forward(self, input_ids, **kwargs):
logits = self.model.forward(input_ids)
return CausalLMOutput(logits=logits)
|