output: tuple: (loss, ) in training
| 9 | WEIGHTS_NAME = "pytorch_model.bin" |
| 10 | |
| 11 | class SQL2TextModel(nn.Module): |
| 12 | """ |
| 13 | output: tuple: (loss, ) in training |
| 14 | """ |
| 15 | def __init__(self): |
| 16 | super().__init__() |
| 17 | self.bert = BartForConditionalGeneration.from_pretrained("facebook/bart-large") |
| 18 | |
| 19 | |
| 20 | def forward(self, *input, **kwargs): |
| 21 | input_ids = kwargs.pop("input_ids") |
| 22 | |
| 23 | pad_token_id = kwargs.pop("pad_token_id") |
| 24 | attention_mask = (input_ids != pad_token_id).long() |
| 25 | |
| 26 | if self.training: |
| 27 | output_ids = kwargs.pop('labels') |
| 28 | y_ids = output_ids[:, :-1].contiguous() |
| 29 | lm_labels = output_ids[:, 1:].clone() |
| 30 | lm_labels[output_ids[:, 1:] == pad_token_id] = -100 |
| 31 | |
| 32 | outputs = self.bert(input_ids, |
| 33 | attention_mask=attention_mask, decoder_input_ids=y_ids, lm_labels=lm_labels, ) |
| 34 | return (outputs[0],) |
| 35 | |
| 36 | else: |
| 37 | label_eos_id = kwargs.pop("label_eos_id") |
| 38 | label_bos_id = kwargs.pop("label_bos_id") |
| 39 | label_padding_id = kwargs.pop("label_padding_id") |
| 40 | generated_ids = self.bert.generate( |
| 41 | input_ids=input_ids, |
| 42 | attention_mask=attention_mask, |
| 43 | num_beams=3, |
| 44 | max_length=60, |
| 45 | length_penalty=2.0, |
| 46 | early_stopping=True, |
| 47 | use_cache=True, |
| 48 | decoder_start_token_id=label_bos_id, |
| 49 | eos_token_id=label_eos_id, |
| 50 | pad_token_id=label_padding_id |
| 51 | ) |
| 52 | |
| 53 | output_ids = kwargs.pop('labels') |
| 54 | y_ids = output_ids[:, :-1].contiguous() |
| 55 | lm_labels = output_ids[:, 1:].clone() |
| 56 | lm_labels[output_ids[:, 1:] == pad_token_id] = -100 |
| 57 | |
| 58 | outputs = self.bert(input_ids, |
| 59 | attention_mask=attention_mask, decoder_input_ids=y_ids, lm_labels=lm_labels, ) |
| 60 | return (outputs[0].detach(), generated_ids) |
| 61 | |
| 62 | def save_pretrained(self, save_directory): |
| 63 | """ Save a model and its configuration file to a directory, so that it |
| 64 | can be re-loaded using the `:func:`~transformers.PreTrainedModel.from_pretrained`` class method. |
| 65 | |
| 66 | Arguments: |
| 67 | save_directory: directory to which to save. |
| 68 | """ |