| import torch
|
| from transformers import PreTrainedModel, PretrainedConfig, GPT2TokenizerFast, \
|
| AutoConfig, AutoModelForCausalLM
|
| from transformers.modeling_outputs import CausalLMOutput
|
| from model import FriendsTransformer
|
|
|
| class PerkLMConfig(PretrainedConfig):
|
| model_type = "perklm"
|
|
|
| def __init__(self, d_model = 512, n_heads = 8, n_layers = 6,
|
| d_ff = 2048, maxt = 512, dropout = 0.1, tokenizer_path = None, **kwargs):
|
| super().__init__(**kwargs)
|
| self.d_model = d_model
|
| self.n_heads = n_heads
|
| self.n_layers = n_layers
|
| self.d_ff = d_ff
|
| self.maxt = maxt
|
| self.dropout = dropout
|
| self.tokenizer_path = tokenizer_path
|
|
|
| class PerkLM(PreTrainedModel):
|
| config_class = PerkLMConfig
|
| _tied_weights_keys = ["transformer.lm_head.weight"]
|
|
|
| def __init__(self, config):
|
| super().__init__(config)
|
| tokenizer = GPT2TokenizerFast.from_pretrained(config.tokenizer_path)
|
| self.transformer = FriendsTransformer(d_model = config.d_model,
|
| n_heads = config.n_heads,
|
| n_layers = config.n_layers,
|
| d_ff = config.d_ff,
|
| dropout = config.dropout,
|
| maxt = config.maxt,
|
| tokenizer = tokenizer)
|
| self.post_init()
|
|
|
| def forward(self, input_ids, attention_mask = None,
|
| responder = None, labels = None, **kwargs):
|
| batch = {'input_ids': input_ids,
|
| 'attention_mask': attention_mask if attention_mask is not None \
|
| else torch.ones_like(input_ids),
|
| 'responder': responder}
|
| logits = self.transformer(batch)
|
|
|
| loss = None
|
| if labels is not None:
|
| loss = torch.nn.functional.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)),
|
| labels[:, 1:].reshape(-1), ignore_index = -100)
|
|
|
| return CausalLMOutput(loss = loss, logits = logits)
|
|
|
| def _save_pretrained_hook(self, *args, **kwargs):
|
| pass
|
|
|
| def state_dict(self, *args, **kwargs):
|
| sd = super().state_dict(*args, **kwargs)
|
|
|
| return {k: v for k, v in sd.items() if v.dtype != torch.complex64}
|
|
|
| def tie_weights(self):
|
| self.transformer.lm_head.weight = self.transformer.embedder.embedding.weight
|
|
|
| AutoConfig.register("perklm", PerkLMConfig)
|
| AutoModelForCausalLM.register(PerkLMConfig, PerkLM)
|
|
|
