def global_distributed_rouge( global_rank: int, world_size: int, gpu_id: Union[int, torch.device], metric: ROUGE, metric_kwargs: Dict[str, Any], desired_values: Dict[str, Any], ): kwargs = {} # Use the arguments meant for the process with rank `global_rank`. for argname in metric_kwargs: kwargs[argname] = metric_kwargs[argname][global_rank] metric(**kwargs) metrics = metric.get_metric() # Unigram unigram_recall = metric._total_rouge_n_recalls[1] assert_allclose(unigram_recall, desired_values["unigram_recall"]) unigram_precision = metric._total_rouge_n_precisions[1] assert_allclose(unigram_precision, desired_values["unigram_precision"]) unigram_f1 = metric._total_rouge_n_f1s[1] assert_allclose(unigram_f1, desired_values["unigram_f1"]) assert metrics[ "ROUGE-1_R"] == unigram_recall / metric._total_sequence_count assert metrics[ "ROUGE-1_P"] == unigram_precision / metric._total_sequence_count assert metrics["ROUGE-1_F1"] == unigram_f1 / metric._total_sequence_count # Bigram bigram_recall = metric._total_rouge_n_recalls[2] assert_allclose(bigram_recall, desired_values["bigram_recall"]) bigram_precision = metric._total_rouge_n_precisions[2] assert_allclose(bigram_precision, desired_values["bigram_precision"]) bigram_f1 = metric._total_rouge_n_f1s[2] assert_allclose(bigram_f1, desired_values["bigram_f1"]) assert metrics["ROUGE-2_R"] == bigram_recall / metric._total_sequence_count assert metrics[ "ROUGE-2_P"] == bigram_precision / metric._total_sequence_count assert metrics["ROUGE-2_F1"] == bigram_f1 / metric._total_sequence_count # ROUGE-L assert_allclose(metric._total_rouge_l_f1, desired_values["total_rouge_l_f1"]) assert metrics[ "ROUGE-L"] == metric._total_rouge_l_f1 / metric._total_sequence_count
class Bart(Model): """ BART model from the paper "BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension" (https://arxiv.org/abs/1910.13461). The Bart model here uses a language modeling head and thus can be used for text generation. """ def __init__( self, model_name: str, vocab: Vocabulary, indexer: PretrainedTransformerIndexer = None, max_decoding_steps: int = 140, beam_size: int = 4, encoder: Seq2SeqEncoder = None, ): """ # Parameters model_name : `str`, required Name of the pre-trained BART model to use. Available options can be found in `transformers.models.bart.modeling_bart.BART_PRETRAINED_MODEL_ARCHIVE_MAP`. vocab : `Vocabulary`, required Vocabulary containing source and target vocabularies. indexer : `PretrainedTransformerIndexer`, optional (default = `None`) Indexer to be used for converting decoded sequences of ids to to sequences of tokens. max_decoding_steps : `int`, optional (default = `128`) Number of decoding steps during beam search. beam_size : `int`, optional (default = `5`) Number of beams to use in beam search. The default is from the BART paper. encoder : `Seq2SeqEncoder`, optional (default = `None`) Encoder to used in BART. By default, the original BART encoder is used. """ super().__init__(vocab) self.bart = BartForConditionalGeneration.from_pretrained(model_name) self._indexer = indexer or PretrainedTransformerIndexer( model_name, namespace="tokens") self._start_id = self.bart.config.bos_token_id # CLS self._decoder_start_id = self.bart.config.decoder_start_token_id or self._start_id self._end_id = self.bart.config.eos_token_id # SEP self._pad_id = self.bart.config.pad_token_id # PAD self._max_decoding_steps = max_decoding_steps self._beam_search = BeamSearch(self._end_id, max_steps=max_decoding_steps, beam_size=beam_size or 1) self._rouge = ROUGE( exclude_indices={self._start_id, self._pad_id, self._end_id}) self._bleu = BLEU( exclude_indices={self._start_id, self._pad_id, self._end_id}) # Replace bart encoder with given encoder. We need to extract the two embedding layers so that # we can use them in the encoder wrapper if encoder is not None: assert (encoder.get_input_dim() == encoder.get_output_dim() == self.bart.config.hidden_size) self.bart.model.encoder = _BartEncoderWrapper( encoder, self.bart.model.encoder.embed_tokens, self.bart.model.encoder.embed_positions, ) @overrides def forward( self, source_tokens: TextFieldTensors, target_tokens: TextFieldTensors = None) -> Dict[str, torch.Tensor]: """ Performs the forward step of Bart. # Parameters source_tokens : `TextFieldTensors`, required The source tokens for the encoder. We assume they are stored under the `tokens` key. target_tokens : `TextFieldTensors`, optional (default = `None`) The target tokens for the decoder. We assume they are stored under the `tokens` key. If no target tokens are given, the source tokens are shifted to the right by 1. # Returns `Dict[str, torch.Tensor]` During training, this dictionary contains the `decoder_logits` of shape `(batch_size, max_target_length, target_vocab_size)` and the `loss`. During inference, it contains `predictions` of shape `(batch_size, max_decoding_steps)` and `log_probabilities` of shape `(batch_size,)`. """ inputs = source_tokens targets = target_tokens input_ids, input_mask = inputs["tokens"]["token_ids"], inputs[ "tokens"]["mask"] outputs = {} # If no targets are provided, then shift input to right by 1. Bart already does this internally # but it does not use them for loss calculation. if targets is not None: target_ids, target_mask = targets["tokens"]["token_ids"], targets[ "tokens"]["mask"] else: target_ids = input_ids[:, 1:] target_mask = input_mask[:, 1:] if self.training: decoder_logits = self.bart( input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=target_ids[:, :-1].contiguous(), decoder_attention_mask=target_mask[:, :-1].contiguous(), use_cache=False, )[0] outputs["decoder_logits"] = decoder_logits # The BART paper mentions label smoothing of 0.1 for sequence generation tasks outputs["loss"] = sequence_cross_entropy_with_logits( decoder_logits, target_ids[:, 1:].contiguous(), target_mask[:, 1:].contiguous(), label_smoothing=0.1, average="token", ) else: # Use decoder start id and start of sentence to start decoder initial_decoder_ids = torch.tensor( [[self._decoder_start_id, self._start_id]], dtype=input_ids.dtype, device=input_ids.device, ).repeat(input_ids.shape[0], 1) inital_state = { "input_ids": input_ids, "input_mask": input_mask, "encoder_states": None, } beam_result = self._beam_search.search(initial_decoder_ids, inital_state, self.take_step) predictions = beam_result[0] max_pred_indices = (beam_result[1].argmax(dim=-1).view( -1, 1, 1).expand(-1, -1, predictions.shape[-1])) predictions = predictions.gather( dim=1, index=max_pred_indices).squeeze(dim=1) self._rouge(predictions, target_ids) self._bleu(predictions, target_ids) outputs["predictions"] = predictions outputs["log_probabilities"] = (beam_result[1].gather( dim=-1, index=max_pred_indices[..., 0]).squeeze(dim=-1)) self.make_output_human_readable(outputs) return outputs @staticmethod def _decoder_cache_to_dict(decoder_cache): cache_dict = {} for layer_index, layer_cache in enumerate(decoder_cache): for attention_name, attention_cache in layer_cache.items(): for tensor_name, cache_value in attention_cache.items(): key = (layer_index, attention_name, tensor_name) cache_dict[key] = cache_value return cache_dict @staticmethod def _dict_to_decoder_cache(cache_dict): decoder_cache = [] for key, cache_value in cache_dict.items(): # Split key and extract index and dict keys layer_idx, attention_name, tensor_name = key # Extend decoder_cache to fit layer_idx + 1 layers decoder_cache = decoder_cache + [ {} for _ in range(layer_idx + 1 - len(decoder_cache)) ] cache = decoder_cache[layer_idx] if attention_name not in cache: cache[attention_name] = {} assert tensor_name not in cache[attention_name] cache[attention_name][tensor_name] = cache_value return decoder_cache def take_step(self, last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """ Take step during beam search. # Parameters last_predictions : `torch.Tensor` The predicted token ids from the previous step. Shape: `(group_size,)` state : `Dict[str, torch.Tensor]` State required to generate next set of predictions step : `int` The time step in beam search decoding. # Returns `Tuple[torch.Tensor, Dict[str, torch.Tensor]]` A tuple containing logits for the next tokens of shape `(group_size, target_vocab_size)` and an updated state dictionary. """ if len(last_predictions.shape) == 1: last_predictions = last_predictions.unsqueeze(-1) # Only the last predictions are needed for the decoder, but we need to pad the decoder ids # to not mess up the positional embeddings in the decoder. padding_size = 0 if step > 0: padding_size = step + 1 padding = torch.full( (last_predictions.shape[0], padding_size), self._pad_id, dtype=last_predictions.dtype, device=last_predictions.device, ) last_predictions = torch.cat([padding, last_predictions], dim=-1) decoder_cache = None decoder_cache_dict = { k: (state[k].contiguous() if state[k] is not None else None) for k in state if k not in {"input_ids", "input_mask", "encoder_states"} } if len(decoder_cache_dict) != 0: decoder_cache = self._dict_to_decoder_cache(decoder_cache_dict) log_probabilities = None for i in range(padding_size, last_predictions.shape[1]): encoder_outputs = ((state["encoder_states"], ) if state["encoder_states"] is not None else None) outputs = self.bart( input_ids=state["input_ids"], attention_mask=state["input_mask"], encoder_outputs=encoder_outputs, decoder_input_ids=last_predictions[:, :i + 1], past_key_values=decoder_cache, use_cache=True, ) decoder_log_probabilities = F.log_softmax(outputs[0][:, 0], dim=-1) if log_probabilities is None: log_probabilities = decoder_log_probabilities else: idx = last_predictions[:, i].view(-1, 1) log_probabilities = decoder_log_probabilities + log_probabilities.gather( dim=-1, index=idx) decoder_cache = outputs[1] state["encoder_states"] = outputs[2] if decoder_cache is not None: decoder_cache_dict = self._decoder_cache_to_dict(decoder_cache) state.update(decoder_cache_dict) return log_probabilities, state @overrides def make_output_human_readable( self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, Any]: """ # Parameters output_dict : `Dict[str, torch.Tensor]` A dictionary containing a batch of predictions with key `predictions`. The tensor should have shape `(batch_size, max_sequence_length)` # Returns `Dict[str, Any]` Original `output_dict` with an additional `predicted_tokens` key that maps to a list of lists of tokens. """ predictions = output_dict["predictions"] predicted_tokens = [None] * predictions.shape[0] for i in range(predictions.shape[0]): predicted_tokens[i] = self._indexer.indices_to_tokens( {"token_ids": predictions[i].tolist()}, self.vocab) output_dict["predicted_tokens"] = predicted_tokens return output_dict @overrides def get_metrics(self, reset: bool = False) -> Dict[str, float]: metrics: Dict[str, float] = {} if not self.training: metrics.update(self._rouge.get_metric(reset=reset)) metrics.update(self._bleu.get_metric(reset=reset)) return metrics
class Seq2seqPlmsGenerator(Model): def __init__(self, vocab: Vocabulary, pretrained_model_path, beam_size=5, max_decoding_steps=140, indexer=None): super().__init__(vocab) self.plm = MT5ForConditionalGeneration.from_pretrained(pretrained_model_path) self._indexer = indexer or PretrainedTransformerIndexer(pretrained_model_path, namespace="tokens") ## self._start_id = self.plm.config.decoder_start_token_id ## self._end_id = self.plm.config.eos_token_id # self._decoder_start_id = self.plm.config.decoder_start_token_id self._end_id = self.plm.config.eos_token_id # self._pad_id = self.plm.config.pad_token_id # self._beam_search = BeamSearch( self._end_id, max_steps=max_decoding_steps, beam_size=beam_size or 1 ) self._rouge = ROUGE(exclude_indices={self._start_id, self._pad_id, self._end_id}) self._bleu = BLEU(exclude_indices={self._start_id, self._pad_id, self._end_id}) @overrides def forward(self, source_tokens, target_tokens=None) -> Dict[str, torch.Tensor]: inputs = source_tokens targets = target_tokens input_ids, input_mask = inputs["tokens"]["token_ids"], inputs["tokens"]["mask"] outputs = {} # If no targets are provided, then shift input to right by 1. Bart already does this internally # but it does not use them for loss calculation. if targets is not None: target_ids, target_mask = targets["tokens"]["token_ids"], targets["tokens"]["mask"] else: target_ids = input_ids[:, 1:] target_mask = input_mask[:, 1:] if self.training: # training outputs = self.plm(input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=target_ids[:, :-1].contiguous(), decoder_attention_mask=target_mask[:, :-1].contiguous(), use_cache=False, return_dict=True) outputs["decoder_logits"] = outputs.logits outputs["loss"] = sequence_cross_entropy_with_logits( outputs.logits, cast(torch.LongTensor, target_ids[:, 1:].contiguous()), cast(torch.BoolTensor, target_mask[:, 1:].contiguous()), label_smoothing=0.1, average="token", ) elif targets is not None: # validation outputs = self.plm(input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=target_ids[:, :-1].contiguous(), decoder_attention_mask=target_mask[:, :-1].contiguous(), use_cache=False, return_dict=True) outputs["decoder_logits"] = outputs.logits outputs["loss"] = sequence_cross_entropy_with_logits( outputs.logits, cast(torch.LongTensor, target_ids[:, 1:].contiguous()), cast(torch.BoolTensor, target_mask[:, 1:].contiguous()), label_smoothing=0.1, ) self._rouge(torch.argmax(outputs.logits, -1), target_ids) self._bleu(torch.argmax(outputs.logits, -1), target_ids) else: #prediction # Use decoder start id and start of sentence to start decoder initial_decoder_ids = torch.tensor( [[self._decoder_start_id]], dtype=input_ids.dtype, device=input_ids.device, ).repeat(input_ids.shape[0], 1) inital_state = { "input_ids": input_ids, "input_mask": input_mask, } beam_result = self._beam_search.search( initial_decoder_ids, inital_state, self.take_step ) predictions = beam_result[0] logger.info(beam_result) max_pred_indices = ( beam_result[1].argmax(dim=-1).view(-1, 1, 1).expand(-1, -1, predictions.shape[-1]) ) predictions = predictions.gather(dim=1, index=max_pred_indices).squeeze(dim=1) self._rouge(predictions, target_ids) self._bleu(predictions, target_ids) outputs["predictions"] = predictions outputs["log_probabilities"] = ( beam_result[1].gather(dim=-1, index=max_pred_indices[..., 0]).squeeze(dim=-1) ) self.make_output_human_readable(outputs) return outputs @overrides def get_metrics(self, reset: bool = False) -> Dict[str, float]: metrics: Dict[str, float] = {} if not self.training: metrics.update(self._rouge.get_metric(reset=reset)) metrics.update(self._bleu.get_metric(reset=reset)) return metrics @staticmethod def _decoder_cache_to_dict(decoder_cache: DecoderCacheType) -> Dict[str, torch.Tensor]: cache_dict = {} for layer_index, layer_cache in enumerate(decoder_cache): # Each layer caches the key and value tensors for its self-attention and cross-attention. # Hence the `layer_cache` tuple has 4 elements. assert len(layer_cache) == 4 for tensor_index, tensor in enumerate(layer_cache): key = f"decoder_cache_{layer_index}_{tensor_index}" cache_dict[key] = tensor return cache_dict def _dict_to_decoder_cache(self, cache_dict: Dict[str, torch.Tensor]) -> DecoderCacheType: decoder_cache = [] for layer_index in range(self.plm.config.num_layers): base_key = f"decoder_cache_{layer_index}_" layer_cache = ( cache_dict[base_key + "0"], cache_dict[base_key + "1"], cache_dict[base_key + "2"], cache_dict[base_key + "3"], ) decoder_cache.append(layer_cache) assert decoder_cache return tuple(decoder_cache) def take_step( self, last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], step: int ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """ Take step during beam search. # Parameters last_predictions : `torch.Tensor` The predicted token ids from the previous step. Shape: `(group_size,)` state : `Dict[str, torch.Tensor]` State required to generate next set of predictions step : `int` The time step in beam search decoding. # Returns `Tuple[torch.Tensor, Dict[str, torch.Tensor]]` A tuple containing logits for the next tokens of shape `(group_size, target_vocab_size)` and an updated state dictionary. """ if len(last_predictions.shape) == 1: last_predictions = last_predictions.unsqueeze(-1) decoder_cache = None decoder_cache_dict = { k: state[k].contiguous() for k in state if k not in {"input_ids", "input_mask", "encoder_states"} } if len(decoder_cache_dict) != 0: decoder_cache = self._dict_to_decoder_cache(decoder_cache_dict) encoder_outputs = (state["encoder_states"],) if "encoder_states" in state else None outputs = self.plm( input_ids=state["input_ids"] if encoder_outputs is None else None, attention_mask=state["input_mask"], encoder_outputs=encoder_outputs, decoder_input_ids=last_predictions, past_key_values=decoder_cache, use_cache=True, return_dict=True, ) logits = outputs.logits[:, -1, :] log_probabilities = F.log_softmax(logits, dim=-1) decoder_cache = outputs.past_key_values if decoder_cache is not None: decoder_cache_dict = self._decoder_cache_to_dict(decoder_cache) state.update(decoder_cache_dict) state["encoder_states"] = outputs.encoder_last_hidden_state return log_probabilities, state @overrides def make_output_human_readable(self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, Any]: """ # Parameters output_dict : `Dict[str, torch.Tensor]` A dictionary containing a batch of predictions with key `predictions`. The tensor should have shape `(batch_size, max_sequence_length)` # Returns `Dict[str, Any]` Original `output_dict` with an additional `predicted_tokens` key that maps to a list of lists of tokens. """ predictions = output_dict["predictions"] predicted_tokens = [None] * predictions.shape[0] for i in range(predictions.shape[0]): predicted_tokens[i] = self._indexer.indices_to_tokens( {"token_ids": predictions[i].tolist()}, self.vocab, ) output_dict["predicted_tokens"] = predicted_tokens # type: ignore output_dict["predicted_text"] = self._indexer._tokenizer.batch_decode( predictions.tolist(), skip_special_tokens=True ) return output_dict
class RougeTest(AllenNlpTestCase): def setup_method(self): super().setup_method() self.metric = ROUGE(exclude_indices={0}) def f1(self, r, p): if r == p == 0: return 0 return 2 * r * p / (r + p) @multi_device def test_rouge(self, device: str): self.metric.reset() predictions = torch.tensor([[1, 0, 1, 2], [1, 0, 3, 0], [1, 2, 3, 0]], device=device) targets = torch.tensor([[2, 0, 1, 2], [1, 2, 1, 0], [1, 0, 2, 3]], device=device) self.metric(predictions, targets) metrics = self.metric.get_metric() assert self.metric._total_sequence_count == 3 # ROUGE-N # Unigram unigram_recall = self.metric._total_rouge_n_recalls[1] assert unigram_recall == 2 / 3 + 1 / 3 + 3 / 3 unigram_precision = self.metric._total_rouge_n_precisions[1] assert unigram_precision == 2 / 3 + 1 / 2 + 3 / 3 unigram_f1 = self.metric._total_rouge_n_f1s[1] assert unigram_f1 == self.f1(2 / 3, 2 / 3) + self.f1( 1 / 2, 1 / 3) + self.f1(3 / 3, 3 / 3) assert metrics[ "ROUGE-1_R"] == unigram_recall / self.metric._total_sequence_count assert metrics[ "ROUGE-1_P"] == unigram_precision / self.metric._total_sequence_count assert metrics[ "ROUGE-1_F1"] == unigram_f1 / self.metric._total_sequence_count # Bigram bigram_recall = self.metric._total_rouge_n_recalls[2] assert bigram_recall == 1 / 1 + 0 / 2 + 1 / 1 bigram_precision = self.metric._total_rouge_n_precisions[2] assert bigram_precision == 1 / 1 + 0 + 1 / 2 bigram_f1 = self.metric._total_rouge_n_f1s[2] assert bigram_f1 == self.f1(1 / 1, 1 / 1) + self.f1( 0, 0 / 2) + self.f1(1 / 2, 1 / 1) assert metrics[ "ROUGE-2_R"] == bigram_recall / self.metric._total_sequence_count assert metrics[ "ROUGE-2_P"] == bigram_precision / self.metric._total_sequence_count assert metrics[ "ROUGE-2_F1"] == bigram_f1 / self.metric._total_sequence_count # ROUGE-L assert self.metric._total_rouge_l_f1 == self.f1( 2 / 3, 2 / 3) + self.f1(1 / 3, 1 / 2) + self.f1(3 / 3, 3 / 3) assert (metrics["ROUGE-L"] == self.metric._total_rouge_l_f1 / self.metric._total_sequence_count) def test_rouge_with_zero_counts(self): self.metric.reset() metrics = self.metric.get_metric() for score in metrics.values(): assert score == 0.0
class Bart(Model): """ BART model from the paper "BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension" (https://arxiv.org/abs/1910.13461). The Bart model here uses a language modeling head and thus can be used for text generation. # Parameters model_name : `str`, required Name of the pre-trained BART model to use. Available options can be found in `transformers.models.bart.modeling_bart.BART_PRETRAINED_MODEL_ARCHIVE_MAP`. vocab : `Vocabulary`, required Vocabulary containing source and target vocabularies. beam_search : `Lazy[BeamSearch]`, optional (default = `Lazy(BeamSearch)`) This is used to during inference to select the tokens of the decoded output sequence. indexer : `PretrainedTransformerIndexer`, optional (default = `None`) Indexer to be used for converting decoded sequences of ids to to sequences of tokens. encoder : `Seq2SeqEncoder`, optional (default = `None`) Encoder to used in BART. By default, the original BART encoder is used. """ def __init__( self, model_name: str, vocab: Vocabulary, beam_search: Lazy[BeamSearch] = Lazy(BeamSearch), indexer: PretrainedTransformerIndexer = None, encoder: Seq2SeqEncoder = None, **kwargs, ): super().__init__(vocab) self.bart = BartForConditionalGeneration.from_pretrained(model_name) self._indexer = indexer or PretrainedTransformerIndexer(model_name, namespace="tokens") self._start_id = self.bart.config.bos_token_id # CLS self._decoder_start_id = self.bart.config.decoder_start_token_id or self._start_id self._end_id = self.bart.config.eos_token_id # SEP self._pad_id = self.bart.config.pad_token_id # PAD # At prediction time, we'll use a beam search to find the best target sequence. # For backwards compatibility, check if beam_size or max_decoding_steps were passed in as # kwargs. If so, update the BeamSearch object before constructing and raise a DeprecationWarning deprecation_warning = ( "The parameter {} has been deprecated." " Provide this parameter as argument to beam_search instead." ) beam_search_extras = {} if "beam_size" in kwargs: beam_search_extras["beam_size"] = kwargs["beam_size"] warnings.warn(deprecation_warning.format("beam_size"), DeprecationWarning) if "max_decoding_steps" in kwargs: beam_search_extras["max_steps"] = kwargs["max_decoding_steps"] warnings.warn(deprecation_warning.format("max_decoding_steps"), DeprecationWarning) self._beam_search = beam_search.construct( end_index=self._end_id, vocab=self.vocab, **beam_search_extras ) self._rouge = ROUGE(exclude_indices={self._start_id, self._pad_id, self._end_id}) self._bleu = BLEU(exclude_indices={self._start_id, self._pad_id, self._end_id}) # Replace bart encoder with given encoder. We need to extract the two embedding layers so that # we can use them in the encoder wrapper if encoder is not None: assert ( encoder.get_input_dim() == encoder.get_output_dim() == self.bart.config.hidden_size ) self.bart.model.encoder = _BartEncoderWrapper( encoder, self.bart.model.encoder.embed_tokens, self.bart.model.encoder.embed_positions, ) def forward( self, source_tokens: TextFieldTensors, target_tokens: TextFieldTensors = None ) -> Dict[str, torch.Tensor]: """ Performs the forward step of Bart. # Parameters source_tokens : `TextFieldTensors`, required The source tokens for the encoder. We assume they are stored under the `tokens` key. target_tokens : `TextFieldTensors`, optional (default = `None`) The target tokens for the decoder. We assume they are stored under the `tokens` key. If no target tokens are given, the source tokens are shifted to the right by 1. # Returns `Dict[str, torch.Tensor]` During training, this dictionary contains the `decoder_logits` of shape `(batch_size, max_target_length, target_vocab_size)` and the `loss`. During inference, it contains `predictions` of shape `(batch_size, max_decoding_steps)` and `log_probabilities` of shape `(batch_size,)`. """ inputs = source_tokens targets = target_tokens input_ids, input_mask = inputs["tokens"]["token_ids"], inputs["tokens"]["mask"] outputs = {} # If no targets are provided, then shift input to right by 1. Bart already does this internally # but it does not use them for loss calculation. if targets is not None: target_ids, target_mask = targets["tokens"]["token_ids"], targets["tokens"]["mask"] else: target_ids = input_ids[:, 1:] target_mask = input_mask[:, 1:] if self.training: bart_outputs = self.bart( input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=target_ids[:, :-1].contiguous(), decoder_attention_mask=target_mask[:, :-1].contiguous(), use_cache=False, return_dict=True, ) outputs["decoder_logits"] = bart_outputs.logits # The BART paper mentions label smoothing of 0.1 for sequence generation tasks outputs["loss"] = sequence_cross_entropy_with_logits( bart_outputs.logits, cast(torch.LongTensor, target_ids[:, 1:].contiguous()), cast(torch.BoolTensor, target_mask[:, 1:].contiguous()), label_smoothing=0.1, average="token", ) else: # Use decoder start id and start of sentence to start decoder initial_decoder_ids = torch.tensor( [[self._decoder_start_id]], dtype=input_ids.dtype, device=input_ids.device, ).repeat(input_ids.shape[0], 1) inital_state = { "input_ids": input_ids, "input_mask": input_mask, } beam_result = self._beam_search.search( initial_decoder_ids, inital_state, self.take_step ) predictions = beam_result[0] max_pred_indices = ( beam_result[1].argmax(dim=-1).view(-1, 1, 1).expand(-1, -1, predictions.shape[-1]) ) predictions = predictions.gather(dim=1, index=max_pred_indices).squeeze(dim=1) self._rouge(predictions, target_ids) self._bleu(predictions, target_ids) outputs["predictions"] = predictions outputs["log_probabilities"] = ( beam_result[1].gather(dim=-1, index=max_pred_indices[..., 0]).squeeze(dim=-1) ) self.make_output_human_readable(outputs) return outputs @staticmethod def _decoder_cache_to_dict(decoder_cache: DecoderCacheType) -> Dict[str, torch.Tensor]: cache_dict = {} for layer_index, layer_cache in enumerate(decoder_cache): # Each layer caches the key and value tensors for its self-attention and cross-attention. # Hence the `layer_cache` tuple has 4 elements. assert len(layer_cache) == 4 for tensor_index, tensor in enumerate(layer_cache): key = f"decoder_cache_{layer_index}_{tensor_index}" cache_dict[key] = tensor return cache_dict def _dict_to_decoder_cache(self, cache_dict: Dict[str, torch.Tensor]) -> DecoderCacheType: decoder_cache = [] for layer_index in range(len(self.bart.model.decoder.layers)): base_key = f"decoder_cache_{layer_index}_" layer_cache = ( cache_dict[base_key + "0"], cache_dict[base_key + "1"], cache_dict[base_key + "2"], cache_dict[base_key + "3"], ) decoder_cache.append(layer_cache) assert decoder_cache return tuple(decoder_cache) def take_step( self, last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], step: int ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """ Take step during beam search. # Parameters last_predictions : `torch.Tensor` The predicted token ids from the previous step. Shape: `(group_size,)` state : `Dict[str, torch.Tensor]` State required to generate next set of predictions step : `int` The time step in beam search decoding. # Returns `Tuple[torch.Tensor, Dict[str, torch.Tensor]]` A tuple containing logits for the next tokens of shape `(group_size, target_vocab_size)` and an updated state dictionary. """ if len(last_predictions.shape) == 1: last_predictions = last_predictions.unsqueeze(-1) decoder_cache = None decoder_cache_dict = { k: state[k].contiguous() for k in state if k not in {"input_ids", "input_mask", "encoder_states"} } if len(decoder_cache_dict) != 0: decoder_cache = self._dict_to_decoder_cache(decoder_cache_dict) encoder_outputs = (state["encoder_states"],) if "encoder_states" in state else None outputs = self.bart( input_ids=state["input_ids"] if encoder_outputs is None else None, attention_mask=state["input_mask"], encoder_outputs=encoder_outputs, decoder_input_ids=last_predictions, past_key_values=decoder_cache, use_cache=True, return_dict=True, ) logits = outputs.logits[:, -1, :] log_probabilities = F.log_softmax(logits, dim=-1) decoder_cache = outputs.past_key_values if decoder_cache is not None: decoder_cache_dict = self._decoder_cache_to_dict(decoder_cache) state.update(decoder_cache_dict) state["encoder_states"] = outputs.encoder_last_hidden_state return log_probabilities, state def make_output_human_readable(self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, Any]: """ # Parameters output_dict : `Dict[str, torch.Tensor]` A dictionary containing a batch of predictions with key `predictions`. The tensor should have shape `(batch_size, max_sequence_length)` # Returns `Dict[str, Any]` Original `output_dict` with an additional `predicted_tokens` key that maps to a list of lists of tokens. """ predictions = output_dict["predictions"] predicted_tokens = [None] * predictions.shape[0] for i in range(predictions.shape[0]): predicted_tokens[i] = self._indexer.indices_to_tokens( {"token_ids": predictions[i].tolist()}, self.vocab, ) output_dict["predicted_tokens"] = predicted_tokens # type: ignore output_dict["predicted_text"] = self._indexer._tokenizer.batch_decode( predictions.tolist(), skip_special_tokens=True ) return output_dict def get_metrics(self, reset: bool = False) -> Dict[str, float]: metrics: Dict[str, float] = {} if not self.training: metrics.update(self._rouge.get_metric(reset=reset)) metrics.update(self._bleu.get_metric(reset=reset)) return metrics default_predictor = "seq2seq"