constrained_seq2seq.py (8092B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 import logging 4 import os 5 6 import torch 7 import torch.nn as nn 8 from dataclasses import dataclass, field 9 from typing import Union, List, Callable, Dict, Tuple, Any, Optional 10 import numpy as np 11 from torch.cuda.amp import autocast 12 13 from transformers import ( 14 PreTrainedTokenizer, 15 EvalPrediction, 16 Seq2SeqTrainer, 17 Seq2SeqTrainingArguments, 18 ) 19 20 from extraction.event_schema import EventSchema 21 from extraction.extract_constraint import get_constraint_decoder 22 from extraction.extraction_metrics import get_extract_metrics 23 from seq2seq.label_smoother_sum import SumLabelSmoother 24 from seq2seq.utils import lmap 25 26 logger = logging.getLogger(__name__) 27 28 29 def add_logging_file(training_args): 30 fh = logging.FileHandler(os.path.join(training_args.output_dir.rstrip(os.sep) + '.log')) 31 fh.setLevel(logging.DEBUG) 32 logger.addHandler(fh) 33 34 35 def decode_tree_str(sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor"], 36 tokenizer: PreTrainedTokenizer) -> List[str]: 37 def clean_tree_text(x): 38 return x.replace('<pad>', '').replace('<s>', '').replace('</s>', '').strip() 39 40 sequences = np.where(sequences != -100, sequences, tokenizer.pad_token_id) 41 42 str_list = tokenizer.batch_decode(sequences, skip_special_tokens=False) 43 return lmap(clean_tree_text, str_list) 44 45 46 def build_compute_extract_metrics_event_fn(decoding_type_schema: EventSchema, 47 decoding_format: str, 48 tokenizer: PreTrainedTokenizer) -> Callable[[EvalPrediction], Dict]: 49 def non_pad_len(tokens: np.ndarray) -> int: 50 return np.count_nonzero(tokens != tokenizer.pad_token_id) 51 52 def decode_pred(pred: EvalPrediction) -> Tuple[List[str], List[str]]: 53 return decode_tree_str(pred.predictions, tokenizer), decode_tree_str(pred.label_ids, tokenizer) 54 55 def extraction_metrics(pred: EvalPrediction) -> Dict: 56 pred_str, label_str = decode_pred(pred) 57 extraction = get_extract_metrics(pred_lns=pred_str, tgt_lns=label_str, label_constraint=decoding_type_schema, 58 decoding_format=decoding_format) 59 # rouge: Dict = calculate_rouge(pred_str, label_str) 60 summ_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1) 61 extraction.update({"gen_len": summ_len}) 62 # extraction.update( ) 63 return extraction 64 65 compute_metrics_fn = extraction_metrics 66 return compute_metrics_fn 67 68 69 @dataclass 70 class ConstraintSeq2SeqTrainingArguments(Seq2SeqTrainingArguments): 71 """ 72 Parameters: 73 constraint_decoding (:obj:`bool`, `optional`, defaults to :obj:`False`): 74 Whether to use Constraint Decoding 75 structure_weight (:obj:`float`, `optional`, defaults to :obj:`None`): 76 """ 77 constraint_decoding: bool = field(default=False, metadata={"help": "Whether to Constraint Decoding or not."}) 78 label_smoothing_sum: bool = field(default=False, 79 metadata={"help": "Whether to use sum token loss for label smoothing"}) 80 81 82 class ConstraintSeq2SeqTrainer(Seq2SeqTrainer): 83 def __init__(self, decoding_type_schema=None, decoding_format='tree', source_prefix=None, *args, **kwargs): 84 super().__init__(*args, **kwargs) 85 86 self.decoding_format = decoding_format 87 self.decoding_type_schema = decoding_type_schema 88 89 # Label smoothing by sum token loss, different from different Label smootheing 90 if self.args.label_smoothing_sum and self.args.label_smoothing_factor != 0: 91 self.label_smoother = SumLabelSmoother(epsilon=self.args.label_smoothing_factor) 92 print('Using %s' % self.label_smoother) 93 elif self.args.label_smoothing_factor != 0: 94 print('Using %s' % self.label_smoother) 95 else: 96 self.label_smoother = None 97 98 if self.args.constraint_decoding: 99 self.constraint_decoder = get_constraint_decoder(tokenizer=self.tokenizer, 100 type_schema=self.decoding_type_schema, 101 decoding_schema=self.decoding_format, 102 source_prefix=source_prefix) 103 else: 104 self.constraint_decoder = None 105 106 def prediction_step( 107 self, 108 model: nn.Module, 109 inputs: Dict[str, Union[torch.Tensor, Any]], 110 prediction_loss_only: bool, 111 ignore_keys: Optional[List[str]] = None, 112 ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: 113 """ 114 Perform an evaluation step on :obj:`model` using obj:`inputs`. 115 116 Subclass and override to inject custom behavior. 117 118 Args: 119 model (:obj:`nn.Module`): 120 The model to evaluate. 121 inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`): 122 The inputs and targets of the model. 123 124 The dictionary will be unpacked before being fed to the model. Most models expect the targets under the 125 argument :obj:`labels`. Check your model's documentation for all accepted arguments. 126 prediction_loss_only (:obj:`bool`): 127 Whether or not to return the loss only. 128 129 Return: 130 Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and 131 labels (each being optional). 132 """ 133 134 def prefix_allowed_tokens_fn(batch_id, sent): 135 # print(self.tokenizer.convert_ids_to_tokens(inputs['labels'][batch_id])) 136 src_sentence = inputs['input_ids'][batch_id] 137 return self.constraint_decoder.constraint_decoding(src_sentence=src_sentence, 138 tgt_generated=sent) 139 140 if not self.args.predict_with_generate or prediction_loss_only: 141 return super().prediction_step( 142 model=model, 143 inputs=inputs, 144 prediction_loss_only=prediction_loss_only, 145 ignore_keys=ignore_keys, 146 prefix_allowed_tokens_fn=prefix_allowed_tokens_fn if self.constraint_decoder else None, 147 ) 148 149 has_labels = "labels" in inputs 150 inputs = self._prepare_inputs(inputs) 151 152 gen_kwargs = { 153 "max_length": self._max_length if self._max_length is not None else self.model.config.max_length, 154 "num_beams": self._num_beams if self._num_beams is not None else self.model.config.num_beams, 155 "prefix_allowed_tokens_fn": prefix_allowed_tokens_fn if self.constraint_decoder else None, 156 } 157 158 generated_tokens = self.model.generate( 159 inputs["input_ids"], 160 attention_mask=inputs["attention_mask"], 161 **gen_kwargs, 162 ) 163 164 # in case the batch is shorter than max length, the output should be padded 165 if generated_tokens.shape[-1] < gen_kwargs["max_length"]: 166 generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"]) 167 168 with torch.no_grad(): 169 if self.use_amp: 170 with autocast(): 171 outputs = model(**inputs) 172 else: 173 outputs = model(**inputs) 174 if has_labels: 175 if self.label_smoother is not None: 176 loss = self.label_smoother(outputs, inputs["labels"]).mean().detach() 177 else: 178 loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach() 179 else: 180 loss = None 181 182 if self.args.prediction_loss_only: 183 return loss, None, None 184 185 labels = inputs["labels"] 186 if labels.shape[-1] < gen_kwargs["max_length"]: 187 labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"]) 188 189 return loss, generated_tokens, labels 190 191 192 def main(): pass 193 194 195 if __name__ == "__main__": 196 main()