| import torch |
| 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) |
| self.post_init() |
|
|
|
|
| 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) |
| self.post_init() |
|
|
|
|
| def forward(self, input_ids, attention_mask=None, labels=None, **kwargs): |
| logits = self.model.forward(input_ids) |
|
|
| loss = None |
| if labels is not None: |
| shifted_logits = logits[:, :-1, :] |
| shifted_labels = labels[:, 1:] |
|
|
| if attention_mask is not None: |
| shifted_mask = attention_mask[:, 1:] |
| shifted_labels = shifted_labels.masked_fill( |
| shifted_mask == 0, -100 |
| ) |
|
|
| loss = torch.nn.functional.cross_entropy( |
| shifted_logits.flatten(0, 1), shifted_labels.flatten(), |
| ignore_index=-100 |
| ) |
|
|
| return CausalLMOutput(logits=logits, loss=loss) |
|
|
|
|
