Class BARTParser

relogic/pretrainkit/models/semparse/bart_parser.py:11–92

output: tuple: (loss, ) in training

WEIGHTS_NAME = "pytorch_model.bin"

class BARTParser(nn.Module):
  """
  output: tuple: (loss, ) in training
  """
  def __init__(self):
    super().__init__()
    self.bert = BartForTextToSQL.from_pretrained("facebook/bart-large")


  def forward(self, *input, **kwargs):
    input_token_ids = kwargs.pop("input_ids")
    column_spans = kwargs.pop("column_spans")
    input_padding_id = kwargs.pop("input_padding_id")
    # copy_span = kwargs.pop("copy_span", None)
    copy_span = None
    attention_mask = (input_token_ids != input_padding_id).long()
    # relation_ids = None
    if self.training:
      label_ids = kwargs.pop("labels")
      label_padding_id = kwargs.pop("label_padding_id")
      # encoded = self.bert.encoder(input_token_ids)[0].contiguous()
      y_ids = label_ids[:, :-1].contiguous()
      lm_labels = label_ids[:, 1:].clone()
      lm_labels[label_ids[:, 1:] == label_padding_id] = -100
      outputs = self.bert(input_token_ids, column_spans=column_spans, copy_span=copy_span,
                          attention_mask=attention_mask, decoder_input_ids=y_ids, lm_labels=lm_labels, )
      return (outputs[0],)

    else:
      label_eos_id = kwargs.pop("label_eos_id")
      label_bos_id = kwargs.pop("label_bos_id")
      label_padding_id = kwargs.pop("label_padding_id")
      generated_ids = self.bert.generate(
        input_ids=input_token_ids,
        column_spans=column_spans,
        copy_span=copy_span,
        attention_mask=attention_mask,
        num_beams=1,
        max_length=30,
        length_penalty=2.0,
        early_stopping=True,
        use_cache=True,
        decoder_start_token_id=label_bos_id,
        eos_token_id=label_eos_id,
        pad_token_id=label_padding_id,
        vocab_size=len(KEYWORDS)
      )
      # label_ids = kwargs.pop("label_ids")
      # label_padding_id = kwargs.pop("label_padding_id")
      # # encoded = self.bert.encoder(input_token_ids)[0].contiguous()
      # y_ids = label_ids[:, :-1].contiguous()
      # lm_labels = label_ids[:, 1:].clone()
      # lm_labels[label_ids[:, 1:] == label_padding_id] = -100
      # outputs = self.bert(input_token_ids, column_spans=column_spans,
      #                     attention_mask=attention_mask, decoder_input_ids=y_ids, lm_labels=lm_labels, )
      # generated_ids = outputs[-1]
      # raise NotImplementedError()
      return (torch.zeros(1), generated_ids)