From b35aa4c1e19bcce3174d308040e52fe9ecbdca90 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 9 Jan 2022 12:11:36 +0000 Subject: [PATCH 01/34] added classes to get started with constrained beam search --- .../generation_beam_constraints.py | 194 ++++++++++ src/transformers/generation_beam_search.py | 1 + src/transformers/generation_utils.py | 347 +++++++++++++++++- tests/test_generation_beam_constraints.py | 250 +++++++++++++ 4 files changed, 791 insertions(+), 1 deletion(-) create mode 100644 src/transformers/generation_beam_constraints.py create mode 100644 tests/test_generation_beam_constraints.py diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py new file mode 100644 index 000000000000..03d078a41ab9 --- /dev/null +++ b/src/transformers/generation_beam_constraints.py @@ -0,0 +1,194 @@ +from abc import ABC + +from collections import Counter +from typing import List, Optional, Set, Tuple + +import torch + +from .file_utils import add_start_docstrings +from .utils.logging import get_logger + + +logger = get_logger(__name__) + + + +class Constraint(ABC): + r"""Abstract base class for all constraints that can be applied during generation. + It must define how the constraint can be satisfied. + + All classes that inherit Constraint must follow the requirement that + + ``` + completed = False + while(not completed): + _, completed = constraint.update(constraint.advance()) + ``` + + will always terminate (halt). + + """ + def advance(self): + ''' + When called, returns the token that would take this constraint + one step closer to being fulfilled. + ''' + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + def does_advance(self, token_id: int): + """ + Reads in a token and returns whether it creates progress. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + def update(self, token_id: int): + """ + Reads in a token and returns booleans that indicate the progress made by it. + This function will update the state of this object unlikes `does_advance(self, token_id: int)`. + + This function assumes that token_id is sure to be generated. This isn't to test whether + a certain token will advnace the progress; so we update the states accordingly + if it's already been generated. This becomes important if token_id != desired token (refer to else statement in PhrasalConstraint) + + Args: + token_id(`int`): + The id of a newly generated token in the beam search. + returns: + stepped(`boolean`): + Whether this constraint has become one step closer to being fulfuilled. + completed(`stepped`): + Whether this constraint has been completely fulfilled by this token being generated. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + +class TokenConstraint(Constraint): + r""" + [`Constraint`] enforcing that a specific token is generated. + + Args: + token_id (`int`): + The token that must be generated by the output. + """ + def __init__(self, token_id: int): + if not isinstance(token_id, int) or token_id < 0: + raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") + self.token_id = token_id + + def advance(self): + return self.token_id + + def does_advance(self, token_id: int): + return token_id == self.token_id + + def update(self, token_id: int): + if not isinstance(token_id, int) or token_id < 0: + raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") + + if self.does_advance(token_id): + return True, True # stepped, completed + else: + return False, False + + +class PhrasalConstraint(Constraint): + r""" + [`Constraint`] enforcing that an ordered sequence of tokens is generated. + + Args: + token_ids (`torch.Tensor`): + The sequence of tokens that must be generated by the output. + """ + def __init__(self, token_ids: torch.Tensor): + if not isinstance(token_ids, torch.Tensor): + raise ValueError(f"`token_ids` has to be a tensor, but is {type(token_ids)}") + self.token_ids = token_ids + + # the index of the currently fulfilled step + self.seqlen = self.token_ids.size(0) + self.fulfilled_idx = -1 + self.completed = False + + def advance(self): + return self.token_ids[self.fulfilled_idx + 1] + + def does_advance(self, token_id): + return token_id == self.token_ids[self.fulfilled_idx + 1] + + def update(self, token_id: int): + + stepped = False + completed = False + reset = False + + if self.does_advance(token_id): + self.fulfilled_idx += 1 + stepped = True + if self.fulfilled_idx == (self.seqlen - 1): + completed = True + self.completed = completed + else: + # failed to make progress. + reset = True + self.reset() + + return stepped, completed + + def reset(self): + self.completed = False + self.fulfilled_idx = 0 + + +class ConstraintListState: + def __init__(self, constraints: List[Constraint]): + self.complete_constraints = [] + self.inprogress_constraint = None + self.pending_constraints = constraints + + def advance(self): + '''The list of tokens to generate such that we can make progress. + + Though we don't care which constraint is fulfilled first, + if we are in the progress of fulfilling a constraint, that's the only one + we'll return. + ''' + if self.inprogress_constraint is None: + token_list = [ + constraint.advance() for constraint in self.pending_constraints + ] + + else: + token_list = [self.inprogress_constraint.advance()] + + return token_list + + + def update(self, token_id: int): + if self.inprogress_constraint is None: + stepped, complete, reset = self.inprogress_constraint.update(token_id) + if reset: + self.pending_constraints.append(self.inprogress_constraint) + self.inprogress_constraint = None + if complete: + self.complete_constraints.append(self.inprogress_constraint) + self.inprogress_constraint = None + else: + for cidx, pending_constraint in enumerate(self.pending_constraints): + if pending_constraint.does_advance(token_id): + stepped, complete, reset = pending_constraint.update(token_id) + if complete: + self.complete_constraints.append(pending_constraint) + + elif stepped: + self.inprogress_constraint = pending_constraint + + if complete or stepped: + self.pending_constraints = pending_constraints[:idx] + pending_constraints[idx+1:] + break + + \ No newline at end of file diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 3c4f259b0047..582560c3cde3 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -395,3 +395,4 @@ def is_done(self, best_sum_logprobs: float, cur_len: int) -> bool: cur_score = best_sum_logprobs / cur_len ** self.length_penalty ret = self.worst_score >= cur_score return ret + diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 1906235ae1cb..daa8f5831cf4 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -24,7 +24,7 @@ from torch import nn from .file_utils import ModelOutput -from .generation_beam_search import BeamScorer, BeamSearchScorer +from .generation_beam_search import BeamScorer, BeamSearchScorer, BeamConstraintsList from .generation_logits_process import ( EncoderNoRepeatNGramLogitsProcessor, ForcedBOSTokenLogitsProcessor, @@ -2676,7 +2676,352 @@ def group_beam_search( ) else: return sequence_outputs["sequences"] + + def constrained_beam_search( + self, + input_ids: torch.LongTensor, + beam_scorer: BeamScorer, + constraints: BeamConstraint, + logits_processor: Optional[LogitsProcessorList] = None, + stopping_criteria: Optional[StoppingCriteriaList] = None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_scores: Optional[bool] = None, + return_dict_in_generate: Optional[bool] = None, + synced_gpus: Optional[bool] = None, + **model_kwargs, + ): + r""" + Generates constrained sequences for models with a language modeling head using beam search decoding + while requiring to satisfy a list of constraints. + + Parameters: + + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + beam_scorer (`BeamScorer`): + An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and + sorted during generation. For more information, the documentation of [`BeamScorer`] should be read. + logits_processor (`LogitsProcessorList`, *optional*): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`, *optional*): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + max_length (`int`, *optional*, defaults to 20): + **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated + tokens. The maximum length of the sequence to be generated. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`int`, *optional*): + The id of the *end-of-sequence* token. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more details. + output_scores (`bool`, *optional*, defaults to `False`): + Whether or not to return the prediction scores. See `scores` under returned tensors for more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. + synced_gpus (`bool`, *optional*, defaults to `False`): + Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + + model_kwargs: + Additional model specific kwargs that will be forwarded to the `forward` function of the model. If + model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation_utils.BeamSearchDecoderOnlyOutput`], [`~generation_utils.BeamSearchEncoderDecoderOutput`] or + `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation_utils.BeamSearchDecoderOnlyOutput`] if [`~generation_utils.BeamSearchDecoderOnlyOutput`] if + `model.config.is_encoder_decoder=False` and `return_dict_in_generate=True` or a + [`~generation_utils.BeamSearchEncoderDecoderOutput`] if `model.config.is_encoder_decoder=True`. + + Examples: + + ```python + >>> from transformers import ( + ... AutoTokenizer, + ... AutoModelForSeq2SeqLM, + ... LogitsProcessorList, + ... MinLengthLogitsProcessor, + ... HammingDiversityLogitsProcessor, + ... BeamSearchScorer, + ... ) + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("t5-base") + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") + + >>> encoder_input_str = "translate English to German: How old are you?" + >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids + + + >>> # lets run diverse beam search using 6 beams + >>> num_beams = 6 + >>> # define decoder start token ids + >>> input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long) + >>> input_ids = input_ids * model.config.decoder_start_token_id + + >>> # add encoder_outputs to model keyword arguments + >>> model_kwargs = { + ... "encoder_outputs": model.get_encoder()( + ... encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True + ... ) + ... } + + >>> # instantiate beam scorer + >>> beam_scorer = BeamSearchScorer( + ... batch_size=1, + ... max_length=model.config.max_length, + ... num_beams=num_beams, + ... device=model.device, + ... num_beam_groups=3, + ... ) + + >>> # instantiate logits processors + >>> logits_processor = LogitsProcessorList( + ... [ + ... HammingDiversityLogitsProcessor(5.5, num_beams=6, num_beam_groups=3), + ... MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id), + ... ] + ... ) + + >>> outputs = model.group_beam_search( + ... input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs + ... ) + + >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) + ```""" + # init values + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() + if max_length is not None: + warnings.warn( + "`max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.", + UserWarning, + ) + stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length) + pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id + output_scores = output_scores if output_scores is not None else self.config.output_scores + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict_in_generate = ( + return_dict_in_generate if return_dict_in_generate is not None else self.config.return_dict_in_generate + ) + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + batch_size = len(beam_scorer._beam_hyps) + num_beams = beam_scorer.num_beams + num_beam_groups = beam_scorer.num_beam_groups + num_sub_beams = num_beams // num_beam_groups + device = input_ids.device + + batch_beam_size, cur_len = input_ids.shape + + if num_beams * batch_size != batch_beam_size: + raise ValueError( + f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." + ) + + beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device) + # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in + # the same group don't produce same tokens everytime. + beam_scores[:, ::num_sub_beams] = 0 + beam_scores = beam_scores.view((batch_size * num_beams,)) + + this_peer_finished = False # used by synced_gpus only + while True: + + if synced_gpus: + # Under synced_gpus the `forward` call must continue until all gpus complete their sequence. + # The following logic allows an early break if all peers finished generating their sequence + this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device) + # send 0.0 if we finished, 1.0 otherwise + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) + # did all peers finish? the reduced sum will be 0.0 then + if this_peer_finished_flag.item() == 0.0: + break + + # predicted tokens in cur_len step + current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device) + + # indices which will form the beams in the next time step + reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device) + + # do one decoder step on all beams of all sentences in batch + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + outputs = self( + **model_inputs, + return_dict=True, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + if synced_gpus and this_peer_finished: + cur_len = cur_len + 1 + continue # don't waste resources running the code we don't need + + if output_scores: + processed_score = torch.zeros_like(outputs.logits[:, -1, :]) + + for beam_group_idx in range(num_beam_groups): + group_start_idx = beam_group_idx * num_sub_beams + group_end_idx = min(group_start_idx + num_sub_beams, num_beams) + group_size = group_end_idx - group_start_idx + + # indices of beams of current group among all sentences in batch + batch_group_indices = [] + + for batch_idx in range(batch_size): + batch_group_indices.extend( + [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)] + ) + group_input_ids = input_ids[batch_group_indices] + + # select outputs of beams of current group only + next_token_logits = outputs.logits[batch_group_indices, -1, :] + + # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id` + # cannot be generated both before and after the `nn.functional.log_softmax` operation. + next_token_logits = self.adjust_logits_during_generation(next_token_logits, cur_len=cur_len) + next_token_scores = nn.functional.log_softmax( + next_token_logits, dim=-1 + ) # (batch_size * group_size, vocab_size) + vocab_size = next_token_scores.shape[-1] + + next_token_scores = logits_processor( + group_input_ids, next_token_scores, current_tokens=current_tokens, beam_group_idx=beam_group_idx + ) + next_token_scores = next_token_scores + beam_scores[batch_group_indices].unsqueeze(-1).expand_as( + next_token_scores + ) + + if output_scores: + processed_score[batch_group_indices] = next_token_scores + + # reshape for beam search + next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size) + + next_token_scores, next_tokens = torch.topk( + next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True + ) + + next_indices = next_tokens // vocab_size + next_tokens = next_tokens % vocab_size + + # stateless + beam_outputs = beam_scorer.process( + group_input_ids, + next_token_scores, + next_tokens, + next_indices, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + ) + beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"] + beam_next_tokens = beam_outputs["next_beam_tokens"] + beam_idx = beam_outputs["next_beam_indices"] + + input_ids[batch_group_indices] = group_input_ids[beam_idx] + group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + current_tokens[batch_group_indices] = group_input_ids[:, -1] + + # (beam_idx // group_size) -> batch_idx + # (beam_idx % group_size) -> offset of idx inside the group + reordering_indices[batch_group_indices] = ( + num_beams * (beam_idx // group_size) + group_start_idx + (beam_idx % group_size) + ) + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (processed_score,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1) + + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + if model_kwargs["past"] is not None: + model_kwargs["past"] = self._reorder_cache(model_kwargs["past"], reordering_indices) + + # increase cur_len + cur_len = cur_len + 1 + + if beam_scorer.is_done or stopping_criteria(input_ids, scores): + if not synced_gpus: + break + else: + this_peer_finished = True + + sequence_outputs = beam_scorer.finalize( + input_ids, + beam_scores, + next_tokens, + next_indices, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + max_length=stopping_criteria.max_length, + ) + + if return_dict_in_generate: + if not output_scores: + sequence_outputs["sequence_scores"] = None + if self.config.is_encoder_decoder: + return BeamSearchEncoderDecoderOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + ) + else: + return BeamSearchDecoderOnlyOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + ) + else: + return sequence_outputs["sequences"] def top_k_top_p_filtering( logits: torch.FloatTensor, diff --git a/tests/test_generation_beam_constraints.py b/tests/test_generation_beam_constraints.py new file mode 100644 index 000000000000..fdbe35eafaa4 --- /dev/null +++ b/tests/test_generation_beam_constraints.py @@ -0,0 +1,250 @@ +# coding=utf-8 +# Copyright 2020 The HuggingFace Team Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a clone of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import unittest + +from transformers import is_torch_available +from transformers.testing_utils import require_torch, torch_device + +from .test_modeling_common import floats_tensor, ids_tensor + + +if is_torch_available(): + import torch + + from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer + + +class BeamSearchTester: + def __init__( + self, + parent, + batch_size=3, + sequence_length=10, + vocab_size=99, + pad_token_id=0, + max_length=20, + num_beams=4, + length_penalty=2.0, + do_early_stopping=True, + num_beam_hyps_to_keep=2, + ): + self.parent = parent + self.batch_size = batch_size + self.sequence_length = sequence_length + self.vocab_size = vocab_size + self.pad_token_id = pad_token_id + self.max_length = max_length + self.num_beams = num_beams + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + + # cannot be randomely generated + self.eos_token_id = vocab_size + 1 + + def prepare_beam_scorer(self, **kwargs): + return BeamSearchScorer( + batch_size=kwargs.get("batch_size", self.batch_size), + num_beams=kwargs.get("num_beams", self.num_beams), + device=torch_device, + length_penalty=kwargs.get("length_penalty", self.length_penalty), + do_early_stopping=kwargs.get("do_early_stopping", self.do_early_stopping), + num_beam_hyps_to_keep=kwargs.get("num_beam_hyps_to_keep", self.num_beam_hyps_to_keep), + ) + + def prepare_inputs(self): + input_ids = ids_tensor((self.batch_size * self.num_beams, self.sequence_length), self.vocab_size) + next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) + next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) + next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) + return (input_ids, next_tokens, next_indices, next_scores) + + def check_beam_hypotheses(self, input_ids, *args): + # check that correct number of beam hypotheses is set in beam scorer + beam_scorer = self.prepare_beam_scorer(do_early_stopping=True) + beam_hyp = beam_scorer._beam_hyps[0] + + self.parent.assertEqual(len(beam_scorer._beam_hyps), self.batch_size) + + # check correct type + self.parent.assertTrue(isinstance(beam_hyp, BeamHypotheses)) + + # check that num_beams is correctly set + self.parent.assertEqual(beam_hyp.num_beams, self.num_beams) + + # check for early stopping deactivated + for beam_idx in range(self.num_beams): + beam_hyp.add(input_ids[beam_idx], -10.0) + + # if early stopping True -> score does not matter + self.parent.assertTrue(beam_hyp.is_done(-10.0, 5)) + + # re-init + beam_scorer = self.prepare_beam_scorer(do_early_stopping=False) + beam_hyp = beam_scorer._beam_hyps[0] + + # add `num_beams + 1` beams to change `worst_score` + for beam_idx in range(self.num_beams + 1): + beam_hyp.add(input_ids[beam_idx], -10.0 + float(beam_idx)) + + # -10.0 is removed => -9.0 is worst score + self.parent.assertAlmostEqual(beam_hyp.worst_score, -9.0 / (self.sequence_length ** beam_hyp.length_penalty)) + + # -5.0 is better than worst score => should not be finished + self.parent.assertFalse(beam_hyp.is_done(-5.0, self.sequence_length)) + + # -20.0 is worse than worst score => should be finished + self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) + + def check_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores): + # check too many eos tokens + beam_scorer = self.prepare_beam_scorer() + + tokens = next_tokens.clone() + tokens[0, :] = self.eos_token_id + + with self.parent.assertRaises(ValueError): + beam_scorer.process(input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id) + + # check all batches are done + beam_scorer = self.prepare_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, : self.num_beams] = self.eos_token_id + beam_scorer.process(input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id) + # beam scorer should be done + self.parent.assertTrue(beam_scorer.is_done) + + # check + beam_scorer = self.prepare_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, 1] = self.eos_token_id + beam_outputs = beam_scorer.process( + input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + + def cut_expected_tensor(tensor): + return torch.cat([tensor[:, :1], tensor[:, 2 : self.num_beams + 1]], dim=1).flatten() + + # check all outptus + # cut out id of eos token and take best `num_beams` outputs + expected_output_tokens = cut_expected_tensor(tokens) + expected_output_scores = cut_expected_tensor(next_scores) + + # add num_beams * batch_idx + expected_output_indices = ( + cut_expected_tensor(next_indices) + + (torch.arange(self.num_beams * self.batch_size, device=torch_device) // self.num_beams) * self.num_beams + ) + + self.parent.assertListEqual(expected_output_tokens.tolist(), output_tokens.tolist()) + self.parent.assertListEqual(expected_output_indices.tolist(), output_indices.tolist()) + self.parent.assertTrue(torch.allclose(expected_output_scores, output_scores, atol=1e-3)) + + # make sure ids of eos token are correctly saved in beam_hyps of beam scorer + for batch_idx in range(self.batch_size): + correct_idx = batch_idx * self.num_beams + next_indices[batch_idx, 1] + self.parent.assertListEqual( + input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() + ) + + def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_scores): + # max_length should be only one more than current input_ids to check that eos is correctly appended + max_length = self.sequence_length + 1 + beam_scorer = self.prepare_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) + + # update beams and append to input_ids + tokens = next_tokens.clone() + # first batch, first output has to finish with eos token id since scores are correctly sorted + tokens[0, 0] = self.eos_token_id + # make sure corresponding score is as good as possible to surely be picked first + next_scores[0, 0] = 0.0 + beam_outputs = beam_scorer.process( + input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + + input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) + + # finalize + sequence_output = beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length` + self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size]) + + # check sequence_scores + self.parent.assertFalse((sequence_scores > 0).any().item()) + + # first batch has to finish with eos_token + self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id) + + # other batches cannot finish with eos token + self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id) + self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id) + + # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned + beam_scorer.num_beam_hyps_to_keep = self.num_beams + sequence_output = beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) + + +@require_torch +class BeamSearchTest(unittest.TestCase): + def setUp(self): + self.beam_search_tester = BeamSearchTester(self) + + def test_beam_hypotheses(self): + inputs = self.beam_search_tester.prepare_inputs() + self.beam_search_tester.check_beam_hypotheses(*inputs) + + def test_beam_scorer_update(self): + inputs = self.beam_search_tester.prepare_inputs() + self.beam_search_tester.check_beam_scorer_update(*inputs) + + def test_beam_scorer_finalize(self): + inputs = self.beam_search_tester.prepare_inputs() + self.beam_search_tester.check_beam_scores_finalize(*inputs) From a2ba6c4ff1753f3af896224318c37465ee43738b Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sat, 15 Jan 2022 12:17:34 +0000 Subject: [PATCH 02/34] in progress, think i can directly force tokens now but not yet with the round robin --- .../generation_beam_constraints.py | 64 +++- src/transformers/generation_beam_search.py | 318 +++++++++++++++++- src/transformers/generation_utils.py | 266 ++++++++------- test.py | 27 ++ test.sh | 1 + tests/test.py | 33 ++ tests/test_generation_beam_constraints.py | 250 -------------- tests/test_generation_beam_search.py | 248 +++++++++++++- tests/test_generation_utils.py | 158 ++++++++- 9 files changed, 976 insertions(+), 389 deletions(-) create mode 100644 test.py create mode 100644 test.sh create mode 100644 tests/test.py delete mode 100644 tests/test_generation_beam_constraints.py diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 03d078a41ab9..d1863887c840 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -1,5 +1,6 @@ from abc import ABC +from itertools import chain from collections import Counter from typing import List, Optional, Set, Tuple @@ -28,6 +29,19 @@ class Constraint(ABC): will always terminate (halt). """ + def __init__(self): + # test for the above condition + counter = 0 + completed = False + while not completed: + advance = self.advance() + _, completed = self.update(advance) + counter += 1 + + if counter > 10000: + raise Exception("update() does not fulfill the constraint.") + + def advance(self): ''' When called, returns the token that would take this constraint @@ -76,10 +90,14 @@ class TokenConstraint(Constraint): The token that must be generated by the output. """ def __init__(self, token_id: int): + super(Constraint, self).__init__() if not isinstance(token_id, int) or token_id < 0: raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") self.token_id = token_id + def copy(self): + return TokenConstraint(self.token_id) + def advance(self): return self.token_id @@ -91,9 +109,9 @@ def update(self, token_id: int): raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") if self.does_advance(token_id): - return True, True # stepped, completed + return True, True, True # stepped, completed, reset else: - return False, False + return False, False, False class PhrasalConstraint(Constraint): @@ -105,15 +123,19 @@ class PhrasalConstraint(Constraint): The sequence of tokens that must be generated by the output. """ def __init__(self, token_ids: torch.Tensor): + super(Constraint, self).__init__() if not isinstance(token_ids, torch.Tensor): raise ValueError(f"`token_ids` has to be a tensor, but is {type(token_ids)}") self.token_ids = token_ids - # the index of the currently fulfilled step + self.seqlen = self.token_ids.size(0) - self.fulfilled_idx = -1 + self.fulfilled_idx = -1 # the index of the currently fulfilled step self.completed = False + def copy(self): + return PhrasalConstraint(self.token_ids) + def advance(self): return self.token_ids[self.fulfilled_idx + 1] @@ -136,8 +158,7 @@ def update(self, token_id: int): # failed to make progress. reset = True self.reset() - - return stepped, completed + return stepped, completed, reset def reset(self): self.completed = False @@ -149,6 +170,7 @@ def __init__(self, constraints: List[Constraint]): self.complete_constraints = [] self.inprogress_constraint = None self.pending_constraints = constraints + self.completed = False def advance(self): '''The list of tokens to generate such that we can make progress. @@ -158,18 +180,22 @@ def advance(self): we'll return. ''' if self.inprogress_constraint is None: - token_list = [ - constraint.advance() for constraint in self.pending_constraints - ] - + token_list = [] + for constraint in self.pending_constraints: + advance = constraint.advance() + token_list.append(advance) else: token_list = [self.inprogress_constraint.advance()] - - return token_list + + if len(token_list) == 0: + return None + else: + return torch.stack(token_list)[:1] def update(self, token_id: int): - if self.inprogress_constraint is None: + complete, stepped = False, False + if self.inprogress_constraint is not None: stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: self.pending_constraints.append(self.inprogress_constraint) @@ -177,6 +203,9 @@ def update(self, token_id: int): if complete: self.complete_constraints.append(self.inprogress_constraint) self.inprogress_constraint = None + + if len(self.pending_constraints) == 0: + self.completed = True else: for cidx, pending_constraint in enumerate(self.pending_constraints): if pending_constraint.does_advance(token_id): @@ -188,7 +217,12 @@ def update(self, token_id: int): self.inprogress_constraint = pending_constraint if complete or stepped: - self.pending_constraints = pending_constraints[:idx] + pending_constraints[idx+1:] + self.pending_constraints = self.pending_constraints[:cidx] + self.pending_constraints[cidx+1:] + if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: + self.completed = True + break - + + + \ No newline at end of file diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 582560c3cde3..7a67f6139bf7 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -16,10 +16,11 @@ import warnings from abc import ABC, abstractmethod from collections import UserDict -from typing import Optional, Tuple +from typing import List, Optional, Tuple import torch +from .generation_beam_constraints import Constraint, ConstraintListState from .file_utils import add_start_docstrings @@ -350,6 +351,321 @@ def finalize( ) +class ConstrainedBeamSearchScorer(BeamScorer): + r""" + [`BeamScorer`] implementing constrained beam search decoding. + + + Args: + batch_size (`int`): + Batch Size of `input_ids` for which standard beam search decoding is run in parallel. + max_length (`int`): + The maximum length of the sequence to be generated. + num_beams (`int`): + Number of beams for beam search. + constraints (`List[Constraint]`): + A list of positive constraints represented as `Constraint` objects that must be fulfilled in + the generation output. For more information, the documentation of [`Constraint`] should be read. + device (`torch.device`): + Defines the device type (*e.g.*, `"cpu"` or `"cuda"`) on which this instance of `BeamSearchScorer` will be + allocated. + length_penalty (`float`, *optional*, defaults to 1.0): + Exponential penalty to the length. 1.0 means no penalty. Set to values < 1.0 in order to encourage the + model to generate shorter sequences, to a value > 1.0 in order to encourage the model to produce longer + sequences. + do_early_stopping (`bool`, *optional*, defaults to `False`): + Whether to stop the beam search when at least `num_beams` sentences are finished per batch or not. + num_beam_hyps_to_keep (`int`, *optional*, defaults to 1): + The number of beam hypotheses that shall be returned upon calling + [`~transformer.BeamSearchScorer.finalize`]. + num_beam_groups (`int`): + Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams. + See [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details. + """ + + def __init__( + self, + batch_size: int, + num_beams: int, + constraints: List[Constraint], + device: torch.device, + length_penalty: Optional[float] = 1.0, + do_early_stopping: Optional[bool] = False, + num_beam_hyps_to_keep: Optional[int] = 1, + num_beam_groups: Optional[int] = 1, + **kwargs, + ): + self.num_beams = num_beams + self.device = device + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + self.num_beam_groups = num_beam_groups + self.group_size = self.num_beams // self.num_beam_groups + + self._is_init = False + self._beam_hyps = [ + BeamHypotheses( + num_beams=self.num_beams, + length_penalty=self.length_penalty, + early_stopping=self.do_early_stopping, + ) + for _ in range(batch_size) + ] + self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device) + + if not isinstance(num_beams, int) or num_beams <= 1: + raise ValueError( + f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1, one should make use of `greedy_search` instead." + ) + + if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0): + raise ValueError( + f"`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` " + f"has to be divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}." + ) + + if "max_length" in kwargs: + warnings.warn( + "Passing `max_length` to BeamSearchScorer is deprecated and has no effect. " + "`max_length` should be passed directly to `beam_search(...)`, `beam_sample(...)`" + ", or `group_beam_search(...)`." + ) + + @property + def is_done(self) -> bool: + return self._done.all() + + def process( + self, + input_ids: torch.LongTensor, + next_scores: torch.FloatTensor, + next_tokens: torch.LongTensor, + next_indices: torch.LongTensor, + constraint_states: List[ConstraintListState], + scores_for_all_vocab: torch.FloatTensor, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + ) -> Tuple[torch.Tensor]: + cur_len = input_ids.shape[-1] + batch_size = len(self._beam_hyps) + if not (batch_size == (input_ids.shape[0] // self.group_size)): + if self.num_beam_groups > 1: + raise ValueError( + f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam " + f"size of {self.group_size} is expected by the beam scorer." + ) + else: + raise ValueError( + f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of " + f"{self.group_size} is expected by the beam scorer." + ) + + device = input_ids.device + + next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device) + next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device) + next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device) + + for batch_idx, beam_hyp in enumerate(self._beam_hyps): + if self._done[batch_idx]: + if self.num_beams < len(beam_hyp): + raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated") + if eos_token_id is None or pad_token_id is None: + raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined") + # pad the batch + next_beam_scores[batch_idx, :] = 0 + next_beam_tokens[batch_idx, :] = pad_token_id + next_beam_indices[batch_idx, :] = 0 + continue + + + # next tokens for this sentence. + beam_idx = 0 + for beam_token_rank, (next_token, next_score, next_index) in enumerate( + zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx]) + ): + batch_beam_idx = batch_idx * self.group_size + next_index + # add to generated hypotheses if end of sentence + if (eos_token_id is not None) and (next_token.item() == eos_token_id): + # if beam_token does not belong to top num_beams tokens, it should not be added + is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size + if is_beam_token_worse_than_top_num_beams: + continue + beam_hyp.add( + input_ids[batch_beam_idx].clone(), + next_score.item(), + ) + else: + # add next predicted token since it is not eos_token + next_beam_scores[batch_idx, beam_idx] = next_score + next_beam_tokens[batch_idx, beam_idx] = next_token + next_beam_indices[batch_idx, beam_idx] = batch_beam_idx + beam_idx += 1 + + # once the beam for next step is full, don't add more tokens to it. + if beam_idx == self.group_size: + break + + print("!!input_ids", input_ids.size()) + + print("!!next_beam_tokens", next_beam_tokens[batch_idx]) + print("!!next_beam_scores", next_beam_scores[batch_idx]) + new_constraint_states, new_scores, new_tokens = self.step_sentence_constraint( + constraint_states[batch_idx], + next_beam_scores[batch_idx].clone(), + next_beam_tokens[batch_idx].clone(), + next_beam_indices[batch_idx].clone(), + scores_for_all_vocab[batch_idx].clone() + ) + print("!!new_scores", new_scores) + print("!!new_tokens", new_tokens) + print() + + constraint_states[batch_idx] = new_constraint_states + next_beam_scores[batch_idx] = new_scores + next_beam_tokens[batch_idx] = new_tokens + # next_beam_indices[batch_idx] = new_indices + + + if beam_idx < self.group_size: + raise ValueError( + f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id: {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected." + ) + + # Check if we are done so that we can save a pad step if all(done) + self._done[batch_idx] = self._done[batch_idx] or beam_hyp.is_done( + next_scores[batch_idx].max().item(), cur_len + ) + + return UserDict( + { + "next_beam_scores": next_beam_scores.view(-1), + "next_beam_tokens": next_beam_tokens.view(-1), + "next_beam_indices": next_beam_indices.view(-1), + } + ) + + def step_sentence_constraint( + self, + sent_constraint_states, + sent_beam_scores, + sent_beam_tokens, + sent_beam_indices, + sent_vocab_scores + ): + orig_len = sent_beam_indices.size(0) + device = sent_beam_indices.get_device() + dtype = sent_beam_indices.dtype + + beam_idx = 0 + for _, _ in enumerate(sent_beam_scores): + if not sent_constraint_states[beam_idx].completed: + advance_tokens = sent_constraint_states[beam_idx].advance() + print(">>>>advance_tokens", advance_tokens) + if advance_tokens.numel() != 0: + additional_num = advance_tokens.size(0) + + # assume constraint ends up being added for now + sent_constraint_states[beam_idx].update(advance_tokens[0]) + # print("!!", stepped, completed) + + next_beam_tokens = ( + advance_tokens + .repeat(additional_num) + .long() + .to(device) + ) + sent_beam_tokens = torch.cat((sent_beam_tokens, next_beam_tokens)) + + sent_beam_scores = torch.cat(( + sent_beam_scores, + sent_vocab_scores.take(advance_tokens) + )) + + + + sent_beam_tokens = sent_beam_tokens[additional_num:] + sent_beam_scores = sent_beam_scores[additional_num:] + + beam_idx += 1 + + break + else: + print("NO MORE ADVANCE") + + sent_beam_scores = sent_beam_scores[-orig_len:] + sent_beam_tokens = sent_beam_tokens[-orig_len:] + assert sent_beam_scores.size(0) == orig_len + + return sent_constraint_states, sent_beam_scores, sent_beam_tokens + + + def finalize( + self, + input_ids: torch.LongTensor, + final_beam_scores: torch.FloatTensor, + final_beam_tokens: torch.LongTensor, + final_beam_indices: torch.LongTensor, + max_length: int, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + ) -> Tuple[torch.LongTensor]: + batch_size = len(self._beam_hyps) + + # finalize all open beam hypotheses and add to generated hypotheses + for batch_idx, beam_hyp in enumerate(self._beam_hyps): + if self._done[batch_idx]: + continue + + # all open beam hypotheses are added to the beam hypothesis + # beam hypothesis class automatically keeps the best beams + for beam_id in range(self.num_beams): + batch_beam_idx = batch_idx * self.num_beams + beam_id + final_score = final_beam_scores[batch_beam_idx].item() + final_tokens = input_ids[batch_beam_idx] + beam_hyp.add(final_tokens, final_score) + + # select the best hypotheses + sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) + best = [] + best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32) + + # retrieve best hypotheses + for i, beam_hyp in enumerate(self._beam_hyps): + sorted_hyps = sorted(beam_hyp.beams, key=lambda x: x[0]) + for j in range(self.num_beam_hyps_to_keep): + best_hyp_tuple = sorted_hyps.pop() + best_score = best_hyp_tuple[0] + best_hyp = best_hyp_tuple[1] + sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp) + + # append to lists + best.append(best_hyp) + best_scores[i * self.num_beam_hyps_to_keep + j] = best_score + + # prepare for adding eos + sent_max_len = min(sent_lengths.max().item() + 1, max_length) + decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) + # shorter batches are padded if needed + if sent_lengths.min().item() != sent_lengths.max().item(): + assert pad_token_id is not None, "`pad_token_id` has to be defined" + decoded.fill_(pad_token_id) + + # fill with hypotheses and eos_token_id if the latter fits in + for i, hypo in enumerate(best): + decoded[i, : sent_lengths[i]] = hypo + if sent_lengths[i] < max_length: + decoded[i, sent_lengths[i]] = eos_token_id + return UserDict( + { + "sequences": decoded, + "sequence_scores": best_scores, + } + ) + + + class BeamHypotheses: def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool): """ diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index daa8f5831cf4..803bd5796f73 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -24,7 +24,7 @@ from torch import nn from .file_utils import ModelOutput -from .generation_beam_search import BeamScorer, BeamSearchScorer, BeamConstraintsList +from .generation_beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from .generation_logits_process import ( EncoderNoRepeatNGramLogitsProcessor, ForcedBOSTokenLogitsProcessor, @@ -48,6 +48,10 @@ StoppingCriteriaList, validate_stopping_criteria, ) +from .generation_beam_constraints import ( + Constraint, + ConstraintListState +) from .utils import logging @@ -773,6 +777,7 @@ def generate( prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, logits_processor: Optional[LogitsProcessorList] = LogitsProcessorList(), stopping_criteria: Optional[StoppingCriteriaList] = StoppingCriteriaList(), + constraints: Optional[List[Constraint]] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_scores: Optional[bool] = None, @@ -1069,11 +1074,12 @@ def generate( ) # 6. determine generation mode - is_greedy_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is False - is_sample_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is True - is_beam_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is False - is_beam_sample_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is True - is_group_beam_gen_mode = (num_beams > 1) and (num_beam_groups > 1) + is_constraint_gen_mode = constraints is not None + is_greedy_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is False and constraints is None + is_sample_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is True and constraints is None + is_beam_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is False and constraints is None + is_beam_sample_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is True and constraints is None + is_group_beam_gen_mode = (num_beams > 1) and (num_beam_groups > 1) and constraints is None if num_beam_groups > num_beams: raise ValueError("`num_beam_groups` has to be smaller or equal to `num_beams`") @@ -1268,6 +1274,43 @@ def generate( **model_kwargs, ) + elif is_constraint_gen_mode: + if num_return_sequences > num_beams: + raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.") + + if stopping_criteria.max_length is None: + raise ValueError("`max_length` needs to be a stopping_criteria for now.") + + # 10. prepare beam search scorer + constrained_beam_scorer = ConstrainedBeamSearchScorer( + constraints=constraints, + batch_size=batch_size, + num_beams=num_beams, + device=self.device, + length_penalty=length_penalty, + do_early_stopping=early_stopping, + num_beam_hyps_to_keep=num_return_sequences, + ) + # 11. interleave input_ids with `num_beams` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids, expand_size=num_beams, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs + ) + # 12. run beam search + return self.constrained_beam_search( + input_ids, + constraints=constraints, + constrained_beam_scorer=constrained_beam_scorer, + logits_processor=logits_processor, + stopping_criteria=stopping_criteria, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + output_scores=output_scores, + return_dict_in_generate=return_dict_in_generate, + synced_gpus=synced_gpus, + **model_kwargs, + ) + + def greedy_search( self, input_ids: torch.LongTensor, @@ -2676,12 +2719,12 @@ def group_beam_search( ) else: return sequence_outputs["sequences"] - + def constrained_beam_search( self, input_ids: torch.LongTensor, - beam_scorer: BeamScorer, - constraints: BeamConstraint, + constraints: List[Constraint], + constrained_beam_scorer: ConstrainedBeamSearchScorer, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, max_length: Optional[int] = None, @@ -2693,18 +2736,21 @@ def constrained_beam_search( return_dict_in_generate: Optional[bool] = None, synced_gpus: Optional[bool] = None, **model_kwargs, - ): + ) -> Union[BeamSearchOutput, torch.LongTensor]: r""" - Generates constrained sequences for models with a language modeling head using beam search decoding - while requiring to satisfy a list of constraints. + Generates sequences for models with a language modeling head using beam search decoding. Parameters: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. - beam_scorer (`BeamScorer`): + constraints (`List[Constraint]`): + A list of positive constraints represented as `Constraint` objects that must be fulfilled in + the generation output. For more information, the documentation of [`Constraint`] should be read. + constrained_beam_scorer (`ConstrainedBeamScorer`): An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and - sorted during generation. For more information, the documentation of [`BeamScorer`] should be read. + sorted during generation, while satisfying a list of positive constraints. For more information, the + documentation of [`ConstrainedBeamScorer`] should be read. logits_processor (`LogitsProcessorList`, *optional*): An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] used to modify the prediction scores of the language modeling head applied at each generation step. @@ -2730,17 +2776,17 @@ def constrained_beam_search( Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. synced_gpus (`bool`, *optional*, defaults to `False`): Whether to continue running the while loop until max_length (needed for ZeRO stage 3) - model_kwargs: - Additional model specific kwargs that will be forwarded to the `forward` function of the model. If - model is an encoder-decoder model the kwargs should include `encoder_outputs`. + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. Return: - [`~generation_utils.BeamSearchDecoderOnlyOutput`], [`~generation_utils.BeamSearchEncoderDecoderOutput`] or + [`generation_utilsBeamSearchDecoderOnlyOutput`], [`~generation_utils.BeamSearchEncoderDecoderOutput`] or `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a - [`~generation_utils.BeamSearchDecoderOnlyOutput`] if [`~generation_utils.BeamSearchDecoderOnlyOutput`] if - `model.config.is_encoder_decoder=False` and `return_dict_in_generate=True` or a - [`~generation_utils.BeamSearchEncoderDecoderOutput`] if `model.config.is_encoder_decoder=True`. + [`~generation_utils.BeamSearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation_utils.BeamSearchEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + Examples: @@ -2750,8 +2796,8 @@ def constrained_beam_search( ... AutoModelForSeq2SeqLM, ... LogitsProcessorList, ... MinLengthLogitsProcessor, - ... HammingDiversityLogitsProcessor, - ... BeamSearchScorer, + ... ConstrainedBeamSearchScorer, + ... PhrasalConstraint ... ) >>> import torch @@ -2762,8 +2808,8 @@ def constrained_beam_search( >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids - >>> # lets run diverse beam search using 6 beams - >>> num_beams = 6 + >>> # lets run beam search using 3 beams + >>> num_beams = 3 >>> # define decoder start token ids >>> input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long) >>> input_ids = input_ids * model.config.decoder_start_token_id @@ -2775,30 +2821,42 @@ def constrained_beam_search( ... ) ... } + >>> constraints = [ + ... PhrasalConstraint(tokenizer.encode("required phrase")) + ... ] + + >>> # instantiate beam scorer - >>> beam_scorer = BeamSearchScorer( + >>> beam_scorer = ConstrainedBeamSearchScorer( ... batch_size=1, - ... max_length=model.config.max_length, ... num_beams=num_beams, ... device=model.device, - ... num_beam_groups=3, + ... constraints=constraints ... ) >>> # instantiate logits processors >>> logits_processor = LogitsProcessorList( ... [ - ... HammingDiversityLogitsProcessor(5.5, num_beams=6, num_beam_groups=3), ... MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id), ... ] ... ) - >>> outputs = model.group_beam_search( - ... input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs - ... ) + >>> outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) ```""" # init values + print("constraints", constraints) + constraint_states = [ + [ + ConstraintListState([ + constraint.copy() + for constraint in constraints + ]) + for _ in range(constrained_beam_scorer.num_beams) + ] for _ in range(len(constrained_beam_scorer._beam_hyps)) + ] # need a constraint tracker for each sentence output, for each batch + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() if max_length is not None: @@ -2807,6 +2865,8 @@ def constrained_beam_search( UserWarning, ) stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length) + if len(stopping_criteria) == 0: + warnings.warn("You don't have defined any stopping_criteria, this will likely loop forever", UserWarning) pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id output_scores = output_scores if output_scores is not None else self.config.output_scores @@ -2831,11 +2891,8 @@ def constrained_beam_search( model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None ) - batch_size = len(beam_scorer._beam_hyps) - num_beams = beam_scorer.num_beams - num_beam_groups = beam_scorer.num_beam_groups - num_sub_beams = num_beams // num_beam_groups - device = input_ids.device + batch_size = len(constrained_beam_scorer._beam_hyps) + num_beams = constrained_beam_scorer.num_beams batch_beam_size, cur_len = input_ids.shape @@ -2844,12 +2901,12 @@ def constrained_beam_search( f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." ) - beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device) - # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in - # the same group don't produce same tokens everytime. - beam_scores[:, ::num_sub_beams] = 0 + beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) + beam_scores[:, 1:] = -1e9 beam_scores = beam_scores.view((batch_size * num_beams,)) + print("<><><>input_ids", input_ids) + this_peer_finished = False # used by synced_gpus only while True: @@ -2863,14 +2920,9 @@ def constrained_beam_search( if this_peer_finished_flag.item() == 0.0: break - # predicted tokens in cur_len step - current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device) - - # indices which will form the beams in the next time step - reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device) - - # do one decoder step on all beams of all sentences in batch + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + outputs = self( **model_inputs, return_dict=True, @@ -2882,81 +2934,26 @@ def constrained_beam_search( cur_len = cur_len + 1 continue # don't waste resources running the code we don't need - if output_scores: - processed_score = torch.zeros_like(outputs.logits[:, -1, :]) - - for beam_group_idx in range(num_beam_groups): - group_start_idx = beam_group_idx * num_sub_beams - group_end_idx = min(group_start_idx + num_sub_beams, num_beams) - group_size = group_end_idx - group_start_idx - - # indices of beams of current group among all sentences in batch - batch_group_indices = [] - - for batch_idx in range(batch_size): - batch_group_indices.extend( - [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)] - ) - group_input_ids = input_ids[batch_group_indices] - - # select outputs of beams of current group only - next_token_logits = outputs.logits[batch_group_indices, -1, :] - - # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id` - # cannot be generated both before and after the `nn.functional.log_softmax` operation. - next_token_logits = self.adjust_logits_during_generation(next_token_logits, cur_len=cur_len) - next_token_scores = nn.functional.log_softmax( - next_token_logits, dim=-1 - ) # (batch_size * group_size, vocab_size) - vocab_size = next_token_scores.shape[-1] - - next_token_scores = logits_processor( - group_input_ids, next_token_scores, current_tokens=current_tokens, beam_group_idx=beam_group_idx - ) - next_token_scores = next_token_scores + beam_scores[batch_group_indices].unsqueeze(-1).expand_as( - next_token_scores - ) - - if output_scores: - processed_score[batch_group_indices] = next_token_scores + next_token_logits = outputs.logits[:, -1, :] + # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id` + # cannot be generated both before and after the `nn.functional.log_softmax` operation. + next_token_logits = self.adjust_logits_during_generation(next_token_logits, cur_len=cur_len) + next_token_scores = nn.functional.log_softmax( + next_token_logits, dim=-1 + ) # (batch_size * num_beams, vocab_size) - # reshape for beam search - next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size) - next_token_scores, next_tokens = torch.topk( - next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True - ) + next_token_scores = logits_processor(input_ids, next_token_scores) - next_indices = next_tokens // vocab_size - next_tokens = next_tokens % vocab_size + scores_for_all_vocab = next_token_scores.clone() - # stateless - beam_outputs = beam_scorer.process( - group_input_ids, - next_token_scores, - next_tokens, - next_indices, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - ) - beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"] - beam_next_tokens = beam_outputs["next_beam_tokens"] - beam_idx = beam_outputs["next_beam_indices"] + next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores) - input_ids[batch_group_indices] = group_input_ids[beam_idx] - group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - current_tokens[batch_group_indices] = group_input_ids[:, -1] - - # (beam_idx // group_size) -> batch_idx - # (beam_idx % group_size) -> offset of idx inside the group - reordering_indices[batch_group_indices] = ( - num_beams * (beam_idx // group_size) + group_start_idx + (beam_idx % group_size) - ) # Store scores, attentions and hidden_states when required if return_dict_in_generate: if output_scores: - scores += (processed_score,) + scores += (next_token_scores,) if output_attentions: decoder_attentions += ( (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) @@ -2971,24 +2968,55 @@ def constrained_beam_search( else (outputs.hidden_states,) ) - input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1) + # reshape for beam search + vocab_size = next_token_scores.shape[-1] + next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) + + + next_token_scores, next_tokens = torch.topk( + next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True + ) + + next_indices = (next_tokens / vocab_size).long() + next_tokens = next_tokens % vocab_size + + print("!!!!!!scores_for_all_vocab", scores_for_all_vocab.size()) + # stateless + beam_outputs = constrained_beam_scorer.process( + input_ids, + next_token_scores, + next_tokens, + next_indices, + constraint_states, + scores_for_all_vocab, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + ) + beam_scores = beam_outputs["next_beam_scores"] + beam_next_tokens = beam_outputs["next_beam_tokens"] + beam_idx = beam_outputs["next_beam_indices"] + + print("beam_idx", beam_idx) + print("input_ids", input_ids.size()) + print("beam_next_tokens", beam_next_tokens.size()) + input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder ) if model_kwargs["past"] is not None: - model_kwargs["past"] = self._reorder_cache(model_kwargs["past"], reordering_indices) + model_kwargs["past"] = self._reorder_cache(model_kwargs["past"], beam_idx) # increase cur_len cur_len = cur_len + 1 - if beam_scorer.is_done or stopping_criteria(input_ids, scores): + if constrained_beam_scorer.is_done or stopping_criteria(input_ids, scores): if not synced_gpus: break else: this_peer_finished = True - sequence_outputs = beam_scorer.finalize( + sequence_outputs = constrained_beam_scorer.finalize( input_ids, beam_scores, next_tokens, @@ -3022,7 +3050,7 @@ def constrained_beam_search( ) else: return sequence_outputs["sequences"] - + def top_k_top_p_filtering( logits: torch.FloatTensor, top_k: int = 0, diff --git a/test.py b/test.py new file mode 100644 index 000000000000..75e1dc07ca6e --- /dev/null +++ b/test.py @@ -0,0 +1,27 @@ +import unittest + +from transformers import BartForConditionalGeneration, BartTokenizer +from transformers.generation_beam_constraints import ( + PhrasalConstraint +) +model = BartForConditionalGeneration.from_pretrained("facebook/bart-base") +tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") + +force_text = "forced" +force_tokens = tokenizer.encode(force_text, return_tensors="pt") +print("force_tokens", force_tokens) +constraints = [PhrasalConstraint(force_tokens[1:])] + +input_text = ["This feels a little"] + +model_inputs = tokenizer(input_text, return_tensors="pt") + +print("input_ids", input_ids) +k = model.generate( + **model_inputs, + constraints=constraints +) + +print(k) + +assert False diff --git a/test.sh b/test.sh new file mode 100644 index 000000000000..a7bfb7d17e4c --- /dev/null +++ b/test.sh @@ -0,0 +1 @@ +CUDA_LAUNCH_BLOCKING=1, pytest -s tests/test_modeling_bart.py::BartStandaloneDecoderModelTest::test_constrained_beam_search_generate --capture=sys diff --git a/tests/test.py b/tests/test.py new file mode 100644 index 000000000000..8ca1ac3b9424 --- /dev/null +++ b/tests/test.py @@ -0,0 +1,33 @@ +import unittest + +from transformers import BartForConditionalGeneration, BartTokenizer +from transformers.generation_beam_constraints import ( + PhrasalConstraint +) +device = "cuda" + +model = BartForConditionalGeneration.from_pretrained("facebook/bart-base").to(device) +tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") + +force_text = "forced a little" +force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device) +print("force_tokens", force_tokens[0][1:-1]) +constraints = [PhrasalConstraint(force_tokens[0][1:-1])] + +input_text = ["this feels very"] * 10 + +model_inputs = tokenizer(input_text, return_tensors="pt") + +for key, value in model_inputs.items(): + model_inputs[key] = value.to(device) + +print("model_inputs", model_inputs) +k = model.generate( + **model_inputs, + constraints=constraints +) + +for out in k: + print(tokenizer.decode(out)) + +assert False diff --git a/tests/test_generation_beam_constraints.py b/tests/test_generation_beam_constraints.py deleted file mode 100644 index fdbe35eafaa4..000000000000 --- a/tests/test_generation_beam_constraints.py +++ /dev/null @@ -1,250 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The HuggingFace Team Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a clone of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -import unittest - -from transformers import is_torch_available -from transformers.testing_utils import require_torch, torch_device - -from .test_modeling_common import floats_tensor, ids_tensor - - -if is_torch_available(): - import torch - - from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer - - -class BeamSearchTester: - def __init__( - self, - parent, - batch_size=3, - sequence_length=10, - vocab_size=99, - pad_token_id=0, - max_length=20, - num_beams=4, - length_penalty=2.0, - do_early_stopping=True, - num_beam_hyps_to_keep=2, - ): - self.parent = parent - self.batch_size = batch_size - self.sequence_length = sequence_length - self.vocab_size = vocab_size - self.pad_token_id = pad_token_id - self.max_length = max_length - self.num_beams = num_beams - self.length_penalty = length_penalty - self.do_early_stopping = do_early_stopping - self.num_beam_hyps_to_keep = num_beam_hyps_to_keep - - # cannot be randomely generated - self.eos_token_id = vocab_size + 1 - - def prepare_beam_scorer(self, **kwargs): - return BeamSearchScorer( - batch_size=kwargs.get("batch_size", self.batch_size), - num_beams=kwargs.get("num_beams", self.num_beams), - device=torch_device, - length_penalty=kwargs.get("length_penalty", self.length_penalty), - do_early_stopping=kwargs.get("do_early_stopping", self.do_early_stopping), - num_beam_hyps_to_keep=kwargs.get("num_beam_hyps_to_keep", self.num_beam_hyps_to_keep), - ) - - def prepare_inputs(self): - input_ids = ids_tensor((self.batch_size * self.num_beams, self.sequence_length), self.vocab_size) - next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) - next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) - next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) - return (input_ids, next_tokens, next_indices, next_scores) - - def check_beam_hypotheses(self, input_ids, *args): - # check that correct number of beam hypotheses is set in beam scorer - beam_scorer = self.prepare_beam_scorer(do_early_stopping=True) - beam_hyp = beam_scorer._beam_hyps[0] - - self.parent.assertEqual(len(beam_scorer._beam_hyps), self.batch_size) - - # check correct type - self.parent.assertTrue(isinstance(beam_hyp, BeamHypotheses)) - - # check that num_beams is correctly set - self.parent.assertEqual(beam_hyp.num_beams, self.num_beams) - - # check for early stopping deactivated - for beam_idx in range(self.num_beams): - beam_hyp.add(input_ids[beam_idx], -10.0) - - # if early stopping True -> score does not matter - self.parent.assertTrue(beam_hyp.is_done(-10.0, 5)) - - # re-init - beam_scorer = self.prepare_beam_scorer(do_early_stopping=False) - beam_hyp = beam_scorer._beam_hyps[0] - - # add `num_beams + 1` beams to change `worst_score` - for beam_idx in range(self.num_beams + 1): - beam_hyp.add(input_ids[beam_idx], -10.0 + float(beam_idx)) - - # -10.0 is removed => -9.0 is worst score - self.parent.assertAlmostEqual(beam_hyp.worst_score, -9.0 / (self.sequence_length ** beam_hyp.length_penalty)) - - # -5.0 is better than worst score => should not be finished - self.parent.assertFalse(beam_hyp.is_done(-5.0, self.sequence_length)) - - # -20.0 is worse than worst score => should be finished - self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) - - def check_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores): - # check too many eos tokens - beam_scorer = self.prepare_beam_scorer() - - tokens = next_tokens.clone() - tokens[0, :] = self.eos_token_id - - with self.parent.assertRaises(ValueError): - beam_scorer.process(input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id) - - # check all batches are done - beam_scorer = self.prepare_beam_scorer() - - tokens = next_tokens.clone() - tokens[:, : self.num_beams] = self.eos_token_id - beam_scorer.process(input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id) - # beam scorer should be done - self.parent.assertTrue(beam_scorer.is_done) - - # check - beam_scorer = self.prepare_beam_scorer() - - tokens = next_tokens.clone() - tokens[:, 1] = self.eos_token_id - beam_outputs = beam_scorer.process( - input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id - ) - output_scores = beam_outputs["next_beam_scores"] - output_tokens = beam_outputs["next_beam_tokens"] - output_indices = beam_outputs["next_beam_indices"] - - def cut_expected_tensor(tensor): - return torch.cat([tensor[:, :1], tensor[:, 2 : self.num_beams + 1]], dim=1).flatten() - - # check all outptus - # cut out id of eos token and take best `num_beams` outputs - expected_output_tokens = cut_expected_tensor(tokens) - expected_output_scores = cut_expected_tensor(next_scores) - - # add num_beams * batch_idx - expected_output_indices = ( - cut_expected_tensor(next_indices) - + (torch.arange(self.num_beams * self.batch_size, device=torch_device) // self.num_beams) * self.num_beams - ) - - self.parent.assertListEqual(expected_output_tokens.tolist(), output_tokens.tolist()) - self.parent.assertListEqual(expected_output_indices.tolist(), output_indices.tolist()) - self.parent.assertTrue(torch.allclose(expected_output_scores, output_scores, atol=1e-3)) - - # make sure ids of eos token are correctly saved in beam_hyps of beam scorer - for batch_idx in range(self.batch_size): - correct_idx = batch_idx * self.num_beams + next_indices[batch_idx, 1] - self.parent.assertListEqual( - input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() - ) - - def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_scores): - # max_length should be only one more than current input_ids to check that eos is correctly appended - max_length = self.sequence_length + 1 - beam_scorer = self.prepare_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) - - # update beams and append to input_ids - tokens = next_tokens.clone() - # first batch, first output has to finish with eos token id since scores are correctly sorted - tokens[0, 0] = self.eos_token_id - # make sure corresponding score is as good as possible to surely be picked first - next_scores[0, 0] = 0.0 - beam_outputs = beam_scorer.process( - input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id - ) - output_scores = beam_outputs["next_beam_scores"] - output_tokens = beam_outputs["next_beam_tokens"] - output_indices = beam_outputs["next_beam_indices"] - - input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) - - # finalize - sequence_output = beam_scorer.finalize( - input_ids, - output_scores, - output_tokens, - output_indices, - pad_token_id=self.pad_token_id, - eos_token_id=self.eos_token_id, - max_length=max_length, - ) - - sequences = sequence_output["sequences"] - sequence_scores = sequence_output["sequence_scores"] - - # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length` - self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length]) - self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size]) - - # check sequence_scores - self.parent.assertFalse((sequence_scores > 0).any().item()) - - # first batch has to finish with eos_token - self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id) - - # other batches cannot finish with eos token - self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id) - self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id) - - # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned - beam_scorer.num_beam_hyps_to_keep = self.num_beams - sequence_output = beam_scorer.finalize( - input_ids, - output_scores, - output_tokens, - output_indices, - pad_token_id=self.pad_token_id, - eos_token_id=self.eos_token_id, - max_length=max_length, - ) - sequences = sequence_output["sequences"] - sequence_scores = sequence_output["sequence_scores"] - - self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) - self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) - - -@require_torch -class BeamSearchTest(unittest.TestCase): - def setUp(self): - self.beam_search_tester = BeamSearchTester(self) - - def test_beam_hypotheses(self): - inputs = self.beam_search_tester.prepare_inputs() - self.beam_search_tester.check_beam_hypotheses(*inputs) - - def test_beam_scorer_update(self): - inputs = self.beam_search_tester.prepare_inputs() - self.beam_search_tester.check_beam_scorer_update(*inputs) - - def test_beam_scorer_finalize(self): - inputs = self.beam_search_tester.prepare_inputs() - self.beam_search_tester.check_beam_scores_finalize(*inputs) diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index fdbe35eafaa4..51df97a22984 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -17,7 +17,10 @@ import unittest from transformers import is_torch_available -from transformers.testing_utils import require_torch, torch_device +# from transformers.testing_utils import require_torch, torch_device +from transformers.testing_utils import require_torch + +torch_device = "cpu" from .test_modeling_common import floats_tensor, ids_tensor @@ -25,8 +28,12 @@ if is_torch_available(): import torch - from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer - + from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer, ConstrainedBeamSearchScorer + from transformers.generation_beam_constraints import ( + Constraint, + PhrasalConstraint, + ConstraintListState + ) class BeamSearchTester: def __init__( @@ -232,6 +239,222 @@ def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_ self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) +class ConstrainedBeamSearchTester: + def __init__( + self, + parent, + batch_size=3, + sequence_length=10, + vocab_size=99, + pad_token_id=0, + max_length=20, + num_beams=4, + length_penalty=2.0, + do_early_stopping=True, + num_beam_hyps_to_keep=2, + ): + self.parent = parent + self.batch_size = batch_size + self.sequence_length = sequence_length + self.vocab_size = vocab_size + self.pad_token_id = pad_token_id + self.max_length = max_length + self.num_beams = num_beams + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + + self.constraints = [ + PhrasalConstraint(ids_tensor((1, 2), self.vocab_size)[0]) + ] + # cannot be randomly generated + self.eos_token_id = vocab_size + 1 + + def prepare_beam_scorer(self, **kwargs): + + return ConstrainedBeamSearchScorer( + batch_size=kwargs.get("batch_size", self.batch_size), + num_beams=kwargs.get("num_beams", self.num_beams), + constraints=kwargs.get("constraints", self.constraints), + device=torch_device, + length_penalty=kwargs.get("length_penalty", self.length_penalty), + do_early_stopping=kwargs.get("do_early_stopping", self.do_early_stopping), + num_beam_hyps_to_keep=kwargs.get("num_beam_hyps_to_keep", self.num_beam_hyps_to_keep), + ) + + def prepare_inputs(self): + constraint_states = [ + ConstraintListState(self.constraints) + for _ in range(self.batch_size) + ] # n + + input_ids = ids_tensor((self.batch_size * self.num_beams, self.sequence_length), self.vocab_size) + next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) + next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) + next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) + return (input_ids, next_tokens, next_indices, next_scores, constraint_states) + + def check_beam_hypotheses(self, input_ids, *args): + # check that correct number of beam hypotheses is set in beam scorer + beam_scorer = self.prepare_beam_scorer(do_early_stopping=True) + beam_hyp = beam_scorer._beam_hyps[0] + + self.parent.assertEqual(len(beam_scorer._beam_hyps), self.batch_size) + + # check correct type + self.parent.assertTrue(isinstance(beam_hyp, BeamHypotheses)) + + # check that num_beams is correctly set + self.parent.assertEqual(beam_hyp.num_beams, self.num_beams) + + # check for early stopping deactivated + for beam_idx in range(self.num_beams): + beam_hyp.add(input_ids[beam_idx], -10.0) + + # if early stopping True -> score does not matter + self.parent.assertTrue(beam_hyp.is_done(-10.0, 5)) + + # re-init + beam_scorer = self.prepare_beam_scorer(do_early_stopping=False) + beam_hyp = beam_scorer._beam_hyps[0] + + # add `num_beams + 1` beams to change `worst_score` + for beam_idx in range(self.num_beams + 1): + beam_hyp.add(input_ids[beam_idx], -10.0 + float(beam_idx)) + + # -10.0 is removed => -9.0 is worst score + self.parent.assertAlmostEqual(beam_hyp.worst_score, -9.0 / (self.sequence_length ** beam_hyp.length_penalty)) + + # -5.0 is better than worst score => should not be finished + self.parent.assertFalse(beam_hyp.is_done(-5.0, self.sequence_length)) + + # -20.0 is worse than worst score => should be finished + self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) + + def check_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores, constraint_states): + # check too many eos tokens + constrained_beam_scorer = self.prepare_beam_scorer() + print("---1next_scores", next_scores) + + tokens = next_tokens.clone() + tokens[0, :] = self.eos_token_id + + with self.parent.assertRaises(ValueError): + constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id) + + # check all batches are done + constrained_beam_scorer = self.prepare_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, : self.num_beams] = self.eos_token_id + constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id) + # beam scorer should be done + self.parent.assertTrue(constrained_beam_scorer.is_done) + + # check + constrained_beam_scorer = self.prepare_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, 1] = self.eos_token_id + beam_outputs = constrained_beam_scorer.process( + input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + + def cut_expected_tensor(tensor): + return torch.cat([tensor[:, :1], tensor[:, 2 : self.num_beams + 1]], dim=1).flatten() + + # check all outptus + # cut out id of eos token and take best `num_beams` outputs + expected_output_tokens = cut_expected_tensor(tokens) + expected_output_scores = cut_expected_tensor(next_scores) + + # add num_beams * batch_idx + expected_output_indices = ( + cut_expected_tensor(next_indices) + + (torch.arange(self.num_beams * self.batch_size, device=torch_device) // self.num_beams) * self.num_beams + ) + + self.parent.assertListEqual(expected_output_tokens.tolist(), output_tokens.tolist()) + self.parent.assertListEqual(expected_output_indices.tolist(), output_indices.tolist()) + self.parent.assertTrue(torch.allclose(expected_output_scores, output_scores, atol=1e-3)) + + # make sure ids of eos token are correctly saved in beam_hyps of beam scorer + for batch_idx in range(self.batch_size): + correct_idx = batch_idx * self.num_beams + next_indices[batch_idx, 1] + self.parent.assertListEqual( + input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() + ) + + def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_scores): + # max_length should be only one more than current input_ids to check that eos is correctly appended + max_length = self.sequence_length + 1 + beam_scorer = self.prepare_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) + + # update beams and append to input_ids + tokens = next_tokens.clone() + # first batch, first output has to finish with eos token id since scores are correctly sorted + tokens[0, 0] = self.eos_token_id + # make sure corresponding score is as good as possible to surely be picked first + next_scores[0, 0] = 0.0 + beam_outputs = beam_scorer.process( + input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + + input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) + + # finalize + sequence_output = beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length` + self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size]) + + # check sequence_scores + self.parent.assertFalse((sequence_scores > 0).any().item()) + + # first batch has to finish with eos_token + self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id) + + # other batches cannot finish with eos token + self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id) + self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id) + + # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned + beam_scorer.num_beam_hyps_to_keep = self.num_beams + sequence_output = beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) + + + @require_torch class BeamSearchTest(unittest.TestCase): def setUp(self): @@ -248,3 +471,22 @@ def test_beam_scorer_update(self): def test_beam_scorer_finalize(self): inputs = self.beam_search_tester.prepare_inputs() self.beam_search_tester.check_beam_scores_finalize(*inputs) + + + +@require_torch +class ConstrainedBeamSearchTest(unittest.TestCase): + def setUp(self): + self.constrained_beam_search_tester = ConstrainedBeamSearchTester(self) + + # def test_beam_hypotheses(self): + # inputs = self.constrained_beam_search_tester.prepare_inputs() + # self.constrained_beam_search_tester.check_beam_hypotheses(*inputs) + + def test_beam_scorer_update(self): + inputs = self.constrained_beam_search_tester.prepare_inputs() + self.constrained_beam_search_tester.check_beam_scorer_update(*inputs) + + # def test_beam_scorer_finalize(self): + # inputs = self.constrained_beam_search_tester.prepare_inputs() + # self.constrained_beam_search_tester.check_beam_scores_finalize(*inputs) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 1ffb02bd4b48..a21f943e61d3 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -37,7 +37,7 @@ VisionEncoderDecoderModel, top_k_top_p_filtering, ) - from transformers.generation_beam_search import BeamSearchScorer + from transformers.generation_beam_search import BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation_logits_process import ( ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor, @@ -63,6 +63,9 @@ SampleDecoderOnlyOutput, SampleEncoderDecoderOutput, ) + from transformers.generation_beam_constraints import ( + PhrasalConstraint + ) class GenerationTesterMixin: @@ -168,6 +171,25 @@ def _get_beam_scorer_and_kwargs(batch_size, max_length, num_return_sequences=1): num_beam_hyps_to_keep=num_return_sequences, ) return beam_kwargs, beam_scorer + + @staticmethod + def _get_constrained_beam_scorer_and_kwargs(batch_size, max_length, constraints, num_return_sequences=1): + beam_kwargs = { + "early_stopping": False, + "length_penalty": 2.0, + "num_beams": 2, + "num_return_sequences": num_return_sequences, + } + beam_scorer = ConstrainedBeamSearchScorer( + batch_size=batch_size, + num_beams=beam_kwargs["num_beams"], + device=torch_device, + constraints=constraints, + length_penalty=beam_kwargs["length_penalty"], + do_early_stopping=beam_kwargs["early_stopping"], + num_beam_hyps_to_keep=num_return_sequences, + ) + return beam_kwargs, beam_scorer @staticmethod def _get_diverse_beam_scorer_and_kwargs(batch_size, max_length, num_return_sequences=1): @@ -398,6 +420,74 @@ def _beam_search_generate( ) return output_generate, output_beam_search + def _constrained_beam_search_generate( + self, + model, + input_ids, + attention_mask, + max_length, + constraints, + constrained_beam_scorer, + beam_kwargs, + logits_processor, + logits_process_kwargs, + output_scores=False, + output_attentions=False, + output_hidden_states=False, + return_dict_in_generate=False, + ): + output_generate = model.generate( + input_ids, + attention_mask=attention_mask, + do_sample=False, + max_length=max_length, + output_scores=output_scores, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict_in_generate=return_dict_in_generate, + remove_invalid_values=True, + constraints=constraints, + **beam_kwargs, + **logits_process_kwargs, + ) + print("output_generate", output_generate) + + # beam_search does not automatically interleave `batch_size` dim for `num_beams` + kwargs = {} + if model.config.is_encoder_decoder: + encoder_outputs, input_ids_clone, attention_mask_clone = self._get_encoder_outputs( + model, + input_ids, + attention_mask, + num_interleave=constrained_beam_scorer.num_beams, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + kwargs["encoder_outputs"] = encoder_outputs + input_ids_clone = input_ids_clone.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + else: + attention_mask_clone = attention_mask.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + input_ids_clone = input_ids.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + + with torch.no_grad(): + output_beam_search = model.constrained_beam_search( + input_ids_clone, + constraints, + constrained_beam_scorer, + max_length=max_length, + attention_mask=attention_mask_clone, + logits_processor=logits_processor, + output_scores=output_scores, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict_in_generate=return_dict_in_generate, + **kwargs, + ) + + print("output_beam_search", output_beam_search) + + return output_generate, output_beam_search + def _beam_sample_generate( self, model, @@ -740,6 +830,72 @@ def test_beam_search_generate(self): logits_processor=logits_processor, ) self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) + + def test_constrained_beam_search_generate(self): + for model_class in self.all_generative_model_classes: + config, input_ids, attention_mask, max_length = self._get_input_ids_and_config() + constraints = [ + PhrasalConstraint(ids_tensor((1, 2), config.vocab_size)[0]) + ] + + # It is important set set the eos_token_id to None to ensure that no sequences + # shorter than `max_length` can be generated which could lead to flaky circle ci + # failures if the top `num_return_sequences` beams are all shorter than the longest beam + config.eos_token_id = None + config.forced_eos_token_id = None + + model = model_class(config).to(torch_device).eval() + if model.config.is_encoder_decoder: + max_length = 4 + + logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( + input_ids.shape[-1], + config.eos_token_id, + config.forced_bos_token_id, + config.forced_eos_token_id, + max_length, + ) + beam_kwargs, constrained_beam_scorer = self._get_constrained_beam_scorer_and_kwargs( + input_ids.shape[0], + max_length, + constraints + ) + + # check `generate()` and `beam_search()` are equal + output_generate, output_beam_search = self._constrained_beam_search_generate( + model=model, + input_ids=input_ids, + constraints=constraints, + constrained_beam_scorer=constrained_beam_scorer, + attention_mask=attention_mask, + max_length=max_length, + beam_kwargs=beam_kwargs, + logits_process_kwargs=logits_process_kwargs, + logits_processor=logits_processor, + ) + assert False + self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) + + # check `generate()` and `beam_search()` are equal for `num_return_sequences` + num_return_sequences = 2 + if model.config.is_encoder_decoder: + max_length = 4 + beam_kwargs, beam_scorer = self._get_beam_scorer_and_kwargs( + input_ids.shape[0], max_length, num_return_sequences=num_return_sequences + ) + + output_generate, output_beam_search = self._beam_search_generate( + model=model, + input_ids=input_ids, + attention_mask=attention_mask, + max_length=max_length, + beam_scorer=beam_scorer, + beam_kwargs=beam_kwargs, + logits_process_kwargs=logits_process_kwargs, + logits_processor=logits_processor, + ) + self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) + def test_beam_search_generate_dict_output(self): for model_class in self.all_generative_model_classes: From c36eaa2a4bc90f9eb0d6eaf795c06b1487da19c2 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Wed, 19 Jan 2022 09:24:18 +0000 Subject: [PATCH 03/34] think now i have total control, now need to code the bank selection --- .../generation_beam_constraints.py | 73 ++++++- src/transformers/generation_beam_search.py | 184 +++++++++++++----- src/transformers/generation_utils.py | 29 +-- tests/test.py | 25 ++- 4 files changed, 226 insertions(+), 85 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index d1863887c840..95986b0feb1d 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -140,10 +140,12 @@ def advance(self): return self.token_ids[self.fulfilled_idx + 1] def does_advance(self, token_id): + if self.completed: + return False + return token_id == self.token_ids[self.fulfilled_idx + 1] def update(self, token_id: int): - stepped = False completed = False reset = False @@ -167,10 +169,17 @@ def reset(self): class ConstraintListState: def __init__(self, constraints: List[Constraint]): + self.constraints = constraints + self.n_constraints = len(constraints) + self.completed = False + + self.init_state() + + + def init_state(self): self.complete_constraints = [] self.inprogress_constraint = None - self.pending_constraints = constraints - self.completed = False + self.pending_constraints = [constraint.copy() for constraint in self.constraints] def advance(self): '''The list of tokens to generate such that we can make progress. @@ -180,6 +189,8 @@ def advance(self): we'll return. ''' if self.inprogress_constraint is None: + print("RUN ADVANCE NO PROGRESS") + print("self.pending_constraints", self.pending_constraints) token_list = [] for constraint in self.pending_constraints: advance = constraint.advance() @@ -190,39 +201,83 @@ def advance(self): if len(token_list) == 0: return None else: - return torch.stack(token_list)[:1] + return torch.stack(token_list) + + def update(self, token_ids: torch.Tensor): + ''' + token_ids: the tokens generated thus far to reset the state of the + progress through constraints. + ''' + self.init_state() + print("\n!!!START UPDATE\n\n") + for token in token_ids: + print("token", token) + complete, stepped = self.add(token) + print("complete, stepped", complete, stepped) + + return self + + def add(self, token_id: int): + if self.completed: + return - def update(self, token_id: int): complete, stepped = False, False if self.inprogress_constraint is not None: + ''' + In the middle of fulfilling a constraint. + + If the token just steps (make but an incremental progress to current job) it, do nothing. + ''' stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: - self.pending_constraints.append(self.inprogress_constraint) - self.inprogress_constraint = None + ''' + 1. If the next token breaks the fulfillment, then we must restart. + e.g. force the sequence "I love pies" and the next token after "I love" is "books". + ''' + print("RESET") + self.pending_constraints.append(self.inprogress_constraint.copy()) + self.inprogress_constraint = None + print("self.pending_constraints", self.pending_constraints) + if complete: + ''' + 2. If the next token completes the constraint, move it to completed list, set + inprogress to None. If there are no pending constraints either, then this + full list of constraints is complete. + ''' self.complete_constraints.append(self.inprogress_constraint) self.inprogress_constraint = None if len(self.pending_constraints) == 0: self.completed = True else: + ''' + Not in the middle of fulfilling a constraint. + ''' + print("NOT IN THE MIDDLE", self.pending_constraints) for cidx, pending_constraint in enumerate(self.pending_constraints): + ''' + 1. Does it advance any of the pending constraints? + ''' if pending_constraint.does_advance(token_id): stepped, complete, reset = pending_constraint.update(token_id) if complete: self.complete_constraints.append(pending_constraint) - - elif stepped: + self.inprogress_constraint = None + if not complete and stepped: self.inprogress_constraint = pending_constraint if complete or stepped: self.pending_constraints = self.pending_constraints[:cidx] + self.pending_constraints[cidx+1:] + print("!!!self.pending_constraints", self.pending_constraints) + print() if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: self.completed = True break + return complete, stepped \ No newline at end of file diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 7a67f6139bf7..23e07efe9445 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -402,7 +402,7 @@ def __init__( self.num_beam_hyps_to_keep = num_beam_hyps_to_keep self.num_beam_groups = num_beam_groups self.group_size = self.num_beams // self.num_beam_groups - + self.constraints = constraints self._is_init = False self._beam_hyps = [ BeamHypotheses( @@ -414,6 +414,8 @@ def __init__( ] self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device) + + if not isinstance(num_beams, int) or num_beams <= 1: raise ValueError( f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1, one should make use of `greedy_search` instead." @@ -436,13 +438,22 @@ def __init__( def is_done(self) -> bool: return self._done.all() + def make_constraint_states(self, n): + return [ + ConstraintListState([ + constraint.copy() + for constraint in self.constraints + ]) + for _ in range(n) + ] + + def process( self, input_ids: torch.LongTensor, next_scores: torch.FloatTensor, next_tokens: torch.LongTensor, next_indices: torch.LongTensor, - constraint_states: List[ConstraintListState], scores_for_all_vocab: torch.FloatTensor, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, @@ -488,6 +499,9 @@ def process( batch_beam_idx = batch_idx * self.group_size + next_index # add to generated hypotheses if end of sentence if (eos_token_id is not None) and (next_token.item() == eos_token_id): + # if constraint not fulfilled, it should not be added. + if not constraint_states[batch_idx].completed: + continue # if beam_token does not belong to top num_beams tokens, it should not be added is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size if is_beam_token_worse_than_top_num_beams: @@ -507,25 +521,25 @@ def process( if beam_idx == self.group_size: break - print("!!input_ids", input_ids.size()) - - print("!!next_beam_tokens", next_beam_tokens[batch_idx]) - print("!!next_beam_scores", next_beam_scores[batch_idx]) - new_constraint_states, new_scores, new_tokens = self.step_sentence_constraint( - constraint_states[batch_idx], + print("scores_for_all_vocab", scores_for_all_vocab.size()) + print("input_ids", input_ids.size()) + print("next_beam_tokens", next_beam_tokens.size()) + + new_scores, new_tokens, new_indices = self.step_sentence_constraint( + batch_idx, + input_ids, + scores_for_all_vocab, next_beam_scores[batch_idx].clone(), next_beam_tokens[batch_idx].clone(), next_beam_indices[batch_idx].clone(), - scores_for_all_vocab[batch_idx].clone() ) print("!!new_scores", new_scores) print("!!new_tokens", new_tokens) print() - constraint_states[batch_idx] = new_constraint_states next_beam_scores[batch_idx] = new_scores next_beam_tokens[batch_idx] = new_tokens - # next_beam_indices[batch_idx] = new_indices + next_beam_indices[batch_idx] = new_indices if beam_idx < self.group_size: @@ -548,57 +562,139 @@ def process( def step_sentence_constraint( self, - sent_constraint_states, + batch_idx, + input_ids, + vocab_scores, sent_beam_scores, sent_beam_tokens, sent_beam_indices, - sent_vocab_scores ): + ''' + sent_beam_tokens are the next {num_beams} number of tokens that are under consideration + for this beam (candidate next tokens) + + 1. Adding "advance_tokens" + using ConstraintStateList.advance(), we propose new tokens to be added into this + "candidate list" that will advance us in fulfilling the constraints. + + 2. Selecting best candidates such that we end up with highest probable candidates + that fulfill our constraints. + ''' orig_len = sent_beam_indices.size(0) device = sent_beam_indices.get_device() - dtype = sent_beam_indices.dtype - - beam_idx = 0 - for _, _ in enumerate(sent_beam_scores): - if not sent_constraint_states[beam_idx].completed: - advance_tokens = sent_constraint_states[beam_idx].advance() - print(">>>>advance_tokens", advance_tokens) + sent_constraint_state = self.make_constraint_states(orig_len) + + start_idx = batch_idx*orig_len + end_idx = (batch_idx+1) * orig_len + + this_batch_input_ids = input_ids[start_idx : end_idx] + this_batch_token_scores = vocab_scores[start_idx : end_idx] + + print("this_batch_input_ids", this_batch_input_ids) + print("sent_beam_tokens", sent_beam_tokens) + + full_hypotheses = torch.cat((this_batch_input_ids, sent_beam_tokens.unsqueeze(-1)), dim=-1) + print("full_hypotheses", full_hypotheses) + + # need to make new hypothesis that advance the constraints + new_indices = [] + new_seqs = [] + new_scores = [] + for seq_idx, pre_seq in enumerate(this_batch_input_ids): + new_state = sent_constraint_state[seq_idx] + print("\bRUN UPDATE\n") + new_state.update(pre_seq) + print("\nDONE UPDATE\n") + if not new_state.completed: + advance_tokens = new_state.advance() + print(">>>>>advance_tokens", advance_tokens) if advance_tokens.numel() != 0: - additional_num = advance_tokens.size(0) + advance_token = advance_tokens[:1] + new_seq = torch.cat((pre_seq, advance_token), dim=0) - # assume constraint ends up being added for now - sent_constraint_states[beam_idx].update(advance_tokens[0]) - # print("!!", stepped, completed) + new_score = this_batch_token_scores[seq_idx].take(advance_token[0]) - next_beam_tokens = ( - advance_tokens - .repeat(additional_num) - .long() - .to(device) - ) - sent_beam_tokens = torch.cat((sent_beam_tokens, next_beam_tokens)) + new_indices.append(seq_idx) + new_seqs.append(new_seq) + new_scores.append(new_score) - sent_beam_scores = torch.cat(( - sent_beam_scores, - sent_vocab_scores.take(advance_tokens) - )) + if len(new_indices) > 0: + new_indices = torch.add(torch.tensor(new_indices), batch_idx * orig_len).to(device) + new_seqs = torch.stack(new_seqs).to(device) + new_scores = torch.stack(new_scores).to(device) + # just force advancing + new_tokens = new_seqs[:, -1] + + sent_beam_tokens = torch.cat((sent_beam_tokens, new_tokens[:1]), -1) + sent_beam_scores = torch.cat((sent_beam_scores, new_scores[:1]), -1) + sent_beam_indices = torch.cat((sent_beam_indices, new_indices[:1]), -1) + print(">>>>>sent_beam_scores", sent_beam_scores) + print(">>>>>sent_beam_tokens", sent_beam_tokens) + print(">>>>>sent_beam_indices", sent_beam_indices) + - sent_beam_tokens = sent_beam_tokens[additional_num:] - sent_beam_scores = sent_beam_scores[additional_num:] + ''' + 2. Compute "banks" for each candidate. + If C is the number of constraints, we construct C banks, where ith bank + is the bank for candidates that have fulfilled ith constraint. + ''' - beam_idx += 1 - - break - else: - print("NO MORE ADVANCE") sent_beam_scores = sent_beam_scores[-orig_len:] sent_beam_tokens = sent_beam_tokens[-orig_len:] - assert sent_beam_scores.size(0) == orig_len + sent_beam_indices = sent_beam_indices[-orig_len:] + + + return sent_beam_scores, sent_beam_tokens, sent_beam_indices + + + + # if not sent_constraint_state.completed: + # advance_tokens = sent_constraint_state.advance() + # print(">>>>advance_tokens", advance_tokens) + # if advance_tokens.numel() != 0: + # advance_tokens = advance_tokens[:1] + + # sent_constraint_state.update(advance_tokens) + + # additional_num = advance_tokens.size(0) + # next_beam_tokens = ( + # advance_tokens + # .repeat(additional_num) + # .long() + # .to(device) + # ) + # sent_beam_tokens = torch.cat((sent_beam_tokens, next_beam_tokens)) + + # sent_beam_scores = torch.cat(( + # sent_beam_scores, + # sent_vocab_scores.take(advance_tokens) + # )) + + # ''' + # 2. Compute "banks" for each candidate. + # If C is the number of constraints, we construct C banks, where ith bank + # is the bank for candidates that have fulfilled ith constraint. + # ''' + + # sent_beam_scores = sent_beam_scores[-orig_len:] + # sent_beam_tokens = sent_beam_tokens[-orig_len:] + # print("???>>", sent_beam_tokens) + # sent_beam_indices = torch.tensor([ + # batch_idx*orig_len + 1, + # batch_idx*orig_len + 2, + # batch_idx*orig_len + 3, + # batch_idx*orig_len + 3, + # ]).to(device) + # print(">>>>>>>>>SELECTIUON INDICES??", sent_beam_indices) + # # sent_beam_indices = sent_beam_indices[-orig_len:] + # assert sent_beam_scores.size(0) == orig_len + + - return sent_constraint_states, sent_beam_scores, sent_beam_tokens + # return sent_constraint_state, sent_beam_scores, sent_beam_tokens, sent_beam_indices def finalize( diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 803bd5796f73..c1ebbf41b177 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -1298,7 +1298,6 @@ def generate( # 12. run beam search return self.constrained_beam_search( input_ids, - constraints=constraints, constrained_beam_scorer=constrained_beam_scorer, logits_processor=logits_processor, stopping_criteria=stopping_criteria, @@ -2723,7 +2722,6 @@ def group_beam_search( def constrained_beam_search( self, input_ids: torch.LongTensor, - constraints: List[Constraint], constrained_beam_scorer: ConstrainedBeamSearchScorer, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, @@ -2744,9 +2742,6 @@ def constrained_beam_search( input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. - constraints (`List[Constraint]`): - A list of positive constraints represented as `Constraint` objects that must be fulfilled in - the generation output. For more information, the documentation of [`Constraint`] should be read. constrained_beam_scorer (`ConstrainedBeamScorer`): An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and sorted during generation, while satisfying a list of positive constraints. For more information, the @@ -2841,22 +2836,11 @@ def constrained_beam_search( ... ] ... ) - >>> outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) + >>> outputs = model.constrained_beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) ```""" - # init values - print("constraints", constraints) - constraint_states = [ - [ - ConstraintListState([ - constraint.copy() - for constraint in constraints - ]) - for _ in range(constrained_beam_scorer.num_beams) - ] for _ in range(len(constrained_beam_scorer._beam_hyps)) - ] # need a constraint tracker for each sentence output, for each batch - + # init values logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() if max_length is not None: @@ -2905,7 +2889,6 @@ def constrained_beam_search( beam_scores[:, 1:] = -1e9 beam_scores = beam_scores.view((batch_size * num_beams,)) - print("<><><>input_ids", input_ids) this_peer_finished = False # used by synced_gpus only while True: @@ -2980,14 +2963,13 @@ def constrained_beam_search( next_indices = (next_tokens / vocab_size).long() next_tokens = next_tokens % vocab_size - print("!!!!!!scores_for_all_vocab", scores_for_all_vocab.size()) + print("<<<<<>>>>beam_next_tokens", beam_next_tokens) input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + print(">>>>>input_ids", input_ids) + model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder diff --git a/tests/test.py b/tests/test.py index 8ca1ac3b9424..288da94dfbb7 100644 --- a/tests/test.py +++ b/tests/test.py @@ -1,20 +1,25 @@ import unittest -from transformers import BartForConditionalGeneration, BartTokenizer +from transformers import GPT2Tokenizer, GPT2LMHeadModel from transformers.generation_beam_constraints import ( PhrasalConstraint ) device = "cuda" -model = BartForConditionalGeneration.from_pretrained("facebook/bart-base").to(device) -tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") +model = GPT2LMHeadModel.from_pretrained("gpt2").to(device) +tokenizer = GPT2Tokenizer.from_pretrained("gpt2") -force_text = "forced a little" -force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device) -print("force_tokens", force_tokens[0][1:-1]) -constraints = [PhrasalConstraint(force_tokens[0][1:-1])] +force_text = "talk" +force_text_2 = "forceful manner" +force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device)[0] +force_tokens_2 = tokenizer.encode(force_text_2, return_tensors="pt").to(device)[0] -input_text = ["this feels very"] * 10 +constraints = [ + PhrasalConstraint(force_tokens), + PhrasalConstraint(force_tokens_2) +] + +input_text = ["He always"] * 2 model_inputs = tokenizer(input_text, return_tensors="pt") @@ -24,7 +29,9 @@ print("model_inputs", model_inputs) k = model.generate( **model_inputs, - constraints=constraints + constraints=constraints, + num_beams=4, + num_return_sequences=3 ) for out in k: From 13c1808f0cd6b01fd0607f40b256e71c1acfcf8e Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Thu, 20 Jan 2022 12:28:17 +0000 Subject: [PATCH 04/34] technically works as desired, need to optimize and fix design choices leading to undersirable outputs --- ai.k | 910 ++++++++++++++++++ k.txt | 0 .../generation_beam_constraints.py | 59 +- src/transformers/generation_beam_search.py | 161 ++-- src/transformers/generation_utils.py | 16 +- tests/test.py | 13 +- 6 files changed, 1079 insertions(+), 80 deletions(-) create mode 100644 ai.k create mode 100644 k.txt diff --git a/ai.k b/ai.k new file mode 100644 index 000000000000..a831bd7f7f96 --- /dev/null +++ b/ai.k @@ -0,0 +1,910 @@ +============================= test session starts ============================== +platform linux -- Python 3.7.9, pytest-6.1.1, py-1.9.0, pluggy-0.13.1 +rootdir: /workspace/chanwookim/transformers, configfile: setup.cfg +plugins: hydra-core-1.0.6, flaky-3.7.0 +collected 0 items / 1 error + +==================================== ERRORS ==================================== +________________________ ERROR collecting tests/test.py ________________________ +tests/test.py:43: in + assert False +E assert False +------------------------------- Captured stdout -------------------------------- +force_tokens tensor([ 1263, 10963], device='cuda:0') +force_tokens_2 tensor([7165], device='cuda:0') +model_inputs {'input_ids': tensor([[ 464, 5156, 318, 13774, 780]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1]], device='cuda:0')} + + +INPUT + +The baby is crying because +The baby is crying because +The baby is crying because +The baby is crying because +The baby is crying because + + + + +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 1263] +new_seqs [] +advance_seq [464, 5156, 318, 13774, 780, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 1263]] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +PREZIP tensor([[-1.4358e+00, 0.0000e+00, 6.7300e+02], + [-2.0147e+00, 0.0000e+00, 3.3900e+02], + [-2.3535e+00, 0.0000e+00, 3.4000e+02], + [-2.4315e+00, 0.0000e+00, 2.8600e+02], + [-2.7982e+00, 0.0000e+00, 6.0700e+02], + [-1.0072e+01, 1.0000e+00, 1.2630e+03], + [-1.2223e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') +POSTZIP tensor([[-1.0072e+01, 1.0000e+00, 1.2630e+03], + [-1.4358e+00, 0.0000e+00, 6.7300e+02], + [-1.2223e+01, 1.0000e+00, 7.1650e+03], + [-2.0147e+00, 0.0000e+00, 3.3900e+02], + [-2.3535e+00, 0.0000e+00, 3.4000e+02], + [-2.4315e+00, 0.0000e+00, 2.8600e+02], + [-2.7982e+00, 0.0000e+00, 6.0700e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([-10.0721, -1.4358, -12.2228, -2.0147, -2.3535], device='cuda:0') +>>>>>sent_beam_tokens tensor([1263, 673, 7165, 339, 340], device='cuda:0') +>>>>>sent_beam_indices tensor([0, 0, 0, 0, 0], device='cuda:0') +>>>>constraints_completed [False, False, False, False, False] + + +OUTPUT + +The baby is crying because big +The baby is crying because she +The baby is crying because crazy +The baby is crying because he +The baby is crying because it + + + + + + +INPUT + +The baby is crying because big +The baby is crying because she +The baby is crying because crazy +The baby is crying because he +The baby is crying because it + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963]] +advance_seq [464, 5156, 318, 13774, 780, 673, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 7165, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165]] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 339, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263]] +advance_seq [464, 5156, 318, 13774, 780, 339, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263]] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 340, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263], [464, 5156, 318, 13774, 780, 339, 7165]] +advance_seq [464, 5156, 318, 13774, 780, 340, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263], [464, 5156, 318, 13774, 780, 339, 7165], [464, 5156, 318, 13774, 780, 340, 1263]] +PREZIP tensor([[-3.3543e+00, 0.0000e+00, 3.3800e+02], + [-3.7313e+00, 0.0000e+00, 3.3800e+02], + [-3.7631e+00, 0.0000e+00, 3.7300e+02], + [-3.8234e+00, 0.0000e+00, 3.3800e+02], + [-3.8247e+00, 0.0000e+00, 3.1800e+02], + [-6.6173e+00, 2.0000e+00, 1.0963e+04], + [-1.2264e+01, 1.0000e+00, 1.2630e+03], + [-1.3972e+01, 1.0000e+00, 7.1650e+03], + [-5.7579e+00, 3.0000e+00, 1.2630e+03], + [-1.1960e+01, 1.0000e+00, 1.2630e+03], + [-1.3683e+01, 1.0000e+00, 7.1650e+03], + [-1.2435e+01, 1.0000e+00, 1.2630e+03], + [-1.3548e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') +POSTZIP tensor([[-5.7579e+00, 3.0000e+00, 1.2630e+03], + [-6.6173e+00, 2.0000e+00, 1.0963e+04], + [-1.1960e+01, 1.0000e+00, 1.2630e+03], + [-3.3543e+00, 0.0000e+00, 3.3800e+02], + [-1.2264e+01, 1.0000e+00, 1.2630e+03], + [-3.7313e+00, 0.0000e+00, 3.3800e+02], + [-1.2435e+01, 1.0000e+00, 1.2630e+03], + [-3.7631e+00, 0.0000e+00, 3.7300e+02], + [-1.3548e+01, 1.0000e+00, 7.1650e+03], + [-3.8234e+00, 0.0000e+00, 3.3800e+02], + [-1.3683e+01, 1.0000e+00, 7.1650e+03], + [-3.8247e+00, 0.0000e+00, 3.1800e+02], + [-1.3972e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -5.7579, -6.6173, -11.9598, -3.3543, -12.2640], device='cuda:0') +>>>>>sent_beam_tokens tensor([ 1263, 10963, 1263, 338, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([2, 0, 3, 1, 1], device='cuda:0') +>>>>constraints_completed [False, False, False, False, False] + + +OUTPUT + +The baby is crying because crazy big +The baby is crying because big monsters +The baby is crying because he big +The baby is crying because she's +The baby is crying because she big + + + + + + +INPUT + +The baby is crying because crazy big +The baby is crying because big monsters +The baby is crying because he big +The baby is crying because she's +The baby is crying because she big + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 7165, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 1263, 10963, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 339, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165]] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963]] +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 1263]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 1263], [464, 5156, 318, 13774, 780, 673, 338, 7165]] +PREZIP tensor([[-5.8691e+00, 0.0000e+00, 5.8700e+02], + [-6.3662e+00, 0.0000e+00, 4.0700e+02], + [-6.4227e+00, 0.0000e+00, 1.2008e+04], + [-6.4356e+00, 0.0000e+00, 7.7870e+03], + [-6.7641e+00, 0.0000e+00, 5.2300e+02], + [-5.7414e+00, 4.0000e+00, 1.0963e+04], + [-1.2468e+01, 4.0000e+00, 7.1650e+03], + [-1.0331e+01, 2.0000e+00, 1.0963e+04], + [-8.2352e+00, 1.0000e+00, 1.2630e+03], + [-7.5897e+00, 1.0000e+00, 7.1650e+03], + [-1.0610e+01, 2.0000e+00, 1.0963e+04]], device='cuda:0') +POSTZIP tensor([[-5.7414e+00, 4.0000e+00, 1.0963e+04], + [-1.0331e+01, 2.0000e+00, 1.0963e+04], + [-7.5897e+00, 1.0000e+00, 7.1650e+03], + [-5.8691e+00, 0.0000e+00, 5.8700e+02], + [-1.2468e+01, 4.0000e+00, 7.1650e+03], + [-1.0610e+01, 2.0000e+00, 1.0963e+04], + [-8.2352e+00, 1.0000e+00, 1.2630e+03], + [-6.3662e+00, 0.0000e+00, 4.0700e+02], + [-6.4227e+00, 0.0000e+00, 1.2008e+04], + [-6.4356e+00, 0.0000e+00, 7.7870e+03], + [-6.7641e+00, 0.0000e+00, 5.2300e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -5.7414, -10.3309, -7.5897, -5.8691, -12.4677], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 10963, 7165, 587, 7165], device='cuda:0') +>>>>>sent_beam_indices tensor([0, 2, 3, 3, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because crazy big monsters +The baby is crying because he big monsters +The baby is crying because she's crazy +The baby is crying because she's been +The baby is crying because big monsters crazy + + + + + + +INPUT + +The baby is crying because crazy big monsters +The baby is crying because he big monsters +The baby is crying because she's crazy +The baby is crying because she's been +The baby is crying because big monsters crazy + + + + +>>>>>advance_tokens tensor([7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165]] +>>>>>advance_tokens tensor([1263, 7165], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263]] +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165] +new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263], [464, 5156, 318, 13774, 780, 673, 338, 587, 1263]] +PREZIP tensor([[-6.8971e+00, 2.0000e+00, 3.8900e+02], + [-8.0195e+00, 0.0000e+00, 4.2300e+02], + [-8.5441e+00, 0.0000e+00, 1.2970e+03], + [-8.9806e+00, 2.0000e+00, 1.3000e+01], + [-9.2211e+00, 0.0000e+00, 1.6110e+04], + [-1.1526e+01, 4.0000e+00, 7.1650e+03], + [-7.5428e+00, 3.0000e+00, 1.2630e+03], + [-1.1152e+01, 1.0000e+00, 1.2630e+03], + [-9.3288e+00, 1.0000e+00, 7.1650e+03]], device='cuda:0') +POSTZIP tensor([[-1.1526e+01, 4.0000e+00, 7.1650e+03], + [-7.5428e+00, 3.0000e+00, 1.2630e+03], + [-6.8971e+00, 2.0000e+00, 3.8900e+02], + [-9.3288e+00, 1.0000e+00, 7.1650e+03], + [-8.0195e+00, 0.0000e+00, 4.2300e+02], + [-8.9806e+00, 2.0000e+00, 1.3000e+01], + [-1.1152e+01, 1.0000e+00, 1.2630e+03], + [-8.5441e+00, 0.0000e+00, 1.2970e+03], + [-9.2211e+00, 0.0000e+00, 1.6110e+04]], device='cuda:0') +>>>>>sent_beam_scores tensor([-11.5259, -7.5428, -6.8971, -9.3288, -8.0195], device='cuda:0') +>>>>>sent_beam_tokens tensor([7165, 1263, 389, 7165, 423], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 2, 0, 3, 0], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because he big monsters crazy +The baby is crying because she's crazy big +The baby is crying because crazy big monsters are +The baby is crying because she's been crazy +The baby is crying because crazy big monsters have + + + + + + +INPUT + +The baby is crying because he big monsters crazy +The baby is crying because she's crazy big +The baby is crying because crazy big monsters are +The baby is crying because she's been crazy +The baby is crying because crazy big monsters have + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 10963]] +PREZIP tensor([[-8.8730e+00, 0.0000e+00, 2.4060e+03], + [-8.9872e+00, 0.0000e+00, 2.9000e+02], + [-9.2539e+00, 0.0000e+00, 1.3000e+01], + [-9.5069e+00, 0.0000e+00, 1.1000e+01], + [-9.6623e+00, 0.0000e+00, 5.5300e+02], + [-1.2619e+01, 4.0000e+00, 1.0963e+04], + [-8.8750e+00, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-1.2619e+01, 4.0000e+00, 1.0963e+04], + [-8.8750e+00, 3.0000e+00, 1.2630e+03], + [-8.8730e+00, 0.0000e+00, 2.4060e+03], + [-8.9872e+00, 0.0000e+00, 2.9000e+02], + [-9.2539e+00, 0.0000e+00, 1.3000e+01], + [-9.5069e+00, 0.0000e+00, 1.1000e+01], + [-9.6623e+00, 0.0000e+00, 5.5300e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([-12.6191, -8.8750, -8.8730, -8.9872, -9.2539], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 2406, 290, 13], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 3, 2, 1, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, True, False] + + +OUTPUT + +The baby is crying because she's crazy big monsters +The baby is crying because she's been crazy big +The baby is crying because crazy big monsters are coming +The baby is crying because she's crazy big and +The baby is crying because she's crazy big. + + + + + + +INPUT + +The baby is crying because she's crazy big monsters +The baby is crying because she's been crazy big +The baby is crying because crazy big monsters are coming +The baby is crying because she's crazy big and +The baby is crying because she's crazy big. + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 10963]] +PREZIP tensor([[-1.0600e+01, 0.0000e+00, 6.7300e+02], + [-1.0826e+01, 2.0000e+00, 2.9000e+02], + [-1.0833e+01, 2.0000e+00, 1.3750e+03], + [-1.0837e+01, 0.0000e+00, 1.3000e+01], + [-1.1092e+01, 4.0000e+00, 1.3000e+01], + [-1.2403e+01, 4.0000e+00, 1.0963e+04], + [-1.5209e+01, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-1.1092e+01, 4.0000e+00, 1.3000e+01], + [-1.5209e+01, 3.0000e+00, 1.2630e+03], + [-1.0826e+01, 2.0000e+00, 2.9000e+02], + [-1.0600e+01, 0.0000e+00, 6.7300e+02], + [-1.2403e+01, 4.0000e+00, 1.0963e+04], + [-1.0833e+01, 2.0000e+00, 1.3750e+03], + [-1.0837e+01, 0.0000e+00, 1.3000e+01]], device='cuda:0') +>>>>>sent_beam_scores tensor([-11.0919, -15.2090, -10.8264, -10.5999, -12.4034], device='cuda:0') +>>>>>sent_beam_tokens tensor([ 13, 1263, 290, 673, 10963], device='cuda:0') +>>>>>sent_beam_indices tensor([2, 4, 1, 3, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because crazy big monsters are coming. +The baby is crying because she's crazy big. big +The baby is crying because she's been crazy big and +The baby is crying because she's crazy big and she +The baby is crying because she's been crazy big monsters + + + + + + +INPUT + +The baby is crying because crazy big monsters are coming. +The baby is crying because she's crazy big. big +The baby is crying because she's been crazy big and +The baby is crying because she's crazy big and she +The baby is crying because she's been crazy big monsters + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263]] +PREZIP tensor([[-1.1419e+01, 2.0000e+00, 3.3800e+02], + [-1.2924e+01, 0.0000e+00, 6.7300e+02], + [-1.3312e+01, 2.0000e+00, 1.9800e+02], + [-1.3459e+01, 2.0000e+00, 3.3820e+03], + [-1.3549e+01, 0.0000e+00, 4.6000e+02], + [-9.9536e+00, 4.0000e+00, 1.0963e+04], + [-3.8742e+00, 3.0000e+00, 1.2630e+03], + [-1.0987e+01, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-9.9536e+00, 4.0000e+00, 1.0963e+04], + [-3.8742e+00, 3.0000e+00, 1.2630e+03], + [-1.1419e+01, 2.0000e+00, 3.3800e+02], + [-1.2924e+01, 0.0000e+00, 6.7300e+02], + [-1.0987e+01, 3.0000e+00, 1.2630e+03], + [-1.3312e+01, 2.0000e+00, 1.9800e+02], + [-1.3549e+01, 0.0000e+00, 4.6000e+02], + [-1.3459e+01, 2.0000e+00, 3.3820e+03]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -9.9536, -3.8742, -11.4194, -12.9239, -10.9866], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 338, 673, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 2, 3, 2, 3], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because she's crazy big. big monsters +The baby is crying because she's been crazy big and big +The baby is crying because she's crazy big and she's +The baby is crying because she's been crazy big and she +The baby is crying because she's crazy big and she big + + + + + + +INPUT + +The baby is crying because she's crazy big. big monsters +The baby is crying because she's been crazy big and big +The baby is crying because she's crazy big and she's +The baby is crying because she's been crazy big and she +The baby is crying because she's crazy big and she big + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 673, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 673, 1263]] +PREZIP tensor([[-5.2463e+00, 0.0000e+00, 2.9000e+02], + [-5.9868e+00, 2.0000e+00, 1.3000e+01], + [-6.1505e+00, 0.0000e+00, 1.1000e+01], + [-6.1889e+00, 0.0000e+00, 5.5300e+02], + [-6.7390e+00, 0.0000e+00, 3.2900e+02], + [-1.1919e+01, 4.0000e+00, 1.0963e+04], + [-5.6936e+00, 3.0000e+00, 1.2630e+03], + [-1.1110e+01, 3.0000e+00, 1.2630e+03], + [-9.3693e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') +POSTZIP tensor([[-9.3693e+00, 4.0000e+00, 1.0963e+04], + [-5.6936e+00, 3.0000e+00, 1.2630e+03], + [-5.9868e+00, 2.0000e+00, 1.3000e+01], + [-5.2463e+00, 0.0000e+00, 2.9000e+02], + [-1.1919e+01, 4.0000e+00, 1.0963e+04], + [-1.1110e+01, 3.0000e+00, 1.2630e+03], + [-6.1505e+00, 0.0000e+00, 1.1000e+01], + [-6.1889e+00, 0.0000e+00, 5.5300e+02], + [-6.7390e+00, 0.0000e+00, 3.2900e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -9.3693, -5.6936, -5.9868, -5.2463, -11.9188], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') +>>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because she's crazy big and she big monsters +The baby is crying because she's crazy big and she's big +The baby is crying because she's been crazy big and big. +The baby is crying because she's been crazy big and big and +The baby is crying because she's been crazy big and big monsters + + + + + + +INPUT + +The baby is crying because she's crazy big and she big monsters +The baby is crying because she's crazy big and she's big +The baby is crying because she's been crazy big and big. +The baby is crying because she's been crazy big and big and +The baby is crying because she's been crazy big and big monsters + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263]] +PREZIP tensor([[-6.3694e+00, 0.0000e+00, 1.2630e+03], + [-6.8644e+00, 2.0000e+00, 2.9000e+02], + [-7.4970e+00, 2.0000e+00, 1.3750e+03], + [-7.5834e+00, 2.0000e+00, 6.7300e+02], + [-7.8789e+00, 0.0000e+00, 1.5760e+03], + [-1.2507e+01, 4.0000e+00, 1.0963e+04], + [-1.6012e+01, 3.0000e+00, 1.2630e+03], + [-1.1232e+00, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-1.2507e+01, 4.0000e+00, 1.0963e+04], + [-1.1232e+00, 3.0000e+00, 1.2630e+03], + [-6.8644e+00, 2.0000e+00, 2.9000e+02], + [-6.3694e+00, 0.0000e+00, 1.2630e+03], + [-1.6012e+01, 3.0000e+00, 1.2630e+03], + [-7.4970e+00, 2.0000e+00, 1.3750e+03], + [-7.8789e+00, 0.0000e+00, 1.5760e+03], + [-7.5834e+00, 2.0000e+00, 6.7300e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([-12.5067, -1.1232, -6.8644, -6.3694, -16.0120], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 290, 1263, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 3, 1, 3, 2], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because she's crazy big and she's big monsters +The baby is crying because she's been crazy big and big and big +The baby is crying because she's crazy big and she's big and +The baby is crying because she's been crazy big and big and big +The baby is crying because she's been crazy big and big. big + + + + + + +INPUT + +The baby is crying because she's crazy big and she's big monsters +The baby is crying because she's been crazy big and big and big +The baby is crying because she's crazy big and she's big and +The baby is crying because she's been crazy big and big and big +The baby is crying because she's been crazy big and big. big + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263]] +PREZIP tensor([[-2.4862e+00, 0.0000e+00, 2.9000e+02], + [-2.8610e+00, 2.0000e+00, 1.3000e+01], + [-3.3225e+00, 0.0000e+00, 5.5300e+02], + [-3.4911e+00, 0.0000e+00, 1.1000e+01], + [-3.9857e+00, 0.0000e+00, 3.2900e+02], + [-1.2293e+01, 4.0000e+00, 1.0963e+04], + [-4.9025e+00, 3.0000e+00, 1.2630e+03], + [-9.6560e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') +POSTZIP tensor([[-9.6560e+00, 4.0000e+00, 1.0963e+04], + [-4.9025e+00, 3.0000e+00, 1.2630e+03], + [-2.8610e+00, 2.0000e+00, 1.3000e+01], + [-2.4862e+00, 0.0000e+00, 2.9000e+02], + [-1.2293e+01, 4.0000e+00, 1.0963e+04], + [-3.3225e+00, 0.0000e+00, 5.5300e+02], + [-3.4911e+00, 0.0000e+00, 1.1000e+01], + [-3.9857e+00, 0.0000e+00, 3.2900e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -9.6560, -4.9025, -2.8610, -2.4862, -12.2932], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') +>>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because she's been crazy big and big. big monsters +The baby is crying because she's crazy big and she's big and big +The baby is crying because she's been crazy big and big and big. +The baby is crying because she's been crazy big and big and big and +The baby is crying because she's been crazy big and big and big monsters + + + + + + +INPUT + +The baby is crying because she's been crazy big and big. big monsters +The baby is crying because she's crazy big and she's big and big +The baby is crying because she's been crazy big and big and big. +The baby is crying because she's been crazy big and big and big and +The baby is crying because she's been crazy big and big and big monsters + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263]] +PREZIP tensor([[-3.3472e+00, 0.0000e+00, 1.2630e+03], + [-4.4136e+00, 0.0000e+00, 1.3750e+03], + [-5.0307e+00, 2.0000e+00, 8.4300e+02], + [-5.0600e+00, 0.0000e+00, 3.1400e+02], + [-5.1705e+00, 0.0000e+00, 1.9800e+02], + [-1.2517e+01, 4.0000e+00, 1.0963e+04], + [-1.6495e+01, 3.0000e+00, 1.2630e+03], + [-8.6100e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-1.2517e+01, 4.0000e+00, 1.0963e+04], + [-8.6100e-01, 3.0000e+00, 1.2630e+03], + [-5.0307e+00, 2.0000e+00, 8.4300e+02], + [-3.3472e+00, 0.0000e+00, 1.2630e+03], + [-1.6495e+01, 3.0000e+00, 1.2630e+03], + [-4.4136e+00, 0.0000e+00, 1.3750e+03], + [-5.0600e+00, 0.0000e+00, 3.1400e+02], + [-5.1705e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([-12.5170, -0.8610, -5.0307, -3.3472, -16.4952], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because she's crazy big and she's big and big monsters +The baby is crying because she's been crazy big and big and big and big +The baby is crying because she's been crazy big and big and big. And +The baby is crying because she's been crazy big and big and big and big +The baby is crying because she's been crazy big and big and big. big + + + + + + +INPUT + +The baby is crying because she's crazy big and she's big and big monsters +The baby is crying because she's been crazy big and big and big and big +The baby is crying because she's been crazy big and big and big. And +The baby is crying because she's been crazy big and big and big and big +The baby is crying because she's been crazy big and big and big. big + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263]] +PREZIP tensor([[-2.1940e+00, 0.0000e+00, 2.9000e+02], + [-2.4520e+00, 2.0000e+00, 1.3000e+01], + [-3.2742e+00, 0.0000e+00, 5.5300e+02], + [-3.3400e+00, 0.0000e+00, 1.1000e+01], + [-3.6658e+00, 0.0000e+00, 5.2600e+02], + [-1.2107e+01, 4.0000e+00, 1.0963e+04], + [-8.4168e+00, 3.0000e+00, 1.2630e+03], + [-9.5848e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') +POSTZIP tensor([[-9.5848e+00, 4.0000e+00, 1.0963e+04], + [-8.4168e+00, 3.0000e+00, 1.2630e+03], + [-2.4520e+00, 2.0000e+00, 1.3000e+01], + [-2.1940e+00, 0.0000e+00, 2.9000e+02], + [-1.2107e+01, 4.0000e+00, 1.0963e+04], + [-3.2742e+00, 0.0000e+00, 5.5300e+02], + [-3.3400e+00, 0.0000e+00, 1.1000e+01], + [-3.6658e+00, 0.0000e+00, 5.2600e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -9.5848, -8.4168, -2.4520, -2.1940, -12.1066], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') +>>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because she's been crazy big and big and big. big monsters +The baby is crying because she's been crazy big and big and big. And big +The baby is crying because she's been crazy big and big and big and big. +The baby is crying because she's been crazy big and big and big and big and +The baby is crying because she's been crazy big and big and big and big monsters + + + + + + +INPUT + +The baby is crying because she's been crazy big and big and big. big monsters +The baby is crying because she's been crazy big and big and big. And big +The baby is crying because she's been crazy big and big and big and big. +The baby is crying because she's been crazy big and big and big and big and +The baby is crying because she's been crazy big and big and big and big monsters + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263]] +PREZIP tensor([[-2.8488e+00, 0.0000e+00, 1.2630e+03], + [-4.0466e+00, 0.0000e+00, 1.3750e+03], + [-4.4985e+00, 2.0000e+00, 8.4300e+02], + [-4.6394e+00, 0.0000e+00, 3.1400e+02], + [-4.7340e+00, 0.0000e+00, 1.9800e+02], + [-8.2042e+00, 4.0000e+00, 1.0963e+04], + [-1.6431e+01, 3.0000e+00, 1.2630e+03], + [-6.5482e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-8.2042e+00, 4.0000e+00, 1.0963e+04], + [-6.5482e-01, 3.0000e+00, 1.2630e+03], + [-4.4985e+00, 2.0000e+00, 8.4300e+02], + [-2.8488e+00, 0.0000e+00, 1.2630e+03], + [-1.6431e+01, 3.0000e+00, 1.2630e+03], + [-4.0466e+00, 0.0000e+00, 1.3750e+03], + [-4.6394e+00, 0.0000e+00, 3.1400e+02], + [-4.7340e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -8.2042, -0.6548, -4.4985, -2.8488, -16.4311], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because she's been crazy big and big and big. And big monsters +The baby is crying because she's been crazy big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big. And +The baby is crying because she's been crazy big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big. big + + + + + + +INPUT + +The baby is crying because she's been crazy big and big and big. And big monsters +The baby is crying because she's been crazy big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big. And +The baby is crying because she's been crazy big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big. big + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963]] +>>>>>advance_tokens tensor([10963], device='cuda:0') +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263, 10963] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263]] +PREZIP tensor([[-1.7995e+00, 0.0000e+00, 2.9000e+02], + [-2.2400e+00, 2.0000e+00, 1.3000e+01], + [-3.2673e+00, 0.0000e+00, 1.1000e+01], + [-3.4144e+00, 0.0000e+00, 5.5300e+02], + [-3.4721e+00, 0.0000e+00, 5.2600e+02], + [-1.1857e+01, 4.0000e+00, 1.0963e+04], + [-8.4926e+00, 3.0000e+00, 1.2630e+03], + [-9.7216e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') +POSTZIP tensor([[-9.7216e+00, 4.0000e+00, 1.0963e+04], + [-8.4926e+00, 3.0000e+00, 1.2630e+03], + [-2.2400e+00, 2.0000e+00, 1.3000e+01], + [-1.7995e+00, 0.0000e+00, 2.9000e+02], + [-1.1857e+01, 4.0000e+00, 1.0963e+04], + [-3.2673e+00, 0.0000e+00, 1.1000e+01], + [-3.4144e+00, 0.0000e+00, 5.5300e+02], + [-3.4721e+00, 0.0000e+00, 5.2600e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -9.7216, -8.4926, -2.2400, -1.7995, -11.8567], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') +>>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') +>>>>constraints_completed [True, False, False, False, True] + + +OUTPUT + +The baby is crying because she's been crazy big and big and big and big. big monsters +The baby is crying because she's been crazy big and big and big and big. And big +The baby is crying because she's been crazy big and big and big and big and big. +The baby is crying because she's been crazy big and big and big and big and big and +The baby is crying because she's been crazy big and big and big and big and big monsters + + + + + + +INPUT + +The baby is crying because she's been crazy big and big and big and big. big monsters +The baby is crying because she's been crazy big and big and big and big. And big +The baby is crying because she's been crazy big and big and big and big and big. +The baby is crying because she's been crazy big and big and big and big and big and +The baby is crying because she's been crazy big and big and big and big and big monsters + + + + +>>>>>advance_tokens tensor([10963], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963] +new_seqs [] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963]] +>>>>>advance_tokens tensor([1263], device='cuda:0') +advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263] +new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263]] +PREZIP tensor([[-2.2907e+00, 0.0000e+00, 1.2630e+03], + [-3.8817e+00, 0.0000e+00, 1.3750e+03], + [-4.2239e+00, 2.0000e+00, 8.4300e+02], + [-4.4378e+00, 0.0000e+00, 3.1400e+02], + [-4.4513e+00, 0.0000e+00, 1.9800e+02], + [-8.3479e+00, 4.0000e+00, 1.0963e+04], + [-1.6134e+01, 3.0000e+00, 1.2630e+03], + [-4.9122e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') +POSTZIP tensor([[-8.3479e+00, 4.0000e+00, 1.0963e+04], + [-4.9122e-01, 3.0000e+00, 1.2630e+03], + [-4.2239e+00, 2.0000e+00, 8.4300e+02], + [-2.2907e+00, 0.0000e+00, 1.2630e+03], + [-1.6134e+01, 3.0000e+00, 1.2630e+03], + [-3.8817e+00, 0.0000e+00, 1.3750e+03], + [-4.4378e+00, 0.0000e+00, 3.1400e+02], + [-4.4513e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') +>>>>>sent_beam_scores tensor([ -8.3479, -0.4912, -4.2239, -2.2907, -16.1338], device='cuda:0') +>>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') +>>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') +>>>>constraints_completed [True, False, False, False, False] + + +OUTPUT + +The baby is crying because she's been crazy big and big and big and big. And big monsters +The baby is crying because she's been crazy big and big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big and big. And +The baby is crying because she's been crazy big and big and big and big and big and big +The baby is crying because she's been crazy big and big and big and big and big. big + + + + +!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, + 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263], + device='cuda:0') +The baby is crying because she's been crazy big and big and big and big and big and big +!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, + 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263], + device='cuda:0') +The baby is crying because she's been crazy big and big and big and big and big and big +!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, + 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843], + device='cuda:0') +The baby is crying because she's been crazy big and big and big and big and big. And +!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, + 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], + device='cuda:0') +The baby is crying because she's been crazy big and big and big and big. And big monsters +!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, + 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263], + device='cuda:0') +The baby is crying because she's been crazy big and big and big and big and big. big +------------------------------- Captured stderr -------------------------------- +2022-01-20 12:25:52.703100: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory +2022-01-20 12:25:52.703133: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine. +Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation. +=============================== warnings summary =============================== +../../../opt/conda/envs/ACW/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py:22 + /opt/conda/envs/ACW/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py:22: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses + import imp + +-- Docs: https://docs.pytest.org/en/stable/warnings.html +=========================== short test summary info ============================ +ERROR tests/test.py - assert False +!!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!! +========================= 1 warning, 1 error in 22.90s ========================= diff --git a/k.txt b/k.txt new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 95986b0feb1d..def4e3783681 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -133,8 +133,18 @@ def __init__(self, token_ids: torch.Tensor): self.fulfilled_idx = -1 # the index of the currently fulfilled step self.completed = False - def copy(self): - return PhrasalConstraint(self.token_ids) + def remaining(self): + return self.seqlen - (self.fulfilled_idx + 1) + + def copy(self, stateful=False): + new_constraint = PhrasalConstraint(self.token_ids) + + if stateful: + new_constraint.seq_len = self.seqlen + new_constraint.fulfilled_idx = self.fulfilled_idx + new_constraint.completed = self.completed + + return new_constraint def advance(self): return self.token_ids[self.fulfilled_idx + 1] @@ -150,6 +160,11 @@ def update(self, token_id: int): completed = False reset = False +# ''' +# ------------------------------- Captured stdout -------------------------------- +# force_tokens tensor([48186, 220], device='cuda:0') +# force_tokens_2 tensor([3174, 913, 220], device='cuda:0') +# ''' if self.does_advance(token_id): self.fulfilled_idx += 1 stepped = True @@ -167,20 +182,48 @@ def reset(self): self.fulfilled_idx = 0 +# force_tokens tensor([3375, 220], device='cuda:0') +# force_tokens_2 tensor([23112], device='cuda:0') + class ConstraintListState: def __init__(self, constraints: List[Constraint]): self.constraints = constraints + + self.max_seqlen = max([c.seqlen for c in constraints]) self.n_constraints = len(constraints) self.completed = False self.init_state() + def copy(self, stateful=True): + new_state = ConstraintListState(self.constraints) + + if stateful: + new_state.complete_constraints = [ + constraint.copy(stateful=True) + for constraint in self.complete_constraints + ] + if self.inprogress_constraint is not None: + new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True) + new_state.pending_constraints = [ + constraint.copy() + for constraint in self.pending_constraints + ] + + return new_state def init_state(self): self.complete_constraints = [] self.inprogress_constraint = None self.pending_constraints = [constraint.copy() for constraint in self.constraints] + def get_bank(self): + add = 0 + if self.inprogress_constraint: + add += self.max_seqlen - self.inprogress_constraint.remaining() + + return (len(self.complete_constraints)*self.max_seqlen) + add + def advance(self): '''The list of tokens to generate such that we can make progress. @@ -189,8 +232,6 @@ def advance(self): we'll return. ''' if self.inprogress_constraint is None: - print("RUN ADVANCE NO PROGRESS") - print("self.pending_constraints", self.pending_constraints) token_list = [] for constraint in self.pending_constraints: advance = constraint.advance() @@ -209,18 +250,15 @@ def update(self, token_ids: torch.Tensor): progress through constraints. ''' self.init_state() - print("\n!!!START UPDATE\n\n") for token in token_ids: - print("token", token) complete, stepped = self.add(token) - print("complete, stepped", complete, stepped) return self def add(self, token_id: int): if self.completed: - return + return True, True complete, stepped = False, False if self.inprogress_constraint is not None: @@ -235,10 +273,8 @@ def add(self, token_id: int): 1. If the next token breaks the fulfillment, then we must restart. e.g. force the sequence "I love pies" and the next token after "I love" is "books". ''' - print("RESET") self.pending_constraints.append(self.inprogress_constraint.copy()) self.inprogress_constraint = None - print("self.pending_constraints", self.pending_constraints) if complete: ''' @@ -255,7 +291,6 @@ def add(self, token_id: int): ''' Not in the middle of fulfilling a constraint. ''' - print("NOT IN THE MIDDLE", self.pending_constraints) for cidx, pending_constraint in enumerate(self.pending_constraints): ''' 1. Does it advance any of the pending constraints? @@ -270,8 +305,6 @@ def add(self, token_id: int): if complete or stepped: self.pending_constraints = self.pending_constraints[:cidx] + self.pending_constraints[cidx+1:] - print("!!!self.pending_constraints", self.pending_constraints) - print() if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: self.completed = True diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 23e07efe9445..5ec0754f1a57 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -19,6 +19,7 @@ from typing import List, Optional, Tuple import torch +import numpy as np from .generation_beam_constraints import Constraint, ConstraintListState from .file_utils import add_start_docstrings @@ -403,6 +404,10 @@ def __init__( self.num_beam_groups = num_beam_groups self.group_size = self.num_beams // self.num_beam_groups self.constraints = constraints + self.constraints_completed = [ + [False] * self.group_size + ] * batch_size + self._is_init = False self._beam_hyps = [ BeamHypotheses( @@ -500,12 +505,18 @@ def process( # add to generated hypotheses if end of sentence if (eos_token_id is not None) and (next_token.item() == eos_token_id): # if constraint not fulfilled, it should not be added. - if not constraint_states[batch_idx].completed: + if not self.constraints_completed[batch_idx][next_index]: continue + else: + print("NOT COMPLETE", self.constraints_completed) + # if beam_token does not belong to top num_beams tokens, it should not be added is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size if is_beam_token_worse_than_top_num_beams: continue + + print("\n!\n!\n!\n!\n!\n!\n!\n!\n!TRULY COMPLETED SEQUENCE!!\n!\n!\n!") + print("input_ids[batch_beam_idx].clone()", input_ids[batch_beam_idx].clone()) beam_hyp.add( input_ids[batch_beam_idx].clone(), next_score.item(), @@ -520,27 +531,23 @@ def process( # once the beam for next step is full, don't add more tokens to it. if beam_idx == self.group_size: break - - print("scores_for_all_vocab", scores_for_all_vocab.size()) - print("input_ids", input_ids.size()) - print("next_beam_tokens", next_beam_tokens.size()) - new_scores, new_tokens, new_indices = self.step_sentence_constraint( + new_completed, new_scores, new_tokens, new_indices = self.step_sentence_constraint( batch_idx, input_ids, scores_for_all_vocab, + self.constraints_completed[batch_idx], next_beam_scores[batch_idx].clone(), next_beam_tokens[batch_idx].clone(), next_beam_indices[batch_idx].clone(), ) - print("!!new_scores", new_scores) - print("!!new_tokens", new_tokens) - print() next_beam_scores[batch_idx] = new_scores next_beam_tokens[batch_idx] = new_tokens next_beam_indices[batch_idx] = new_indices + self.constraints_completed[batch_idx] = new_completed + if beam_idx < self.group_size: raise ValueError( @@ -565,6 +572,7 @@ def step_sentence_constraint( batch_idx, input_ids, vocab_scores, + constraints_completed, sent_beam_scores, sent_beam_tokens, sent_beam_indices, @@ -582,7 +590,9 @@ def step_sentence_constraint( ''' orig_len = sent_beam_indices.size(0) device = sent_beam_indices.get_device() - sent_constraint_state = self.make_constraint_states(orig_len) + + topk_contraint_states = self.make_constraint_states(orig_len) + advance_constraint_states = self.make_constraint_states(orig_len) start_idx = batch_idx*orig_len end_idx = (batch_idx+1) * orig_len @@ -590,64 +600,101 @@ def step_sentence_constraint( this_batch_input_ids = input_ids[start_idx : end_idx] this_batch_token_scores = vocab_scores[start_idx : end_idx] - print("this_batch_input_ids", this_batch_input_ids) - print("sent_beam_tokens", sent_beam_tokens) - full_hypotheses = torch.cat((this_batch_input_ids, sent_beam_tokens.unsqueeze(-1)), dim=-1) - print("full_hypotheses", full_hypotheses) # need to make new hypothesis that advance the constraints - new_indices = [] new_seqs = [] + new_states = [] + new_indices = [] + new_tokens = [] new_scores = [] for seq_idx, pre_seq in enumerate(this_batch_input_ids): - new_state = sent_constraint_state[seq_idx] - print("\bRUN UPDATE\n") - new_state.update(pre_seq) - print("\nDONE UPDATE\n") - if not new_state.completed: - advance_tokens = new_state.advance() + topk_state = topk_contraint_states[sent_beam_indices[seq_idx].item()] + topk_state.update(full_hypotheses[seq_idx]) + + if constraints_completed[seq_idx]: + continue + advance_state = advance_constraint_states[seq_idx] + advance_state.update(pre_seq) + if not advance_state.completed: + advance_tokens = advance_state.advance() print(">>>>>advance_tokens", advance_tokens) if advance_tokens.numel() != 0: - advance_token = advance_tokens[:1] - new_seq = torch.cat((pre_seq, advance_token), dim=0) - - new_score = this_batch_token_scores[seq_idx].take(advance_token[0]) - - new_indices.append(seq_idx) - new_seqs.append(new_seq) - new_scores.append(new_score) - + new_state = advance_state.copy(stateful=True) + new_state.add(advance_tokens[:1]) + + for advance_token in advance_tokens: + advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).cpu().tolist() + if advance_seq not in new_seqs: + print("advance_seq", advance_seq) + print("new_seqs", new_seqs) + new_seqs.append(advance_seq) + new_score = this_batch_token_scores[seq_idx].take(advance_token) + new_indices.append(seq_idx) + new_tokens.append(advance_token) + new_scores.append(new_score) + new_states.append(new_state) + # else: + # # if it's a complete state + # new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) + # new_indices.append(seq_idx) + # new_tokens.append(new_token) + # new_scores.append(new_score) + # new_states.append(advance_state) + if len(new_indices) > 0: - new_indices = torch.add(torch.tensor(new_indices), batch_idx * orig_len).to(device) - new_seqs = torch.stack(new_seqs).to(device) + new_indices = torch.tensor(new_indices).to(device) + new_tokens = torch.stack(new_tokens).to(device) new_scores = torch.stack(new_scores).to(device) - # just force advancing - new_tokens = new_seqs[:, -1] - - sent_beam_tokens = torch.cat((sent_beam_tokens, new_tokens[:1]), -1) - sent_beam_scores = torch.cat((sent_beam_scores, new_scores[:1]), -1) - sent_beam_indices = torch.cat((sent_beam_indices, new_indices[:1]), -1) - - print(">>>>>sent_beam_scores", sent_beam_scores) - print(">>>>>sent_beam_tokens", sent_beam_tokens) - print(">>>>>sent_beam_indices", sent_beam_indices) - - - ''' - 2. Compute "banks" for each candidate. - If C is the number of constraints, we construct C banks, where ith bank - is the bank for candidates that have fulfilled ith constraint. - ''' - - - sent_beam_scores = sent_beam_scores[-orig_len:] - sent_beam_tokens = sent_beam_tokens[-orig_len:] - sent_beam_indices = sent_beam_indices[-orig_len:] - - - return sent_beam_scores, sent_beam_tokens, sent_beam_indices + all_states = topk_contraint_states + new_states + all_tokens = torch.cat((sent_beam_tokens, new_tokens), -1) + all_scores = torch.cat((sent_beam_scores, new_scores), -1) + all_banks = torch.tensor([one.get_bank() for one in all_states]).to(device) + # ! only for testing!! + zipped = torch.stack((all_scores, all_banks, all_tokens)) + zipped = torch.transpose(zipped, 0, 1) + + augmented_zipped = all_banks * 1000 + all_scores*10 + indices = augmented_zipped.sort(descending=True).indices + + sorted_banks = all_banks[indices] # C, C, C-1, C-2, ..., 1, 0, 0 + + ''' + Then we end up with + {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} + ''' + counter = -1 + cur_bank = sorted_banks[0] + increments = [] + for bank in sorted_banks: + if bank == cur_bank: + counter += 1 + else: + counter = 0 + cur_bank = bank + increments.append(counter) + rearrangers = torch.tensor(np.argsort(increments, kind="mergesort")) + + print("PREZIP", zipped) + indices = indices[rearrangers] + print("POSTZIP", zipped[indices]) + + sent_beam_scores = all_scores[indices] + sent_beam_tokens = all_tokens[indices] + sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices] + constraints_completed = [all_states[idx].completed for idx in indices[:orig_len]] + + sent_beam_scores = sent_beam_scores[:orig_len] + sent_beam_tokens = sent_beam_tokens[:orig_len] + sent_beam_indices = sent_beam_indices[:orig_len] + + print(">>>>>sent_beam_scores", sent_beam_scores) + print(">>>>>sent_beam_tokens", sent_beam_tokens) + print(">>>>>sent_beam_indices", sent_beam_indices) + print(">>>>constraints_completed", constraints_completed) + + return constraints_completed, sent_beam_scores, sent_beam_tokens, sent_beam_indices diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index c1ebbf41b177..1eb1bbaad361 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -2735,6 +2735,8 @@ def constrained_beam_search( synced_gpus: Optional[bool] = None, **model_kwargs, ) -> Union[BeamSearchOutput, torch.LongTensor]: + from transformers import GPT2Tokenizer + tokenizer = GPT2Tokenizer.from_pretrained("gpt2") r""" Generates sequences for models with a language modeling head using beam search decoding. @@ -2963,7 +2965,10 @@ def constrained_beam_search( next_indices = (next_tokens / vocab_size).long() next_tokens = next_tokens % vocab_size - print("<<<<<>>>>beam_next_tokens", beam_next_tokens) input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - print(">>>>>input_ids", input_ids) - + + print("\n\nOUTPUT\n") + for one in input_ids: + print(tokenizer.decode(one)) + print("\n\n\n") model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder diff --git a/tests/test.py b/tests/test.py index 288da94dfbb7..77cd4edd1931 100644 --- a/tests/test.py +++ b/tests/test.py @@ -9,17 +9,19 @@ model = GPT2LMHeadModel.from_pretrained("gpt2").to(device) tokenizer = GPT2Tokenizer.from_pretrained("gpt2") -force_text = "talk" -force_text_2 = "forceful manner" +force_text = " big monsters" +force_text_2 = " crazy" force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device)[0] force_tokens_2 = tokenizer.encode(force_text_2, return_tensors="pt").to(device)[0] +print("force_tokens", force_tokens) +print("force_tokens_2", force_tokens_2) constraints = [ PhrasalConstraint(force_tokens), PhrasalConstraint(force_tokens_2) ] -input_text = ["He always"] * 2 +input_text = ["The baby is crying because"] * 1 model_inputs = tokenizer(input_text, return_tensors="pt") @@ -30,11 +32,12 @@ k = model.generate( **model_inputs, constraints=constraints, - num_beams=4, - num_return_sequences=3 + num_beams=5, + num_return_sequences=5 ) for out in k: + print("!!", out) print(tokenizer.decode(out)) assert False From 8a0d87156e8de12f6de14b130daca7f7bb499f3e Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 11:04:50 +0000 Subject: [PATCH 05/34] complete PR #1 without disjunctive decoding --- ai.k | 910 ------------------ .../generation_beam_constraints.py | 248 +++-- src/transformers/generation_beam_search.py | 233 ++--- src/transformers/generation_utils.py | 9 - tests/test.py | 43 - 5 files changed, 266 insertions(+), 1177 deletions(-) delete mode 100644 ai.k delete mode 100644 tests/test.py diff --git a/ai.k b/ai.k deleted file mode 100644 index a831bd7f7f96..000000000000 --- a/ai.k +++ /dev/null @@ -1,910 +0,0 @@ -============================= test session starts ============================== -platform linux -- Python 3.7.9, pytest-6.1.1, py-1.9.0, pluggy-0.13.1 -rootdir: /workspace/chanwookim/transformers, configfile: setup.cfg -plugins: hydra-core-1.0.6, flaky-3.7.0 -collected 0 items / 1 error - -==================================== ERRORS ==================================== -________________________ ERROR collecting tests/test.py ________________________ -tests/test.py:43: in - assert False -E assert False -------------------------------- Captured stdout -------------------------------- -force_tokens tensor([ 1263, 10963], device='cuda:0') -force_tokens_2 tensor([7165], device='cuda:0') -model_inputs {'input_ids': tensor([[ 464, 5156, 318, 13774, 780]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1]], device='cuda:0')} - - -INPUT - -The baby is crying because -The baby is crying because -The baby is crying because -The baby is crying because -The baby is crying because - - - - ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 1263] -new_seqs [] -advance_seq [464, 5156, 318, 13774, 780, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 1263]] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -PREZIP tensor([[-1.4358e+00, 0.0000e+00, 6.7300e+02], - [-2.0147e+00, 0.0000e+00, 3.3900e+02], - [-2.3535e+00, 0.0000e+00, 3.4000e+02], - [-2.4315e+00, 0.0000e+00, 2.8600e+02], - [-2.7982e+00, 0.0000e+00, 6.0700e+02], - [-1.0072e+01, 1.0000e+00, 1.2630e+03], - [-1.2223e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') -POSTZIP tensor([[-1.0072e+01, 1.0000e+00, 1.2630e+03], - [-1.4358e+00, 0.0000e+00, 6.7300e+02], - [-1.2223e+01, 1.0000e+00, 7.1650e+03], - [-2.0147e+00, 0.0000e+00, 3.3900e+02], - [-2.3535e+00, 0.0000e+00, 3.4000e+02], - [-2.4315e+00, 0.0000e+00, 2.8600e+02], - [-2.7982e+00, 0.0000e+00, 6.0700e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([-10.0721, -1.4358, -12.2228, -2.0147, -2.3535], device='cuda:0') ->>>>>sent_beam_tokens tensor([1263, 673, 7165, 339, 340], device='cuda:0') ->>>>>sent_beam_indices tensor([0, 0, 0, 0, 0], device='cuda:0') ->>>>constraints_completed [False, False, False, False, False] - - -OUTPUT - -The baby is crying because big -The baby is crying because she -The baby is crying because crazy -The baby is crying because he -The baby is crying because it - - - - - - -INPUT - -The baby is crying because big -The baby is crying because she -The baby is crying because crazy -The baby is crying because he -The baby is crying because it - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963]] -advance_seq [464, 5156, 318, 13774, 780, 673, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 7165, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165]] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 339, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263]] -advance_seq [464, 5156, 318, 13774, 780, 339, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263]] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 340, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263], [464, 5156, 318, 13774, 780, 339, 7165]] -advance_seq [464, 5156, 318, 13774, 780, 340, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 1263], [464, 5156, 318, 13774, 780, 673, 7165], [464, 5156, 318, 13774, 780, 7165, 1263], [464, 5156, 318, 13774, 780, 339, 1263], [464, 5156, 318, 13774, 780, 339, 7165], [464, 5156, 318, 13774, 780, 340, 1263]] -PREZIP tensor([[-3.3543e+00, 0.0000e+00, 3.3800e+02], - [-3.7313e+00, 0.0000e+00, 3.3800e+02], - [-3.7631e+00, 0.0000e+00, 3.7300e+02], - [-3.8234e+00, 0.0000e+00, 3.3800e+02], - [-3.8247e+00, 0.0000e+00, 3.1800e+02], - [-6.6173e+00, 2.0000e+00, 1.0963e+04], - [-1.2264e+01, 1.0000e+00, 1.2630e+03], - [-1.3972e+01, 1.0000e+00, 7.1650e+03], - [-5.7579e+00, 3.0000e+00, 1.2630e+03], - [-1.1960e+01, 1.0000e+00, 1.2630e+03], - [-1.3683e+01, 1.0000e+00, 7.1650e+03], - [-1.2435e+01, 1.0000e+00, 1.2630e+03], - [-1.3548e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') -POSTZIP tensor([[-5.7579e+00, 3.0000e+00, 1.2630e+03], - [-6.6173e+00, 2.0000e+00, 1.0963e+04], - [-1.1960e+01, 1.0000e+00, 1.2630e+03], - [-3.3543e+00, 0.0000e+00, 3.3800e+02], - [-1.2264e+01, 1.0000e+00, 1.2630e+03], - [-3.7313e+00, 0.0000e+00, 3.3800e+02], - [-1.2435e+01, 1.0000e+00, 1.2630e+03], - [-3.7631e+00, 0.0000e+00, 3.7300e+02], - [-1.3548e+01, 1.0000e+00, 7.1650e+03], - [-3.8234e+00, 0.0000e+00, 3.3800e+02], - [-1.3683e+01, 1.0000e+00, 7.1650e+03], - [-3.8247e+00, 0.0000e+00, 3.1800e+02], - [-1.3972e+01, 1.0000e+00, 7.1650e+03]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -5.7579, -6.6173, -11.9598, -3.3543, -12.2640], device='cuda:0') ->>>>>sent_beam_tokens tensor([ 1263, 10963, 1263, 338, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([2, 0, 3, 1, 1], device='cuda:0') ->>>>constraints_completed [False, False, False, False, False] - - -OUTPUT - -The baby is crying because crazy big -The baby is crying because big monsters -The baby is crying because he big -The baby is crying because she's -The baby is crying because she big - - - - - - -INPUT - -The baby is crying because crazy big -The baby is crying because big monsters -The baby is crying because he big -The baby is crying because she's -The baby is crying because she big - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 7165, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 1263, 10963, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963]] ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 339, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165]] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963]] -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 1263]] ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 7165, 1263, 10963], [464, 5156, 318, 13774, 780, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 339, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 1263], [464, 5156, 318, 13774, 780, 673, 338, 7165]] -PREZIP tensor([[-5.8691e+00, 0.0000e+00, 5.8700e+02], - [-6.3662e+00, 0.0000e+00, 4.0700e+02], - [-6.4227e+00, 0.0000e+00, 1.2008e+04], - [-6.4356e+00, 0.0000e+00, 7.7870e+03], - [-6.7641e+00, 0.0000e+00, 5.2300e+02], - [-5.7414e+00, 4.0000e+00, 1.0963e+04], - [-1.2468e+01, 4.0000e+00, 7.1650e+03], - [-1.0331e+01, 2.0000e+00, 1.0963e+04], - [-8.2352e+00, 1.0000e+00, 1.2630e+03], - [-7.5897e+00, 1.0000e+00, 7.1650e+03], - [-1.0610e+01, 2.0000e+00, 1.0963e+04]], device='cuda:0') -POSTZIP tensor([[-5.7414e+00, 4.0000e+00, 1.0963e+04], - [-1.0331e+01, 2.0000e+00, 1.0963e+04], - [-7.5897e+00, 1.0000e+00, 7.1650e+03], - [-5.8691e+00, 0.0000e+00, 5.8700e+02], - [-1.2468e+01, 4.0000e+00, 7.1650e+03], - [-1.0610e+01, 2.0000e+00, 1.0963e+04], - [-8.2352e+00, 1.0000e+00, 1.2630e+03], - [-6.3662e+00, 0.0000e+00, 4.0700e+02], - [-6.4227e+00, 0.0000e+00, 1.2008e+04], - [-6.4356e+00, 0.0000e+00, 7.7870e+03], - [-6.7641e+00, 0.0000e+00, 5.2300e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -5.7414, -10.3309, -7.5897, -5.8691, -12.4677], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 10963, 7165, 587, 7165], device='cuda:0') ->>>>>sent_beam_indices tensor([0, 2, 3, 3, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because crazy big monsters -The baby is crying because he big monsters -The baby is crying because she's crazy -The baby is crying because she's been -The baby is crying because big monsters crazy - - - - - - -INPUT - -The baby is crying because crazy big monsters -The baby is crying because he big monsters -The baby is crying because she's crazy -The baby is crying because she's been -The baby is crying because big monsters crazy - - - - ->>>>>advance_tokens tensor([7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165]] ->>>>>advance_tokens tensor([1263, 7165], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263]] -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165] -new_seqs [[464, 5156, 318, 13774, 780, 339, 1263, 10963, 7165], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263], [464, 5156, 318, 13774, 780, 673, 338, 587, 1263]] -PREZIP tensor([[-6.8971e+00, 2.0000e+00, 3.8900e+02], - [-8.0195e+00, 0.0000e+00, 4.2300e+02], - [-8.5441e+00, 0.0000e+00, 1.2970e+03], - [-8.9806e+00, 2.0000e+00, 1.3000e+01], - [-9.2211e+00, 0.0000e+00, 1.6110e+04], - [-1.1526e+01, 4.0000e+00, 7.1650e+03], - [-7.5428e+00, 3.0000e+00, 1.2630e+03], - [-1.1152e+01, 1.0000e+00, 1.2630e+03], - [-9.3288e+00, 1.0000e+00, 7.1650e+03]], device='cuda:0') -POSTZIP tensor([[-1.1526e+01, 4.0000e+00, 7.1650e+03], - [-7.5428e+00, 3.0000e+00, 1.2630e+03], - [-6.8971e+00, 2.0000e+00, 3.8900e+02], - [-9.3288e+00, 1.0000e+00, 7.1650e+03], - [-8.0195e+00, 0.0000e+00, 4.2300e+02], - [-8.9806e+00, 2.0000e+00, 1.3000e+01], - [-1.1152e+01, 1.0000e+00, 1.2630e+03], - [-8.5441e+00, 0.0000e+00, 1.2970e+03], - [-9.2211e+00, 0.0000e+00, 1.6110e+04]], device='cuda:0') ->>>>>sent_beam_scores tensor([-11.5259, -7.5428, -6.8971, -9.3288, -8.0195], device='cuda:0') ->>>>>sent_beam_tokens tensor([7165, 1263, 389, 7165, 423], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 2, 0, 3, 0], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because he big monsters crazy -The baby is crying because she's crazy big -The baby is crying because crazy big monsters are -The baby is crying because she's been crazy -The baby is crying because crazy big monsters have - - - - - - -INPUT - -The baby is crying because he big monsters crazy -The baby is crying because she's crazy big -The baby is crying because crazy big monsters are -The baby is crying because she's been crazy -The baby is crying because crazy big monsters have - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 10963]] -PREZIP tensor([[-8.8730e+00, 0.0000e+00, 2.4060e+03], - [-8.9872e+00, 0.0000e+00, 2.9000e+02], - [-9.2539e+00, 0.0000e+00, 1.3000e+01], - [-9.5069e+00, 0.0000e+00, 1.1000e+01], - [-9.6623e+00, 0.0000e+00, 5.5300e+02], - [-1.2619e+01, 4.0000e+00, 1.0963e+04], - [-8.8750e+00, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-1.2619e+01, 4.0000e+00, 1.0963e+04], - [-8.8750e+00, 3.0000e+00, 1.2630e+03], - [-8.8730e+00, 0.0000e+00, 2.4060e+03], - [-8.9872e+00, 0.0000e+00, 2.9000e+02], - [-9.2539e+00, 0.0000e+00, 1.3000e+01], - [-9.5069e+00, 0.0000e+00, 1.1000e+01], - [-9.6623e+00, 0.0000e+00, 5.5300e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([-12.6191, -8.8750, -8.8730, -8.9872, -9.2539], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 2406, 290, 13], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 3, 2, 1, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, True, False] - - -OUTPUT - -The baby is crying because she's crazy big monsters -The baby is crying because she's been crazy big -The baby is crying because crazy big monsters are coming -The baby is crying because she's crazy big and -The baby is crying because she's crazy big. - - - - - - -INPUT - -The baby is crying because she's crazy big monsters -The baby is crying because she's been crazy big -The baby is crying because crazy big monsters are coming -The baby is crying because she's crazy big and -The baby is crying because she's crazy big. - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 10963]] -PREZIP tensor([[-1.0600e+01, 0.0000e+00, 6.7300e+02], - [-1.0826e+01, 2.0000e+00, 2.9000e+02], - [-1.0833e+01, 2.0000e+00, 1.3750e+03], - [-1.0837e+01, 0.0000e+00, 1.3000e+01], - [-1.1092e+01, 4.0000e+00, 1.3000e+01], - [-1.2403e+01, 4.0000e+00, 1.0963e+04], - [-1.5209e+01, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-1.1092e+01, 4.0000e+00, 1.3000e+01], - [-1.5209e+01, 3.0000e+00, 1.2630e+03], - [-1.0826e+01, 2.0000e+00, 2.9000e+02], - [-1.0600e+01, 0.0000e+00, 6.7300e+02], - [-1.2403e+01, 4.0000e+00, 1.0963e+04], - [-1.0833e+01, 2.0000e+00, 1.3750e+03], - [-1.0837e+01, 0.0000e+00, 1.3000e+01]], device='cuda:0') ->>>>>sent_beam_scores tensor([-11.0919, -15.2090, -10.8264, -10.5999, -12.4034], device='cuda:0') ->>>>>sent_beam_tokens tensor([ 13, 1263, 290, 673, 10963], device='cuda:0') ->>>>>sent_beam_indices tensor([2, 4, 1, 3, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because crazy big monsters are coming. -The baby is crying because she's crazy big. big -The baby is crying because she's been crazy big and -The baby is crying because she's crazy big and she -The baby is crying because she's been crazy big monsters - - - - - - -INPUT - -The baby is crying because crazy big monsters are coming. -The baby is crying because she's crazy big. big -The baby is crying because she's been crazy big and -The baby is crying because she's crazy big and she -The baby is crying because she's been crazy big monsters - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 13, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263]] -PREZIP tensor([[-1.1419e+01, 2.0000e+00, 3.3800e+02], - [-1.2924e+01, 0.0000e+00, 6.7300e+02], - [-1.3312e+01, 2.0000e+00, 1.9800e+02], - [-1.3459e+01, 2.0000e+00, 3.3820e+03], - [-1.3549e+01, 0.0000e+00, 4.6000e+02], - [-9.9536e+00, 4.0000e+00, 1.0963e+04], - [-3.8742e+00, 3.0000e+00, 1.2630e+03], - [-1.0987e+01, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-9.9536e+00, 4.0000e+00, 1.0963e+04], - [-3.8742e+00, 3.0000e+00, 1.2630e+03], - [-1.1419e+01, 2.0000e+00, 3.3800e+02], - [-1.2924e+01, 0.0000e+00, 6.7300e+02], - [-1.0987e+01, 3.0000e+00, 1.2630e+03], - [-1.3312e+01, 2.0000e+00, 1.9800e+02], - [-1.3549e+01, 0.0000e+00, 4.6000e+02], - [-1.3459e+01, 2.0000e+00, 3.3820e+03]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -9.9536, -3.8742, -11.4194, -12.9239, -10.9866], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 338, 673, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 2, 3, 2, 3], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because she's crazy big. big monsters -The baby is crying because she's been crazy big and big -The baby is crying because she's crazy big and she's -The baby is crying because she's been crazy big and she -The baby is crying because she's crazy big and she big - - - - - - -INPUT - -The baby is crying because she's crazy big. big monsters -The baby is crying because she's been crazy big and big -The baby is crying because she's crazy big and she's -The baby is crying because she's been crazy big and she -The baby is crying because she's crazy big and she big - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 673, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263]] ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 673, 1263]] -PREZIP tensor([[-5.2463e+00, 0.0000e+00, 2.9000e+02], - [-5.9868e+00, 2.0000e+00, 1.3000e+01], - [-6.1505e+00, 0.0000e+00, 1.1000e+01], - [-6.1889e+00, 0.0000e+00, 5.5300e+02], - [-6.7390e+00, 0.0000e+00, 3.2900e+02], - [-1.1919e+01, 4.0000e+00, 1.0963e+04], - [-5.6936e+00, 3.0000e+00, 1.2630e+03], - [-1.1110e+01, 3.0000e+00, 1.2630e+03], - [-9.3693e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') -POSTZIP tensor([[-9.3693e+00, 4.0000e+00, 1.0963e+04], - [-5.6936e+00, 3.0000e+00, 1.2630e+03], - [-5.9868e+00, 2.0000e+00, 1.3000e+01], - [-5.2463e+00, 0.0000e+00, 2.9000e+02], - [-1.1919e+01, 4.0000e+00, 1.0963e+04], - [-1.1110e+01, 3.0000e+00, 1.2630e+03], - [-6.1505e+00, 0.0000e+00, 1.1000e+01], - [-6.1889e+00, 0.0000e+00, 5.5300e+02], - [-6.7390e+00, 0.0000e+00, 3.2900e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -9.3693, -5.6936, -5.9868, -5.2463, -11.9188], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') ->>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because she's crazy big and she big monsters -The baby is crying because she's crazy big and she's big -The baby is crying because she's been crazy big and big. -The baby is crying because she's been crazy big and big and -The baby is crying because she's been crazy big and big monsters - - - - - - -INPUT - -The baby is crying because she's crazy big and she big monsters -The baby is crying because she's crazy big and she's big -The baby is crying because she's been crazy big and big. -The baby is crying because she's been crazy big and big and -The baby is crying because she's been crazy big and big monsters - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263]] -PREZIP tensor([[-6.3694e+00, 0.0000e+00, 1.2630e+03], - [-6.8644e+00, 2.0000e+00, 2.9000e+02], - [-7.4970e+00, 2.0000e+00, 1.3750e+03], - [-7.5834e+00, 2.0000e+00, 6.7300e+02], - [-7.8789e+00, 0.0000e+00, 1.5760e+03], - [-1.2507e+01, 4.0000e+00, 1.0963e+04], - [-1.6012e+01, 3.0000e+00, 1.2630e+03], - [-1.1232e+00, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-1.2507e+01, 4.0000e+00, 1.0963e+04], - [-1.1232e+00, 3.0000e+00, 1.2630e+03], - [-6.8644e+00, 2.0000e+00, 2.9000e+02], - [-6.3694e+00, 0.0000e+00, 1.2630e+03], - [-1.6012e+01, 3.0000e+00, 1.2630e+03], - [-7.4970e+00, 2.0000e+00, 1.3750e+03], - [-7.8789e+00, 0.0000e+00, 1.5760e+03], - [-7.5834e+00, 2.0000e+00, 6.7300e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([-12.5067, -1.1232, -6.8644, -6.3694, -16.0120], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 290, 1263, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 3, 1, 3, 2], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because she's crazy big and she's big monsters -The baby is crying because she's been crazy big and big and big -The baby is crying because she's crazy big and she's big and -The baby is crying because she's been crazy big and big and big -The baby is crying because she's been crazy big and big. big - - - - - - -INPUT - -The baby is crying because she's crazy big and she's big monsters -The baby is crying because she's been crazy big and big and big -The baby is crying because she's crazy big and she's big and -The baby is crying because she's been crazy big and big and big -The baby is crying because she's been crazy big and big. big - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963]] ->>>>>advance_tokens tensor([10963], device='cuda:0') ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 13, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263]] -PREZIP tensor([[-2.4862e+00, 0.0000e+00, 2.9000e+02], - [-2.8610e+00, 2.0000e+00, 1.3000e+01], - [-3.3225e+00, 0.0000e+00, 5.5300e+02], - [-3.4911e+00, 0.0000e+00, 1.1000e+01], - [-3.9857e+00, 0.0000e+00, 3.2900e+02], - [-1.2293e+01, 4.0000e+00, 1.0963e+04], - [-4.9025e+00, 3.0000e+00, 1.2630e+03], - [-9.6560e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') -POSTZIP tensor([[-9.6560e+00, 4.0000e+00, 1.0963e+04], - [-4.9025e+00, 3.0000e+00, 1.2630e+03], - [-2.8610e+00, 2.0000e+00, 1.3000e+01], - [-2.4862e+00, 0.0000e+00, 2.9000e+02], - [-1.2293e+01, 4.0000e+00, 1.0963e+04], - [-3.3225e+00, 0.0000e+00, 5.5300e+02], - [-3.4911e+00, 0.0000e+00, 1.1000e+01], - [-3.9857e+00, 0.0000e+00, 3.2900e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -9.6560, -4.9025, -2.8610, -2.4862, -12.2932], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') ->>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because she's been crazy big and big. big monsters -The baby is crying because she's crazy big and she's big and big -The baby is crying because she's been crazy big and big and big. -The baby is crying because she's been crazy big and big and big and -The baby is crying because she's been crazy big and big and big monsters - - - - - - -INPUT - -The baby is crying because she's been crazy big and big. big monsters -The baby is crying because she's crazy big and she's big and big -The baby is crying because she's been crazy big and big and big. -The baby is crying because she's been crazy big and big and big and -The baby is crying because she's been crazy big and big and big monsters - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 7165, 1263, 290, 673, 338, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263]] -PREZIP tensor([[-3.3472e+00, 0.0000e+00, 1.2630e+03], - [-4.4136e+00, 0.0000e+00, 1.3750e+03], - [-5.0307e+00, 2.0000e+00, 8.4300e+02], - [-5.0600e+00, 0.0000e+00, 3.1400e+02], - [-5.1705e+00, 0.0000e+00, 1.9800e+02], - [-1.2517e+01, 4.0000e+00, 1.0963e+04], - [-1.6495e+01, 3.0000e+00, 1.2630e+03], - [-8.6100e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-1.2517e+01, 4.0000e+00, 1.0963e+04], - [-8.6100e-01, 3.0000e+00, 1.2630e+03], - [-5.0307e+00, 2.0000e+00, 8.4300e+02], - [-3.3472e+00, 0.0000e+00, 1.2630e+03], - [-1.6495e+01, 3.0000e+00, 1.2630e+03], - [-4.4136e+00, 0.0000e+00, 1.3750e+03], - [-5.0600e+00, 0.0000e+00, 3.1400e+02], - [-5.1705e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([-12.5170, -0.8610, -5.0307, -3.3472, -16.4952], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because she's crazy big and she's big and big monsters -The baby is crying because she's been crazy big and big and big and big -The baby is crying because she's been crazy big and big and big. And -The baby is crying because she's been crazy big and big and big and big -The baby is crying because she's been crazy big and big and big. big - - - - - - -INPUT - -The baby is crying because she's crazy big and she's big and big monsters -The baby is crying because she's been crazy big and big and big and big -The baby is crying because she's been crazy big and big and big. And -The baby is crying because she's been crazy big and big and big and big -The baby is crying because she's been crazy big and big and big. big - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963]] ->>>>>advance_tokens tensor([10963], device='cuda:0') ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263]] -PREZIP tensor([[-2.1940e+00, 0.0000e+00, 2.9000e+02], - [-2.4520e+00, 2.0000e+00, 1.3000e+01], - [-3.2742e+00, 0.0000e+00, 5.5300e+02], - [-3.3400e+00, 0.0000e+00, 1.1000e+01], - [-3.6658e+00, 0.0000e+00, 5.2600e+02], - [-1.2107e+01, 4.0000e+00, 1.0963e+04], - [-8.4168e+00, 3.0000e+00, 1.2630e+03], - [-9.5848e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') -POSTZIP tensor([[-9.5848e+00, 4.0000e+00, 1.0963e+04], - [-8.4168e+00, 3.0000e+00, 1.2630e+03], - [-2.4520e+00, 2.0000e+00, 1.3000e+01], - [-2.1940e+00, 0.0000e+00, 2.9000e+02], - [-1.2107e+01, 4.0000e+00, 1.0963e+04], - [-3.2742e+00, 0.0000e+00, 5.5300e+02], - [-3.3400e+00, 0.0000e+00, 1.1000e+01], - [-3.6658e+00, 0.0000e+00, 5.2600e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -9.5848, -8.4168, -2.4520, -2.1940, -12.1066], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') ->>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because she's been crazy big and big and big. big monsters -The baby is crying because she's been crazy big and big and big. And big -The baby is crying because she's been crazy big and big and big and big. -The baby is crying because she's been crazy big and big and big and big and -The baby is crying because she's been crazy big and big and big and big monsters - - - - - - -INPUT - -The baby is crying because she's been crazy big and big and big. big monsters -The baby is crying because she's been crazy big and big and big. And big -The baby is crying because she's been crazy big and big and big and big. -The baby is crying because she's been crazy big and big and big and big and -The baby is crying because she's been crazy big and big and big and big monsters - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263]] -PREZIP tensor([[-2.8488e+00, 0.0000e+00, 1.2630e+03], - [-4.0466e+00, 0.0000e+00, 1.3750e+03], - [-4.4985e+00, 2.0000e+00, 8.4300e+02], - [-4.6394e+00, 0.0000e+00, 3.1400e+02], - [-4.7340e+00, 0.0000e+00, 1.9800e+02], - [-8.2042e+00, 4.0000e+00, 1.0963e+04], - [-1.6431e+01, 3.0000e+00, 1.2630e+03], - [-6.5482e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-8.2042e+00, 4.0000e+00, 1.0963e+04], - [-6.5482e-01, 3.0000e+00, 1.2630e+03], - [-4.4985e+00, 2.0000e+00, 8.4300e+02], - [-2.8488e+00, 0.0000e+00, 1.2630e+03], - [-1.6431e+01, 3.0000e+00, 1.2630e+03], - [-4.0466e+00, 0.0000e+00, 1.3750e+03], - [-4.6394e+00, 0.0000e+00, 3.1400e+02], - [-4.7340e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -8.2042, -0.6548, -4.4985, -2.8488, -16.4311], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because she's been crazy big and big and big. And big monsters -The baby is crying because she's been crazy big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big. And -The baby is crying because she's been crazy big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big. big - - - - - - -INPUT - -The baby is crying because she's been crazy big and big and big. And big monsters -The baby is crying because she's been crazy big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big. And -The baby is crying because she's been crazy big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big. big - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963]] ->>>>>advance_tokens tensor([10963], device='cuda:0') ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263, 10963] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263]] -PREZIP tensor([[-1.7995e+00, 0.0000e+00, 2.9000e+02], - [-2.2400e+00, 2.0000e+00, 1.3000e+01], - [-3.2673e+00, 0.0000e+00, 1.1000e+01], - [-3.4144e+00, 0.0000e+00, 5.5300e+02], - [-3.4721e+00, 0.0000e+00, 5.2600e+02], - [-1.1857e+01, 4.0000e+00, 1.0963e+04], - [-8.4926e+00, 3.0000e+00, 1.2630e+03], - [-9.7216e+00, 4.0000e+00, 1.0963e+04]], device='cuda:0') -POSTZIP tensor([[-9.7216e+00, 4.0000e+00, 1.0963e+04], - [-8.4926e+00, 3.0000e+00, 1.2630e+03], - [-2.2400e+00, 2.0000e+00, 1.3000e+01], - [-1.7995e+00, 0.0000e+00, 2.9000e+02], - [-1.1857e+01, 4.0000e+00, 1.0963e+04], - [-3.2673e+00, 0.0000e+00, 1.1000e+01], - [-3.4144e+00, 0.0000e+00, 5.5300e+02], - [-3.4721e+00, 0.0000e+00, 5.2600e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -9.7216, -8.4926, -2.2400, -1.7995, -11.8567], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 13, 290, 10963], device='cuda:0') ->>>>>sent_beam_indices tensor([4, 2, 1, 1, 1], device='cuda:0') ->>>>constraints_completed [True, False, False, False, True] - - -OUTPUT - -The baby is crying because she's been crazy big and big and big and big. big monsters -The baby is crying because she's been crazy big and big and big and big. And big -The baby is crying because she's been crazy big and big and big and big and big. -The baby is crying because she's been crazy big and big and big and big and big and -The baby is crying because she's been crazy big and big and big and big and big monsters - - - - - - -INPUT - -The baby is crying because she's been crazy big and big and big and big. big monsters -The baby is crying because she's been crazy big and big and big and big. And big -The baby is crying because she's been crazy big and big and big and big and big. -The baby is crying because she's been crazy big and big and big and big and big and -The baby is crying because she's been crazy big and big and big and big and big monsters - - - - ->>>>>advance_tokens tensor([10963], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963] -new_seqs [] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963]] ->>>>>advance_tokens tensor([1263], device='cuda:0') -advance_seq [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263] -new_seqs [[464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], [464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263]] -PREZIP tensor([[-2.2907e+00, 0.0000e+00, 1.2630e+03], - [-3.8817e+00, 0.0000e+00, 1.3750e+03], - [-4.2239e+00, 2.0000e+00, 8.4300e+02], - [-4.4378e+00, 0.0000e+00, 3.1400e+02], - [-4.4513e+00, 0.0000e+00, 1.9800e+02], - [-8.3479e+00, 4.0000e+00, 1.0963e+04], - [-1.6134e+01, 3.0000e+00, 1.2630e+03], - [-4.9122e-01, 3.0000e+00, 1.2630e+03]], device='cuda:0') -POSTZIP tensor([[-8.3479e+00, 4.0000e+00, 1.0963e+04], - [-4.9122e-01, 3.0000e+00, 1.2630e+03], - [-4.2239e+00, 2.0000e+00, 8.4300e+02], - [-2.2907e+00, 0.0000e+00, 1.2630e+03], - [-1.6134e+01, 3.0000e+00, 1.2630e+03], - [-3.8817e+00, 0.0000e+00, 1.3750e+03], - [-4.4378e+00, 0.0000e+00, 3.1400e+02], - [-4.4513e+00, 0.0000e+00, 1.9800e+02]], device='cuda:0') ->>>>>sent_beam_scores tensor([ -8.3479, -0.4912, -4.2239, -2.2907, -16.1338], device='cuda:0') ->>>>>sent_beam_tokens tensor([10963, 1263, 843, 1263, 1263], device='cuda:0') ->>>>>sent_beam_indices tensor([1, 3, 2, 3, 2], device='cuda:0') ->>>>constraints_completed [True, False, False, False, False] - - -OUTPUT - -The baby is crying because she's been crazy big and big and big and big. And big monsters -The baby is crying because she's been crazy big and big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big and big. And -The baby is crying because she's been crazy big and big and big and big and big and big -The baby is crying because she's been crazy big and big and big and big and big. big - - - - -!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, - 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263], - device='cuda:0') -The baby is crying because she's been crazy big and big and big and big and big and big -!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, - 290, 1263, 290, 1263, 290, 1263, 290, 1263, 290, 1263], - device='cuda:0') -The baby is crying because she's been crazy big and big and big and big and big and big -!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, - 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 843], - device='cuda:0') -The baby is crying because she's been crazy big and big and big and big and big. And -!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, - 290, 1263, 290, 1263, 290, 1263, 13, 843, 1263, 10963], - device='cuda:0') -The baby is crying because she's been crazy big and big and big and big. And big monsters -!! tensor([ 464, 5156, 318, 13774, 780, 673, 338, 587, 7165, 1263, - 290, 1263, 290, 1263, 290, 1263, 290, 1263, 13, 1263], - device='cuda:0') -The baby is crying because she's been crazy big and big and big and big and big. big -------------------------------- Captured stderr -------------------------------- -2022-01-20 12:25:52.703100: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory -2022-01-20 12:25:52.703133: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine. -Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation. -=============================== warnings summary =============================== -../../../opt/conda/envs/ACW/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py:22 - /opt/conda/envs/ACW/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py:22: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses - import imp - --- Docs: https://docs.pytest.org/en/stable/warnings.html -=========================== short test summary info ============================ -ERROR tests/test.py - assert False -!!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!! -========================= 1 warning, 1 error in 22.90s ========================= diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index def4e3783681..2352df1fdf64 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -2,7 +2,7 @@ from itertools import chain from collections import Counter -from typing import List, Optional, Set, Tuple +from typing import List, Optional, Set, Tuple, Union import torch @@ -13,7 +13,6 @@ logger = get_logger(__name__) - class Constraint(ABC): r"""Abstract base class for all constraints that can be applied during generation. It must define how the constraint can be satisfied. @@ -27,25 +26,37 @@ class Constraint(ABC): ``` will always terminate (halt). - """ def __init__(self): # test for the above condition + self.test() + + def test(self): + ''' + Tests whether this constraint has been properly defined. + ''' counter = 0 completed = False while not completed: + if counter == 1: + self.reset() advance = self.advance() - _, completed = self.update(advance) + assert self.does_advance(advance) + stepped, completed, reset = self.update(advance) counter += 1 if counter > 10000: raise Exception("update() does not fulfill the constraint.") - + + assert self.remaining() == 0 def advance(self): ''' When called, returns the token that would take this constraint one step closer to being fulfilled. + + returns: + token_ids(`torch.tensor`): Must be a tensor of a list of indexable tokens, not some integer. ''' raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." @@ -64,9 +75,9 @@ def update(self, token_id: int): Reads in a token and returns booleans that indicate the progress made by it. This function will update the state of this object unlikes `does_advance(self, token_id: int)`. - This function assumes that token_id is sure to be generated. This isn't to test whether - a certain token will advnace the progress; so we update the states accordingly - if it's already been generated. This becomes important if token_id != desired token (refer to else statement in PhrasalConstraint) + This isn't to test whether a certain token will advance the progress; it's to update its state + as if it has been generated. This becomes important if token_id != desired token + (refer to else statement in PhrasalConstraint) Args: token_id(`int`): @@ -74,12 +85,45 @@ def update(self, token_id: int): returns: stepped(`boolean`): Whether this constraint has become one step closer to being fulfuilled. - completed(`stepped`): + completed(`boolean`): Whether this constraint has been completely fulfilled by this token being generated. + reset (`boolean`): + Whether this constraint has reset its progress by this token being generated. """ raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) + + def reset(self): + """ + Resets the state of this constraint to its initialization. + We would call this in cases where the fulfillment of a constraint is abrupted by an unwanted token. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + def remaining(self): + ''' + Returns the number of remaining steps of `advance()` in order to complete this constraint. + ''' + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + def copy(self, stateful=False): + ''' + Creates a new instance of this constraint. + + Args: + stateful(`boolean`): Whether to not only copy the constraint for new instance, but also its state. + Returns: + constraint(`Constraint`): The same constraint as the one being called from. + ''' + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + class TokenConstraint(Constraint): r""" @@ -89,30 +133,45 @@ class TokenConstraint(Constraint): token_id (`int`): The token that must be generated by the output. """ - def __init__(self, token_id: int): + def __init__(self, token_id: Union[int, torch.LongTensor]): super(Constraint, self).__init__() - if not isinstance(token_id, int) or token_id < 0: - raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") + if not (isinstance(token_id, int) or isinstance(token_id, torch.LongTensor)) or token_id < 0: + raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integeter, but is {token_id}") + else: + if isinstance(token_id, torch.LongTensor) and token_id.size(0) > 1: + raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." + "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`.") + self.token_id = token_id - - def copy(self): - return TokenConstraint(self.token_id) + self.completed = False def advance(self): return self.token_id - def does_advance(self, token_id: int): + def does_advance(self, token_id: Union[int, torch.LongTensor]): return token_id == self.token_id - def update(self, token_id: int): + def update(self, token_id: Union[int, torch.LongTensor]): if not isinstance(token_id, int) or token_id < 0: raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") if self.does_advance(token_id): - return True, True, True # stepped, completed, reset + self.completed = True + return True, True, True # stepped, completed, reset else: return False, False, False + def reset(self): + self.completed = False + + def remaining(self): + return 0 if self.completed else 1 + + def copy(self, stateful=False): + constraint = TokenConstraint(self.token_id) + if stateful: + constraint.completed = self.completed + return constraint class PhrasalConstraint(Constraint): r""" @@ -125,26 +184,12 @@ class PhrasalConstraint(Constraint): def __init__(self, token_ids: torch.Tensor): super(Constraint, self).__init__() if not isinstance(token_ids, torch.Tensor): - raise ValueError(f"`token_ids` has to be a tensor, but is {type(token_ids)}") + raise ValueError(f"`token_ids` has to be a `torch.Tensor`, but is {type(token_ids)}") self.token_ids = token_ids - self.seqlen = self.token_ids.size(0) self.fulfilled_idx = -1 # the index of the currently fulfilled step self.completed = False - - def remaining(self): - return self.seqlen - (self.fulfilled_idx + 1) - - def copy(self, stateful=False): - new_constraint = PhrasalConstraint(self.token_ids) - - if stateful: - new_constraint.seq_len = self.seqlen - new_constraint.fulfilled_idx = self.fulfilled_idx - new_constraint.completed = self.completed - - return new_constraint def advance(self): return self.token_ids[self.fulfilled_idx + 1] @@ -160,11 +205,6 @@ def update(self, token_id: int): completed = False reset = False -# ''' -# ------------------------------- Captured stdout -------------------------------- -# force_tokens tensor([48186, 220], device='cuda:0') -# force_tokens_2 tensor([3174, 913, 220], device='cuda:0') -# ''' if self.does_advance(token_id): self.fulfilled_idx += 1 stepped = True @@ -181,59 +221,63 @@ def reset(self): self.completed = False self.fulfilled_idx = 0 - -# force_tokens tensor([3375, 220], device='cuda:0') -# force_tokens_2 tensor([23112], device='cuda:0') + def remaining(self): + return self.seqlen - (self.fulfilled_idx + 1) + + def copy(self, stateful=False): + new_constraint = PhrasalConstraint(self.token_ids) + if stateful: + new_constraint.seq_len = self.seqlen + new_constraint.fulfilled_idx = self.fulfilled_idx + new_constraint.completed = self.completed + + return new_constraint + + +# For beam scorers to track its progress through a list of constraints. class ConstraintListState: def __init__(self, constraints: List[Constraint]): self.constraints = constraints - self.max_seqlen = max([c.seqlen for c in constraints]) + # max # of steps required to fulfill a given constraint + self.max_seqlen = max([c.seqlen for c in constraints]) self.n_constraints = len(constraints) self.completed = False self.init_state() - def copy(self, stateful=True): - new_state = ConstraintListState(self.constraints) - - if stateful: - new_state.complete_constraints = [ - constraint.copy(stateful=True) - for constraint in self.complete_constraints - ] - if self.inprogress_constraint is not None: - new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True) - new_state.pending_constraints = [ - constraint.copy() - for constraint in self.pending_constraints - ] - - return new_state - def init_state(self): self.complete_constraints = [] self.inprogress_constraint = None - self.pending_constraints = [constraint.copy() for constraint in self.constraints] + self.pending_constraints = [constraint.copy(stateful=False) for constraint in self.constraints] def get_bank(self): add = 0 if self.inprogress_constraint: + # extra points for having a constraint mid-fulfilled add += self.max_seqlen - self.inprogress_constraint.remaining() return (len(self.complete_constraints)*self.max_seqlen) + add def advance(self): '''The list of tokens to generate such that we can make progress. + By "list" we don't mean the list of token that will fully fulfill a constraint. + + Given constraints c_i = {t_ij | j == # of tokens}, + If we're not in the middle of progressing through a specific constraint c_i, we return: + + [t_k1 for k in indices of unfulfilled constraints] + + If we are in the middle of a constraint, then we return: + [t_ij], where i == index of the inprogress constraint, j == the next step for the constraint. Though we don't care which constraint is fulfilled first, - if we are in the progress of fulfilling a constraint, that's the only one - we'll return. + if we are in the progress of fulfilling a constraint, that's the only one we'll return. ''' if self.inprogress_constraint is None: token_list = [] - for constraint in self.pending_constraints: + for constraint in self.pending_constraints: # "pending" == "unfulfilled yet" advance = constraint.advance() token_list.append(advance) else: @@ -244,36 +288,49 @@ def advance(self): else: return torch.stack(token_list) - def update(self, token_ids: torch.Tensor): + def reset(self, token_ids: Optional[torch.Tensor]): ''' token_ids: the tokens generated thus far to reset the state of the progress through constraints. ''' self.init_state() - for token in token_ids: - complete, stepped = self.add(token) + + if token_ids is not None: + for token in token_ids: + complete, stepped = self.add(token) return self + def add(self, token_id: Union[int, torch.LongTensor]): + complete, stepped = False, False + + if isinstance(token_id, torch.LongTensor): + if (token_id.size(0)) > 1: + raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." + "It must have length 1.") + else: + token_id = token_id[0] - def add(self, token_id: int): if self.completed: - return True, True + complete = True + stepped = False + return complete, stepped - complete, stepped = False, False if self.inprogress_constraint is not None: ''' In the middle of fulfilling a constraint. - - If the token just steps (make but an incremental progress to current job) it, do nothing. + If the `token_id` *does* makes an incremental progress to current job, simply update the state ''' stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: ''' - 1. If the next token breaks the fulfillment, then we must restart. - e.g. force the sequence "I love pies" and the next token after "I love" is "books". + 1. If the next token breaks the progress, then we must restart. + e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books". + + But that doesn't mean we self.init_state(), since we only reset the state + for this particular constraint, not the full list of constraints. ''' - self.pending_constraints.append(self.inprogress_constraint.copy()) + self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False)) self.inprogress_constraint = None if complete: @@ -286,31 +343,60 @@ def add(self, token_id: int): self.inprogress_constraint = None if len(self.pending_constraints) == 0: + # we're done! self.completed = True + else: ''' - Not in the middle of fulfilling a constraint. + Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards + any of our list of constraints? ''' for cidx, pending_constraint in enumerate(self.pending_constraints): - ''' - 1. Does it advance any of the pending constraints? - ''' if pending_constraint.does_advance(token_id): stepped, complete, reset = pending_constraint.update(token_id) + + if not stepped: + raise Exception("constraint.update(token_id) is not yielding incremental progress, " + "even though constraint.does_advance(token_id) is true.") + if complete: self.complete_constraints.append(pending_constraint) self.inprogress_constraint = None + if not complete and stepped: self.inprogress_constraint = pending_constraint if complete or stepped: + ''' + If we made any progress at all, then it's at least not a "pending constraint". + ''' self.pending_constraints = self.pending_constraints[:cidx] + self.pending_constraints[cidx+1:] + if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: + ''' + If there's no longer any pending after this and no inprogress either, then we must be complete. + ''' self.completed = True - break + break # prevent accidentally stepping through multiple constraints with just one token. return complete, stepped - - \ No newline at end of file + def copy(self, stateful=True): + new_state = ConstraintListState(self.constraints) # we actually never though self.constraints objects + # throughout this process. So it's at initialization state. + + if stateful: + new_state.complete_constraints = [ + constraint.copy(stateful=True) + for constraint in self.complete_constraints + ] + if self.inprogress_constraint is not None: + new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True) + new_state.pending_constraints = [ + constraint.copy() + for constraint in self.pending_constraints + ] + + return new_state + diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 5ec0754f1a57..eeb328ae89d8 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -452,7 +452,6 @@ def make_constraint_states(self, n): for _ in range(n) ] - def process( self, input_ids: torch.LongTensor, @@ -506,17 +505,13 @@ def process( if (eos_token_id is not None) and (next_token.item() == eos_token_id): # if constraint not fulfilled, it should not be added. if not self.constraints_completed[batch_idx][next_index]: - continue - else: - print("NOT COMPLETE", self.constraints_completed) + continue # if beam_token does not belong to top num_beams tokens, it should not be added is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size if is_beam_token_worse_than_top_num_beams: continue - print("\n!\n!\n!\n!\n!\n!\n!\n!\n!TRULY COMPLETED SEQUENCE!!\n!\n!\n!") - print("input_ids[batch_beam_idx].clone()", input_ids[batch_beam_idx].clone()) beam_hyp.add( input_ids[batch_beam_idx].clone(), next_score.item(), @@ -537,15 +532,14 @@ def process( input_ids, scores_for_all_vocab, self.constraints_completed[batch_idx], - next_beam_scores[batch_idx].clone(), - next_beam_tokens[batch_idx].clone(), - next_beam_indices[batch_idx].clone(), + next_beam_scores[batch_idx], + next_beam_tokens[batch_idx], + next_beam_indices[batch_idx], ) next_beam_scores[batch_idx] = new_scores next_beam_tokens[batch_idx] = new_tokens next_beam_indices[batch_idx] = new_indices - self.constraints_completed[batch_idx] = new_completed @@ -569,14 +563,17 @@ def process( def step_sentence_constraint( self, - batch_idx, - input_ids, - vocab_scores, - constraints_completed, - sent_beam_scores, - sent_beam_tokens, - sent_beam_indices, + batch_idx: int, + input_ids: torch.LongTensor, + vocab_scores: torch.FloatTensor, + constraints_completed: List[bool], + sent_beam_scores: torch.FloatTensor, + sent_beam_tokens: torch.LongTensor, + sent_beam_indices: torch.LongTensor, + push_progress: bool = False ): + # from transformers import GPT2Tokenizer + # tokenizer = GPT2Tokenizer.from_pretrained("gpt2") ''' sent_beam_tokens are the next {num_beams} number of tokens that are under consideration for this beam (candidate next tokens) @@ -591,75 +588,102 @@ def step_sentence_constraint( orig_len = sent_beam_indices.size(0) device = sent_beam_indices.get_device() - topk_contraint_states = self.make_constraint_states(orig_len) + # initialize states + topk_contraint_states = self.make_constraint_states(orig_len) advance_constraint_states = self.make_constraint_states(orig_len) - start_idx = batch_idx*orig_len - end_idx = (batch_idx+1) * orig_len + sidx, eidx = batch_idx*orig_len, (batch_idx+1) * orig_len - this_batch_input_ids = input_ids[start_idx : end_idx] - this_batch_token_scores = vocab_scores[start_idx : end_idx] + this_batch_input_ids = input_ids[sidx:eidx] + this_batch_token_scores = vocab_scores[sidx:eidx] - full_hypotheses = torch.cat((this_batch_input_ids, sent_beam_tokens.unsqueeze(-1)), dim=-1) - + full_hypotheses = torch.cat(( + this_batch_input_ids[sent_beam_indices], + sent_beam_tokens.unsqueeze(-1)), + dim=-1 + ) # need to make new hypothesis that advance the constraints - new_seqs = [] - new_states = [] - new_indices = [] - new_tokens = [] - new_scores = [] + track_new = {"new_seqs": [], "new_states": [], "new_indices": [], "new_tokens": [], "new_scores": []} for seq_idx, pre_seq in enumerate(this_batch_input_ids): - topk_state = topk_contraint_states[sent_beam_indices[seq_idx].item()] - topk_state.update(full_hypotheses[seq_idx]) + ''' + pre_seq = ith sequence generated before this step. - if constraints_completed[seq_idx]: - continue - advance_state = advance_constraint_states[seq_idx] - advance_state.update(pre_seq) - if not advance_state.completed: + input_ids -> (topk) generic beam search best model next tokens + -> (advance) constraints forcing the next token + either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of hypotheses. + ''' + topk_state = topk_contraint_states[seq_idx] + topk_state.reset(full_hypotheses[seq_idx]) + + if not constraints_completed[seq_idx]: + advance_state = advance_constraint_states[seq_idx] + advance_state.reset(pre_seq) + # the curerent `pre_seq` does not yet satisfy our list of constraints. advance_tokens = advance_state.advance() - print(">>>>>advance_tokens", advance_tokens) - if advance_tokens.numel() != 0: + for advance_token in advance_tokens: + # since adding each `advance_token` leads to a different hypothesis, create new state instance. new_state = advance_state.copy(stateful=True) - new_state.add(advance_tokens[:1]) - - for advance_token in advance_tokens: - advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).cpu().tolist() - if advance_seq not in new_seqs: - print("advance_seq", advance_seq) - print("new_seqs", new_seqs) - new_seqs.append(advance_seq) - new_score = this_batch_token_scores[seq_idx].take(advance_token) - new_indices.append(seq_idx) - new_tokens.append(advance_token) - new_scores.append(new_score) - new_states.append(new_state) - # else: - # # if it's a complete state - # new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) - # new_indices.append(seq_idx) - # new_tokens.append(new_token) - # new_scores.append(new_score) - # new_states.append(advance_state) + new_state.add(advance_token) + + advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).cpu().tolist() + if advance_seq not in track_new["new_seqs"]: + # prevent duplicates, which are basically bound to happen in this process. + track_new["new_seqs"].append(advance_seq) + track_new["new_indices"].append(seq_idx) + track_new["new_tokens"].append(advance_token) + track_new["new_scores"].append(this_batch_token_scores[seq_idx].take(advance_token)) + track_new["new_states"].append(new_state) + elif push_progress: + ''' + Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences + that actually fulfill our constraints. For example, let constraints == ["loves pies"] and + + pre_seq_1 = "The child loves pies and" + pre_seq_2 = "The child plays in the playground and" + + Without this step, if `sent_beam_indices` is something like [1,1], then + 1. `pre_seq_1` won't be added to the list of (topk) hypothesis since it's not in the indices and + 2. it won't be added to the list of (advance) hypothesis since it's completed already. + (this is the else part of `if constraints_completed[seq_idx]`) + 3. it ends up simply getting removed from consideration. + + #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, especially + if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam search times, + since completed sequences keep getting removed after all this effort for constrained generation. + + Here, we basically take `pre_seq_1` and to "push" it into the considered list of + hypotheses, by simply appending the next likely token in the vocabulary and adding it to the list of hypotheses. + ''' + new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) # some next probable token + advance_seq = torch.cat((pre_seq, new_token.unsqueeze(0)), -1) + + advance_state = advance_constraint_states[seq_idx] + + advance_state.reset(advance_seq) + advance_seq = advance_seq.cpu().tolist() + if advance_seq not in track_new["new_seqs"]: + # but still don't want to have duplicates + track_new["new_seqs"].append(advance_seq) + track_new["new_indices"].append(seq_idx) + track_new["new_tokens"].append(new_token) + track_new["new_scores"].append(new_score) + track_new["new_states"].append(advance_state) + + if len(track_new["new_indices"]) > 0: + new_indices = torch.tensor(track_new["new_indices"]).to(device) + new_tokens = torch.stack(track_new["new_tokens"]).to(device) + new_scores = torch.stack(track_new["new_scores"]).to(device) - if len(new_indices) > 0: - new_indices = torch.tensor(new_indices).to(device) - new_tokens = torch.stack(new_tokens).to(device) - new_scores = torch.stack(new_scores).to(device) - - all_states = topk_contraint_states + new_states + all_states = topk_contraint_states + track_new["new_states"] all_tokens = torch.cat((sent_beam_tokens, new_tokens), -1) all_scores = torch.cat((sent_beam_scores, new_scores), -1) - all_banks = torch.tensor([one.get_bank() for one in all_states]).to(device) - # ! only for testing!! - zipped = torch.stack((all_scores, all_banks, all_tokens)) - zipped = torch.transpose(zipped, 0, 1) + all_banks = torch.tensor([one.get_bank() + for one in all_states + ]).to(device) - augmented_zipped = all_banks * 1000 + all_scores*10 + augmented_zipped = all_banks * 100 + all_scores indices = augmented_zipped.sort(descending=True).indices - - sorted_banks = all_banks[indices] # C, C, C-1, C-2, ..., 1, 0, 0 - + sorted_banks = all_banks[indices] ''' Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} @@ -676,74 +700,15 @@ def step_sentence_constraint( increments.append(counter) rearrangers = torch.tensor(np.argsort(increments, kind="mergesort")) - print("PREZIP", zipped) indices = indices[rearrangers] - print("POSTZIP", zipped[indices]) - sent_beam_scores = all_scores[indices] - sent_beam_tokens = all_tokens[indices] - sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices] + sent_beam_scores = all_scores[indices][:orig_len] + sent_beam_tokens = all_tokens[indices][:orig_len] + sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices][:orig_len] constraints_completed = [all_states[idx].completed for idx in indices[:orig_len]] - sent_beam_scores = sent_beam_scores[:orig_len] - sent_beam_tokens = sent_beam_tokens[:orig_len] - sent_beam_indices = sent_beam_indices[:orig_len] - - print(">>>>>sent_beam_scores", sent_beam_scores) - print(">>>>>sent_beam_tokens", sent_beam_tokens) - print(">>>>>sent_beam_indices", sent_beam_indices) - print(">>>>constraints_completed", constraints_completed) - return constraints_completed, sent_beam_scores, sent_beam_tokens, sent_beam_indices - - - # if not sent_constraint_state.completed: - # advance_tokens = sent_constraint_state.advance() - # print(">>>>advance_tokens", advance_tokens) - # if advance_tokens.numel() != 0: - # advance_tokens = advance_tokens[:1] - - # sent_constraint_state.update(advance_tokens) - - # additional_num = advance_tokens.size(0) - # next_beam_tokens = ( - # advance_tokens - # .repeat(additional_num) - # .long() - # .to(device) - # ) - # sent_beam_tokens = torch.cat((sent_beam_tokens, next_beam_tokens)) - - # sent_beam_scores = torch.cat(( - # sent_beam_scores, - # sent_vocab_scores.take(advance_tokens) - # )) - - # ''' - # 2. Compute "banks" for each candidate. - # If C is the number of constraints, we construct C banks, where ith bank - # is the bank for candidates that have fulfilled ith constraint. - # ''' - - # sent_beam_scores = sent_beam_scores[-orig_len:] - # sent_beam_tokens = sent_beam_tokens[-orig_len:] - # print("???>>", sent_beam_tokens) - # sent_beam_indices = torch.tensor([ - # batch_idx*orig_len + 1, - # batch_idx*orig_len + 2, - # batch_idx*orig_len + 3, - # batch_idx*orig_len + 3, - # ]).to(device) - # print(">>>>>>>>>SELECTIUON INDICES??", sent_beam_indices) - # # sent_beam_indices = sent_beam_indices[-orig_len:] - # assert sent_beam_scores.size(0) == orig_len - - - - # return sent_constraint_state, sent_beam_scores, sent_beam_tokens, sent_beam_indices - - def finalize( self, input_ids: torch.LongTensor, diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 1eb1bbaad361..74e880a287cd 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -2965,10 +2965,6 @@ def constrained_beam_search( next_indices = (next_tokens / vocab_size).long() next_tokens = next_tokens % vocab_size - print("\n\nINPUT\n") - for one in input_ids: - print(tokenizer.decode(one)) - print("\n\n\n") # stateless beam_outputs = constrained_beam_scorer.process( input_ids, @@ -2985,11 +2981,6 @@ def constrained_beam_search( input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - print("\n\nOUTPUT\n") - for one in input_ids: - print(tokenizer.decode(one)) - print("\n\n\n") - model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder ) diff --git a/tests/test.py b/tests/test.py deleted file mode 100644 index 77cd4edd1931..000000000000 --- a/tests/test.py +++ /dev/null @@ -1,43 +0,0 @@ -import unittest - -from transformers import GPT2Tokenizer, GPT2LMHeadModel -from transformers.generation_beam_constraints import ( - PhrasalConstraint -) -device = "cuda" - -model = GPT2LMHeadModel.from_pretrained("gpt2").to(device) -tokenizer = GPT2Tokenizer.from_pretrained("gpt2") - -force_text = " big monsters" -force_text_2 = " crazy" -force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device)[0] -force_tokens_2 = tokenizer.encode(force_text_2, return_tensors="pt").to(device)[0] - -print("force_tokens", force_tokens) -print("force_tokens_2", force_tokens_2) -constraints = [ - PhrasalConstraint(force_tokens), - PhrasalConstraint(force_tokens_2) -] - -input_text = ["The baby is crying because"] * 1 - -model_inputs = tokenizer(input_text, return_tensors="pt") - -for key, value in model_inputs.items(): - model_inputs[key] = value.to(device) - -print("model_inputs", model_inputs) -k = model.generate( - **model_inputs, - constraints=constraints, - num_beams=5, - num_return_sequences=5 -) - -for out in k: - print("!!", out) - print(tokenizer.decode(out)) - -assert False From 125a9aa06f2ed7ffe40efc771668120f7911d417 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 11:07:02 +0000 Subject: [PATCH 06/34] removed incorrect tests --- tests/test_generation_beam_search.py | 234 --------------------------- 1 file changed, 234 deletions(-) diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index 51df97a22984..ac557720b093 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -239,222 +239,6 @@ def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_ self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) -class ConstrainedBeamSearchTester: - def __init__( - self, - parent, - batch_size=3, - sequence_length=10, - vocab_size=99, - pad_token_id=0, - max_length=20, - num_beams=4, - length_penalty=2.0, - do_early_stopping=True, - num_beam_hyps_to_keep=2, - ): - self.parent = parent - self.batch_size = batch_size - self.sequence_length = sequence_length - self.vocab_size = vocab_size - self.pad_token_id = pad_token_id - self.max_length = max_length - self.num_beams = num_beams - self.length_penalty = length_penalty - self.do_early_stopping = do_early_stopping - self.num_beam_hyps_to_keep = num_beam_hyps_to_keep - - self.constraints = [ - PhrasalConstraint(ids_tensor((1, 2), self.vocab_size)[0]) - ] - # cannot be randomly generated - self.eos_token_id = vocab_size + 1 - - def prepare_beam_scorer(self, **kwargs): - - return ConstrainedBeamSearchScorer( - batch_size=kwargs.get("batch_size", self.batch_size), - num_beams=kwargs.get("num_beams", self.num_beams), - constraints=kwargs.get("constraints", self.constraints), - device=torch_device, - length_penalty=kwargs.get("length_penalty", self.length_penalty), - do_early_stopping=kwargs.get("do_early_stopping", self.do_early_stopping), - num_beam_hyps_to_keep=kwargs.get("num_beam_hyps_to_keep", self.num_beam_hyps_to_keep), - ) - - def prepare_inputs(self): - constraint_states = [ - ConstraintListState(self.constraints) - for _ in range(self.batch_size) - ] # n - - input_ids = ids_tensor((self.batch_size * self.num_beams, self.sequence_length), self.vocab_size) - next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) - next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) - next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) - return (input_ids, next_tokens, next_indices, next_scores, constraint_states) - - def check_beam_hypotheses(self, input_ids, *args): - # check that correct number of beam hypotheses is set in beam scorer - beam_scorer = self.prepare_beam_scorer(do_early_stopping=True) - beam_hyp = beam_scorer._beam_hyps[0] - - self.parent.assertEqual(len(beam_scorer._beam_hyps), self.batch_size) - - # check correct type - self.parent.assertTrue(isinstance(beam_hyp, BeamHypotheses)) - - # check that num_beams is correctly set - self.parent.assertEqual(beam_hyp.num_beams, self.num_beams) - - # check for early stopping deactivated - for beam_idx in range(self.num_beams): - beam_hyp.add(input_ids[beam_idx], -10.0) - - # if early stopping True -> score does not matter - self.parent.assertTrue(beam_hyp.is_done(-10.0, 5)) - - # re-init - beam_scorer = self.prepare_beam_scorer(do_early_stopping=False) - beam_hyp = beam_scorer._beam_hyps[0] - - # add `num_beams + 1` beams to change `worst_score` - for beam_idx in range(self.num_beams + 1): - beam_hyp.add(input_ids[beam_idx], -10.0 + float(beam_idx)) - - # -10.0 is removed => -9.0 is worst score - self.parent.assertAlmostEqual(beam_hyp.worst_score, -9.0 / (self.sequence_length ** beam_hyp.length_penalty)) - - # -5.0 is better than worst score => should not be finished - self.parent.assertFalse(beam_hyp.is_done(-5.0, self.sequence_length)) - - # -20.0 is worse than worst score => should be finished - self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) - - def check_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores, constraint_states): - # check too many eos tokens - constrained_beam_scorer = self.prepare_beam_scorer() - print("---1next_scores", next_scores) - - tokens = next_tokens.clone() - tokens[0, :] = self.eos_token_id - - with self.parent.assertRaises(ValueError): - constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id) - - # check all batches are done - constrained_beam_scorer = self.prepare_beam_scorer() - - tokens = next_tokens.clone() - tokens[:, : self.num_beams] = self.eos_token_id - constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id) - # beam scorer should be done - self.parent.assertTrue(constrained_beam_scorer.is_done) - - # check - constrained_beam_scorer = self.prepare_beam_scorer() - - tokens = next_tokens.clone() - tokens[:, 1] = self.eos_token_id - beam_outputs = constrained_beam_scorer.process( - input_ids, next_scores, tokens, next_indices, constraint_states, eos_token_id=self.eos_token_id - ) - output_scores = beam_outputs["next_beam_scores"] - output_tokens = beam_outputs["next_beam_tokens"] - output_indices = beam_outputs["next_beam_indices"] - - def cut_expected_tensor(tensor): - return torch.cat([tensor[:, :1], tensor[:, 2 : self.num_beams + 1]], dim=1).flatten() - - # check all outptus - # cut out id of eos token and take best `num_beams` outputs - expected_output_tokens = cut_expected_tensor(tokens) - expected_output_scores = cut_expected_tensor(next_scores) - - # add num_beams * batch_idx - expected_output_indices = ( - cut_expected_tensor(next_indices) - + (torch.arange(self.num_beams * self.batch_size, device=torch_device) // self.num_beams) * self.num_beams - ) - - self.parent.assertListEqual(expected_output_tokens.tolist(), output_tokens.tolist()) - self.parent.assertListEqual(expected_output_indices.tolist(), output_indices.tolist()) - self.parent.assertTrue(torch.allclose(expected_output_scores, output_scores, atol=1e-3)) - - # make sure ids of eos token are correctly saved in beam_hyps of beam scorer - for batch_idx in range(self.batch_size): - correct_idx = batch_idx * self.num_beams + next_indices[batch_idx, 1] - self.parent.assertListEqual( - input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() - ) - - def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_scores): - # max_length should be only one more than current input_ids to check that eos is correctly appended - max_length = self.sequence_length + 1 - beam_scorer = self.prepare_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) - - # update beams and append to input_ids - tokens = next_tokens.clone() - # first batch, first output has to finish with eos token id since scores are correctly sorted - tokens[0, 0] = self.eos_token_id - # make sure corresponding score is as good as possible to surely be picked first - next_scores[0, 0] = 0.0 - beam_outputs = beam_scorer.process( - input_ids, next_scores, tokens, next_indices, eos_token_id=self.eos_token_id - ) - output_scores = beam_outputs["next_beam_scores"] - output_tokens = beam_outputs["next_beam_tokens"] - output_indices = beam_outputs["next_beam_indices"] - - input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) - - # finalize - sequence_output = beam_scorer.finalize( - input_ids, - output_scores, - output_tokens, - output_indices, - pad_token_id=self.pad_token_id, - eos_token_id=self.eos_token_id, - max_length=max_length, - ) - - sequences = sequence_output["sequences"] - sequence_scores = sequence_output["sequence_scores"] - - # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length` - self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length]) - self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size]) - - # check sequence_scores - self.parent.assertFalse((sequence_scores > 0).any().item()) - - # first batch has to finish with eos_token - self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id) - - # other batches cannot finish with eos token - self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id) - self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id) - - # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned - beam_scorer.num_beam_hyps_to_keep = self.num_beams - sequence_output = beam_scorer.finalize( - input_ids, - output_scores, - output_tokens, - output_indices, - pad_token_id=self.pad_token_id, - eos_token_id=self.eos_token_id, - max_length=max_length, - ) - sequences = sequence_output["sequences"] - sequence_scores = sequence_output["sequence_scores"] - - self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) - self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) - - - @require_torch class BeamSearchTest(unittest.TestCase): def setUp(self): @@ -472,21 +256,3 @@ def test_beam_scorer_finalize(self): inputs = self.beam_search_tester.prepare_inputs() self.beam_search_tester.check_beam_scores_finalize(*inputs) - - -@require_torch -class ConstrainedBeamSearchTest(unittest.TestCase): - def setUp(self): - self.constrained_beam_search_tester = ConstrainedBeamSearchTester(self) - - # def test_beam_hypotheses(self): - # inputs = self.constrained_beam_search_tester.prepare_inputs() - # self.constrained_beam_search_tester.check_beam_hypotheses(*inputs) - - def test_beam_scorer_update(self): - inputs = self.constrained_beam_search_tester.prepare_inputs() - self.constrained_beam_search_tester.check_beam_scorer_update(*inputs) - - # def test_beam_scorer_finalize(self): - # inputs = self.constrained_beam_search_tester.prepare_inputs() - # self.constrained_beam_search_tester.check_beam_scores_finalize(*inputs) From 9fcba0d5f1dda59c029133a601bc4f75c39cc081 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 20:11:10 +0900 Subject: [PATCH 07/34] Delete k.txt --- k.txt | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 k.txt diff --git a/k.txt b/k.txt deleted file mode 100644 index e69de29bb2d1..000000000000 From 377349589e73626c16d7232e20d482774ff8045a Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 20:11:24 +0900 Subject: [PATCH 08/34] Delete test.py --- test.py | 27 --------------------------- 1 file changed, 27 deletions(-) delete mode 100644 test.py diff --git a/test.py b/test.py deleted file mode 100644 index 75e1dc07ca6e..000000000000 --- a/test.py +++ /dev/null @@ -1,27 +0,0 @@ -import unittest - -from transformers import BartForConditionalGeneration, BartTokenizer -from transformers.generation_beam_constraints import ( - PhrasalConstraint -) -model = BartForConditionalGeneration.from_pretrained("facebook/bart-base") -tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") - -force_text = "forced" -force_tokens = tokenizer.encode(force_text, return_tensors="pt") -print("force_tokens", force_tokens) -constraints = [PhrasalConstraint(force_tokens[1:])] - -input_text = ["This feels a little"] - -model_inputs = tokenizer(input_text, return_tensors="pt") - -print("input_ids", input_ids) -k = model.generate( - **model_inputs, - constraints=constraints -) - -print(k) - -assert False From d214c83a4d8441a052c942a90b0e3ffafe3d41b9 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 20:11:33 +0900 Subject: [PATCH 09/34] Delete test.sh --- test.sh | 1 - 1 file changed, 1 deletion(-) delete mode 100644 test.sh diff --git a/test.sh b/test.sh deleted file mode 100644 index a7bfb7d17e4c..000000000000 --- a/test.sh +++ /dev/null @@ -1 +0,0 @@ -CUDA_LAUNCH_BLOCKING=1, pytest -s tests/test_modeling_bart.py::BartStandaloneDecoderModelTest::test_constrained_beam_search_generate --capture=sys From 26255809121a073a54ebc239c0235bfe7971ef5a Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sun, 23 Jan 2022 11:14:57 +0000 Subject: [PATCH 10/34] revert changes to test scripts --- tests/test_generation_beam_search.py | 16 +-- tests/test_generation_utils.py | 160 +-------------------------- 2 files changed, 6 insertions(+), 170 deletions(-) diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index ac557720b093..216284915e93 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -17,10 +17,7 @@ import unittest from transformers import is_torch_available -# from transformers.testing_utils import require_torch, torch_device -from transformers.testing_utils import require_torch - -torch_device = "cpu" +from transformers.testing_utils import require_torch, torch_device from .test_modeling_common import floats_tensor, ids_tensor @@ -28,12 +25,8 @@ if is_torch_available(): import torch - from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer, ConstrainedBeamSearchScorer - from transformers.generation_beam_constraints import ( - Constraint, - PhrasalConstraint, - ConstraintListState - ) + from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer + class BeamSearchTester: def __init__( @@ -254,5 +247,4 @@ def test_beam_scorer_update(self): def test_beam_scorer_finalize(self): inputs = self.beam_search_tester.prepare_inputs() - self.beam_search_tester.check_beam_scores_finalize(*inputs) - + self.beam_search_tester.check_beam_scores_finalize(*inputs) \ No newline at end of file diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index a21f943e61d3..b36fadafd0d7 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -37,7 +37,7 @@ VisionEncoderDecoderModel, top_k_top_p_filtering, ) - from transformers.generation_beam_search import BeamSearchScorer, ConstrainedBeamSearchScorer + from transformers.generation_beam_search import BeamSearchScorer from transformers.generation_logits_process import ( ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor, @@ -63,9 +63,6 @@ SampleDecoderOnlyOutput, SampleEncoderDecoderOutput, ) - from transformers.generation_beam_constraints import ( - PhrasalConstraint - ) class GenerationTesterMixin: @@ -171,25 +168,6 @@ def _get_beam_scorer_and_kwargs(batch_size, max_length, num_return_sequences=1): num_beam_hyps_to_keep=num_return_sequences, ) return beam_kwargs, beam_scorer - - @staticmethod - def _get_constrained_beam_scorer_and_kwargs(batch_size, max_length, constraints, num_return_sequences=1): - beam_kwargs = { - "early_stopping": False, - "length_penalty": 2.0, - "num_beams": 2, - "num_return_sequences": num_return_sequences, - } - beam_scorer = ConstrainedBeamSearchScorer( - batch_size=batch_size, - num_beams=beam_kwargs["num_beams"], - device=torch_device, - constraints=constraints, - length_penalty=beam_kwargs["length_penalty"], - do_early_stopping=beam_kwargs["early_stopping"], - num_beam_hyps_to_keep=num_return_sequences, - ) - return beam_kwargs, beam_scorer @staticmethod def _get_diverse_beam_scorer_and_kwargs(batch_size, max_length, num_return_sequences=1): @@ -420,74 +398,6 @@ def _beam_search_generate( ) return output_generate, output_beam_search - def _constrained_beam_search_generate( - self, - model, - input_ids, - attention_mask, - max_length, - constraints, - constrained_beam_scorer, - beam_kwargs, - logits_processor, - logits_process_kwargs, - output_scores=False, - output_attentions=False, - output_hidden_states=False, - return_dict_in_generate=False, - ): - output_generate = model.generate( - input_ids, - attention_mask=attention_mask, - do_sample=False, - max_length=max_length, - output_scores=output_scores, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict_in_generate=return_dict_in_generate, - remove_invalid_values=True, - constraints=constraints, - **beam_kwargs, - **logits_process_kwargs, - ) - print("output_generate", output_generate) - - # beam_search does not automatically interleave `batch_size` dim for `num_beams` - kwargs = {} - if model.config.is_encoder_decoder: - encoder_outputs, input_ids_clone, attention_mask_clone = self._get_encoder_outputs( - model, - input_ids, - attention_mask, - num_interleave=constrained_beam_scorer.num_beams, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - ) - kwargs["encoder_outputs"] = encoder_outputs - input_ids_clone = input_ids_clone.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) - else: - attention_mask_clone = attention_mask.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) - input_ids_clone = input_ids.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) - - with torch.no_grad(): - output_beam_search = model.constrained_beam_search( - input_ids_clone, - constraints, - constrained_beam_scorer, - max_length=max_length, - attention_mask=attention_mask_clone, - logits_processor=logits_processor, - output_scores=output_scores, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict_in_generate=return_dict_in_generate, - **kwargs, - ) - - print("output_beam_search", output_beam_search) - - return output_generate, output_beam_search - def _beam_sample_generate( self, model, @@ -830,72 +740,6 @@ def test_beam_search_generate(self): logits_processor=logits_processor, ) self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) - - def test_constrained_beam_search_generate(self): - for model_class in self.all_generative_model_classes: - config, input_ids, attention_mask, max_length = self._get_input_ids_and_config() - constraints = [ - PhrasalConstraint(ids_tensor((1, 2), config.vocab_size)[0]) - ] - - # It is important set set the eos_token_id to None to ensure that no sequences - # shorter than `max_length` can be generated which could lead to flaky circle ci - # failures if the top `num_return_sequences` beams are all shorter than the longest beam - config.eos_token_id = None - config.forced_eos_token_id = None - - model = model_class(config).to(torch_device).eval() - if model.config.is_encoder_decoder: - max_length = 4 - - logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( - input_ids.shape[-1], - config.eos_token_id, - config.forced_bos_token_id, - config.forced_eos_token_id, - max_length, - ) - beam_kwargs, constrained_beam_scorer = self._get_constrained_beam_scorer_and_kwargs( - input_ids.shape[0], - max_length, - constraints - ) - - # check `generate()` and `beam_search()` are equal - output_generate, output_beam_search = self._constrained_beam_search_generate( - model=model, - input_ids=input_ids, - constraints=constraints, - constrained_beam_scorer=constrained_beam_scorer, - attention_mask=attention_mask, - max_length=max_length, - beam_kwargs=beam_kwargs, - logits_process_kwargs=logits_process_kwargs, - logits_processor=logits_processor, - ) - assert False - self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) - - # check `generate()` and `beam_search()` are equal for `num_return_sequences` - num_return_sequences = 2 - if model.config.is_encoder_decoder: - max_length = 4 - beam_kwargs, beam_scorer = self._get_beam_scorer_and_kwargs( - input_ids.shape[0], max_length, num_return_sequences=num_return_sequences - ) - - output_generate, output_beam_search = self._beam_search_generate( - model=model, - input_ids=input_ids, - attention_mask=attention_mask, - max_length=max_length, - beam_scorer=beam_scorer, - beam_kwargs=beam_kwargs, - logits_process_kwargs=logits_process_kwargs, - logits_processor=logits_processor, - ) - self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) - def test_beam_search_generate_dict_output(self): for model_class in self.all_generative_model_classes: @@ -2058,4 +1902,4 @@ def test_generate_encoder_outputs_attention_mask(self): output_sequences_with_mask = model.generate(encoder_outputs=encoder_outputs, attention_mask=attention_mask) output_sequences_with_mask = output_sequences_with_mask.cpu() - self.assertListEqual(output_sequences_no_mask.tolist(), output_sequences_with_mask.tolist()) + self.assertListEqual(output_sequences_no_mask.tolist(), output_sequences_with_mask.tolist()) \ No newline at end of file From 97dc8ccc007b94008210b3867d999c7e0cca0bd3 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Sat, 29 Jan 2022 11:28:04 +0000 Subject: [PATCH 11/34] genutils --- tests/test_generation_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index b36fadafd0d7..73f838ebf18b 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -1896,7 +1896,7 @@ def test_generate_encoder_outputs_attention_mask(self): encoder = model.get_encoder() - encoder_outputs = encoder(input_values) + encoder_outputs = encoder(inp/gut_values) output_sequences_no_mask = model.generate(encoder_outputs=encoder_outputs).cpu() output_sequences_with_mask = model.generate(encoder_outputs=encoder_outputs, attention_mask=attention_mask) From 91a6403905525cc4600970e0ca8bcd2081aad8a1 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 06:11:20 +0000 Subject: [PATCH 12/34] full implementation with testing, no disjunctive yet --- .../generation_beam_constraints.py | 27 +- src/transformers/generation_beam_search.py | 63 ++-- src/transformers/generation_utils.py | 29 +- tests/test_generation_beam_search.py | 299 +++++++++++++++++- tests/test_generation_utils.py | 255 ++++++++++++++- 5 files changed, 616 insertions(+), 57 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 2352df1fdf64..1b7b4afe6fc1 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -136,7 +136,7 @@ class TokenConstraint(Constraint): def __init__(self, token_id: Union[int, torch.LongTensor]): super(Constraint, self).__init__() if not (isinstance(token_id, int) or isinstance(token_id, torch.LongTensor)) or token_id < 0: - raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integeter, but is {token_id}") + raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integer, but is {token_id}") else: if isinstance(token_id, torch.LongTensor) and token_id.size(0) > 1: raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." @@ -181,10 +181,18 @@ class PhrasalConstraint(Constraint): token_ids (`torch.Tensor`): The sequence of tokens that must be generated by the output. """ - def __init__(self, token_ids: torch.Tensor): + def __init__(self, token_ids: Union[List[int], torch.LongTensor]): super(Constraint, self).__init__() - if not isinstance(token_ids, torch.Tensor): - raise ValueError(f"`token_ids` has to be a `torch.Tensor`, but is {type(token_ids)}") + + is_int_list = isinstance(token_ids, List) and isinstance(token_ids[0], int) + is_long_tensor = isinstance(token_ids, torch.LongTensor) and len(token_ids.size()) == 1 + if not (is_int_list or is_long_tensor) or torch.any(token_ids < 0): + raise ValueError(f"`token_ids` has to be a single list of positive integers or a `torch.LongTensor` but is {token_ids}") + else: + if (is_int_list or is_long_tensor) and token_ids.size(0) == 1: + raise ValueError(f"`token_ids` has to be list of positive integers or a `torch.LongTensor` but is {token_ids}" + "For single token constraints, refer to `TokenConstraint`.") + self.token_ids = token_ids self.seqlen = self.token_ids.size(0) @@ -194,11 +202,10 @@ def __init__(self, token_ids: torch.Tensor): def advance(self): return self.token_ids[self.fulfilled_idx + 1] - def does_advance(self, token_id): + def does_advance(self, token_id: int): if self.completed: return False - - return token_id == self.token_ids[self.fulfilled_idx + 1] + return token_id.cpu() == self.token_ids[self.fulfilled_idx + 1] def update(self, token_id: int): stepped = False @@ -288,16 +295,18 @@ def advance(self): else: return torch.stack(token_list) - def reset(self, token_ids: Optional[torch.Tensor]): + def reset(self, token_ids: Optional[torch.LongTensor]): ''' token_ids: the tokens generated thus far to reset the state of the progress through constraints. ''' self.init_state() - if token_ids is not None: + if token_ids is not None and token_ids.size(0) > 0: for token in token_ids: complete, stepped = self.add(token) + if complete: + break return self diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index eeb328ae89d8..f20616b54b8c 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -404,9 +404,6 @@ def __init__( self.num_beam_groups = num_beam_groups self.group_size = self.num_beams // self.num_beam_groups self.constraints = constraints - self.constraints_completed = [ - [False] * self.group_size - ] * batch_size self._is_init = False self._beam_hyps = [ @@ -434,7 +431,7 @@ def __init__( if "max_length" in kwargs: warnings.warn( - "Passing `max_length` to BeamSearchScorer is deprecated and has no effect. " + "Passing `max_length` to ConstrainedBeamSearchScorer is deprecated and has no effect. " "`max_length` should be passed directly to `beam_search(...)`, `beam_sample(...)`" ", or `group_beam_search(...)`." ) @@ -452,6 +449,11 @@ def make_constraint_states(self, n): for _ in range(n) ] + def check_completes_constraints(self, sequence): + new_state = self.make_constraint_states(1)[0] + new_state = new_state.reset(sequence) + return new_state.completed + def process( self, input_ids: torch.LongTensor, @@ -503,19 +505,18 @@ def process( batch_beam_idx = batch_idx * self.group_size + next_index # add to generated hypotheses if end of sentence if (eos_token_id is not None) and (next_token.item() == eos_token_id): - # if constraint not fulfilled, it should not be added. - if not self.constraints_completed[batch_idx][next_index]: - continue # if beam_token does not belong to top num_beams tokens, it should not be added is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size if is_beam_token_worse_than_top_num_beams: continue - beam_hyp.add( - input_ids[batch_beam_idx].clone(), - next_score.item(), - ) + completes_constraint = self.check_completes_constraints(input_ids[batch_beam_idx]) + if completes_constraint: + beam_hyp.add( + input_ids[batch_beam_idx].clone(), + next_score.item(), + ) else: # add next predicted token since it is not eos_token next_beam_scores[batch_idx, beam_idx] = next_score @@ -527,11 +528,10 @@ def process( if beam_idx == self.group_size: break - new_completed, new_scores, new_tokens, new_indices = self.step_sentence_constraint( + new_scores, new_tokens, new_indices = self.step_sentence_constraint( batch_idx, input_ids, scores_for_all_vocab, - self.constraints_completed[batch_idx], next_beam_scores[batch_idx], next_beam_tokens[batch_idx], next_beam_indices[batch_idx], @@ -540,7 +540,6 @@ def process( next_beam_scores[batch_idx] = new_scores next_beam_tokens[batch_idx] = new_tokens next_beam_indices[batch_idx] = new_indices - self.constraints_completed[batch_idx] = new_completed if beam_idx < self.group_size: @@ -566,7 +565,6 @@ def step_sentence_constraint( batch_idx: int, input_ids: torch.LongTensor, vocab_scores: torch.FloatTensor, - constraints_completed: List[bool], sent_beam_scores: torch.FloatTensor, sent_beam_tokens: torch.LongTensor, sent_beam_indices: torch.LongTensor, @@ -593,15 +591,14 @@ def step_sentence_constraint( advance_constraint_states = self.make_constraint_states(orig_len) sidx, eidx = batch_idx*orig_len, (batch_idx+1) * orig_len - this_batch_input_ids = input_ids[sidx:eidx] this_batch_token_scores = vocab_scores[sidx:eidx] - full_hypotheses = torch.cat(( - this_batch_input_ids[sent_beam_indices], + input_ids[sent_beam_indices], sent_beam_tokens.unsqueeze(-1)), dim=-1 ) + # need to make new hypothesis that advance the constraints track_new = {"new_seqs": [], "new_states": [], "new_indices": [], "new_tokens": [], "new_scores": []} for seq_idx, pre_seq in enumerate(this_batch_input_ids): @@ -615,12 +612,12 @@ def step_sentence_constraint( topk_state = topk_contraint_states[seq_idx] topk_state.reset(full_hypotheses[seq_idx]) - if not constraints_completed[seq_idx]: - advance_state = advance_constraint_states[seq_idx] - advance_state.reset(pre_seq) - # the curerent `pre_seq` does not yet satisfy our list of constraints. + advance_state = advance_constraint_states[seq_idx] + advance_state.reset(pre_seq) + + if not advance_state.completed: advance_tokens = advance_state.advance() - for advance_token in advance_tokens: + for advance_token in advance_tokens.to(device): # since adding each `advance_token` leads to a different hypothesis, create new state instance. new_state = advance_state.copy(stateful=True) new_state.add(advance_token) @@ -681,8 +678,8 @@ def step_sentence_constraint( for one in all_states ]).to(device) - augmented_zipped = all_banks * 100 + all_scores - indices = augmented_zipped.sort(descending=True).indices + zipped = all_banks * 100 + all_scores + indices = zipped.sort(descending=True).indices sorted_banks = all_banks[indices] ''' Then we end up with @@ -700,14 +697,13 @@ def step_sentence_constraint( increments.append(counter) rearrangers = torch.tensor(np.argsort(increments, kind="mergesort")) - indices = indices[rearrangers] + indices = indices[rearrangers][:orig_len] - sent_beam_scores = all_scores[indices][:orig_len] - sent_beam_tokens = all_tokens[indices][:orig_len] - sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices][:orig_len] - constraints_completed = [all_states[idx].completed for idx in indices[:orig_len]] + sent_beam_scores = all_scores[indices] + sent_beam_tokens = all_tokens[indices] + sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices] - return constraints_completed, sent_beam_scores, sent_beam_tokens, sent_beam_indices + return sent_beam_scores, sent_beam_tokens, sent_beam_indices def finalize( self, @@ -732,7 +728,10 @@ def finalize( batch_beam_idx = batch_idx * self.num_beams + beam_id final_score = final_beam_scores[batch_beam_idx].item() final_tokens = input_ids[batch_beam_idx] - beam_hyp.add(final_tokens, final_score) + + completes_constraint = self.check_completes_constraints(final_tokens) + if completes_constraint: + beam_hyp.add(final_tokens, final_score) # select the best hypotheses sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index b2cb5593b57a..c3629aff8638 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -939,6 +939,9 @@ def generate( Custom stopping criteria that complement the default stopping criteria built from arguments and a model's config. If a stopping criteria is passed that is already created with the arguments or a model's config an error is thrown. This feature is intended for advanced users. + constraints (`List[Constraint]`, *optional*): + Custom constraints that can be added to the generation to ensure that the output will contain the use + of certain tokens as defined by `Constraint` objects, in the most sensible way possible. output_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more details. @@ -960,7 +963,6 @@ def generate( crash. Note that using `remove_invalid_values` can slow down generation. synced_gpus (`bool`, *optional*, defaults to `False`): Whether to continue running the while loop until max_length (needed for ZeRO stage 3) - model_kwargs: Additional model specific kwargs will be forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs @@ -2845,13 +2847,11 @@ def constrained_beam_search( synced_gpus: Optional[bool] = None, **model_kwargs, ) -> Union[BeamSearchOutput, torch.LongTensor]: - from transformers import GPT2Tokenizer - tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + r""" Generates sequences for models with a language modeling head using beam search decoding. Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. constrained_beam_scorer (`ConstrainedBeamScorer`): @@ -2864,6 +2864,10 @@ def constrained_beam_search( stopping_criteria (`StoppingCriteriaList`, *optional*): An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] used to tell if the generation loop should stop. + logits_warper (`LogitsProcessorList`, *optional*): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used + to warp the prediction score distribution of the language modeling head applied before multinomial + sampling at each generation step. max_length (`int`, *optional*, defaults to 20): **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated tokens. The maximum length of the sequence to be generated. @@ -2929,7 +2933,7 @@ def constrained_beam_search( ... } >>> constraints = [ - ... PhrasalConstraint(tokenizer.encode("required phrase")) + ... PhrasalConstraint(tokenizer.encode("required phrase")[0]) ... ] @@ -2948,7 +2952,7 @@ def constrained_beam_search( ... ] ... ) - >>> outputs = model.constrained_beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) + >>> outputs = model.constrained_beam_search(input_ids, beam_scorer, constraints=constraints, logits_processor=logits_processor, **model_kwargs) >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) ```""" @@ -3001,7 +3005,6 @@ def constrained_beam_search( beam_scores[:, 1:] = -1e9 beam_scores = beam_scores.view((batch_size * num_beams,)) - this_peer_finished = False # used by synced_gpus only while True: @@ -3030,6 +3033,10 @@ def constrained_beam_search( continue # don't waste resources running the code we don't need next_token_logits = outputs.logits[:, -1, :] + + # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id` + # cannot be generated both before and after the `nn.functional.log_softmax` operation. + next_token_logits = outputs.logits[:, -1, :] # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id` # cannot be generated both before and after the `nn.functional.log_softmax` operation. next_token_logits = self.adjust_logits_during_generation(next_token_logits, cur_len=cur_len) @@ -3037,12 +3044,11 @@ def constrained_beam_search( next_token_logits, dim=-1 ) # (batch_size * num_beams, vocab_size) + next_token_scores_processed = logits_processor(input_ids, next_token_scores) - next_token_scores = logits_processor(input_ids, next_token_scores) - - scores_for_all_vocab = next_token_scores.clone() + scores_for_all_vocab = next_token_scores_processed.clone() - next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores) + next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores) # Store scores, attentions and hidden_states when required @@ -3090,7 +3096,6 @@ def constrained_beam_search( beam_idx = beam_outputs["next_beam_indices"] input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder ) diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index 216284915e93..139c48ff847e 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -25,7 +25,8 @@ if is_torch_available(): import torch - from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer + from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer, ConstrainedBeamSearchScorer + from transformers.generation_beam_constraints import Constraint, PhrasalConstraint class BeamSearchTester: @@ -232,6 +233,282 @@ def check_beam_scores_finalize(self, input_ids, next_tokens, next_indices, next_ self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) +class ConstrainedBeamSearchTester: + def __init__( + self, + parent, + constraints=None, + batch_size=3, + sequence_length=10, + vocab_size=99, + pad_token_id=0, + max_length=20, + num_beams=4, + length_penalty=2.0, + do_early_stopping=True, + num_beam_hyps_to_keep=2, + ): + self.parent = parent + self.batch_size = batch_size + self.sequence_length = sequence_length + self.vocab_size = vocab_size + self.pad_token_id = pad_token_id + self.max_length = max_length + self.num_beams = num_beams + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + + if constraints is None: + force_tokens = torch.randint(10, 50, (1, 2)).type(torch.LongTensor)[0] + constraints = [ + PhrasalConstraint(force_tokens), + ] + self.constraints = constraints + # cannot be randomely generated + self.eos_token_id = vocab_size + 1 + + def prepare_constrained_beam_scorer(self, **kwargs): + return ConstrainedBeamSearchScorer( + constraints=kwargs.get("constraints", self.constraints), + batch_size=kwargs.get("batch_size", self.batch_size), + num_beams=kwargs.get("num_beams", self.num_beams), + device=torch_device, + length_penalty=kwargs.get("length_penalty", self.length_penalty), + do_early_stopping=kwargs.get("do_early_stopping", self.do_early_stopping), + num_beam_hyps_to_keep=kwargs.get("num_beam_hyps_to_keep", self.num_beam_hyps_to_keep), + ) + + def prepare_inputs(self): + input_ids = ids_tensor((self.batch_size * self.num_beams, self.sequence_length), self.vocab_size) + next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) + next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) + next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) + scores_for_all_vocab, _ = (-floats_tensor((self.batch_size * self.num_beams, self.vocab_size)).to(torch_device)).sort(descending=True) + return (input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab) + + def check_beam_hypotheses(self, input_ids, *args): + # check that correct number of beam hypotheses is set in beam scorer + constrained_beam_scorer = self.prepare_constrained_beam_scorer(do_early_stopping=True) + beam_hyp = constrained_beam_scorer._beam_hyps[0] + + self.parent.assertEqual(len(constrained_beam_scorer._beam_hyps), self.batch_size) + + # check correct type + self.parent.assertTrue(isinstance(beam_hyp, BeamHypotheses)) + + # check that num_beams is correctly set + self.parent.assertEqual(beam_hyp.num_beams, self.num_beams) + + # check for early stopping deactivated + for beam_idx in range(self.num_beams): + beam_hyp.add(input_ids[beam_idx], -10.0) + + # if early stopping True -> score does not matter + self.parent.assertTrue(beam_hyp.is_done(-10.0, 5)) + + # re-init + constrained_beam_scorer = self.prepare_constrained_beam_scorer(do_early_stopping=False) + beam_hyp = constrained_beam_scorer._beam_hyps[0] + + # add `num_beams + 1` beams to change `worst_score` + for beam_idx in range(self.num_beams + 1): + beam_hyp.add(input_ids[beam_idx], -10.0 + float(beam_idx)) + + # -10.0 is removed => -9.0 is worst score + self.parent.assertAlmostEqual(beam_hyp.worst_score, -9.0 / (self.sequence_length ** beam_hyp.length_penalty)) + + # -5.0 is better than worst score => should not be finished + self.parent.assertFalse(beam_hyp.is_done(-5.0, self.sequence_length)) + + # -20.0 is worse than worst score => should be finished + self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) + + def check_constrained_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab): + # check too many eos tokens + constrained_beam_scorer = self.prepare_constrained_beam_scorer() + fulfilling_sequence = torch.stack( + [constraint.token_ids for constraint in self.constraints] + ).flatten() + fulfill_len = fulfilling_sequence.size(0) + input_ids[:, :fulfill_len] = fulfilling_sequence + + tokens = next_tokens.clone() + tokens[0, :] = self.eos_token_id + + with self.parent.assertRaises(ValueError): + constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id) + + # check all batches are done + constrained_beam_scorer = self.prepare_constrained_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, :self.num_beams] = self.eos_token_id + constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id) + # beam scorer should be done + self.parent.assertTrue(constrained_beam_scorer.is_done) + + # check + constrained_beam_scorer = self.prepare_constrained_beam_scorer() + + tokens = next_tokens.clone() + tokens[:, 1] = self.eos_token_id + beam_outputs = constrained_beam_scorer.process( + input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + + def cut_expected_tensor(tensor): + return torch.cat([tensor[:, :1], tensor[:, 2 : self.num_beams + 1]], dim=1).flatten() + + # check all outptus + # cut out id of eos token and take best `num_beams` outputs + expected_output_tokens = cut_expected_tensor(tokens) + expected_output_scores = cut_expected_tensor(next_scores) + + # add num_beams * batch_idx + expected_output_indices = ( + cut_expected_tensor(next_indices) + + (torch.arange(self.num_beams * self.batch_size, device=torch_device) // self.num_beams) * self.num_beams + ) + + self.parent.assertListEqual(expected_output_tokens.tolist(), output_tokens.tolist()) + self.parent.assertListEqual(expected_output_indices.tolist(), output_indices.tolist()) + self.parent.assertTrue(torch.allclose(expected_output_scores, output_scores, atol=1e-3)) + + # make sure ids of eos token are correctly saved in beam_hyps of beam scorer + for batch_idx in range(self.batch_size): + correct_idx = batch_idx * self.num_beams + next_indices[batch_idx, 1] + self.parent.assertListEqual( + input_ids[correct_idx].tolist(), constrained_beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() + ) + + def check_constrained_beam_scorer_finalize(self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab): + # max_length should be only one more than current input_ids to check that eos is correctly appended + max_length = self.sequence_length + 1 + + # for testing finalize, we do want to have fulfilled constraints + fulfilling_sequence = torch.stack( + [constraint.token_ids for constraint in self.constraints] + ).flatten() + fulfill_len = fulfilling_sequence.size(0) + input_ids[:, :fulfill_len] = fulfilling_sequence + + constrained_beam_scorer = self.prepare_constrained_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) + + constraints = constrained_beam_scorer.constraints + # update beams and append to input_ids + tokens = next_tokens.clone() + # first batch, first output has to finish with eos token id since scores are correctly sorted + tokens[0, 0] = self.eos_token_id + # make sure corresponding score is as good as possible to surely be picked first + next_scores[0, 0] = 0.0 + + # # because constrainted beam search can't possibly fulfill the constraints in one pass + # # especially if constraints involves more than one token, so we repeat this several times. + # # doesn't *really* matter that we're not adjusting the scores & tokens TBH. + # repeat = 10 + # for _ in range(repeat): + beam_outputs = constrained_beam_scorer.process( + input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id + ) + output_scores = beam_outputs["next_beam_scores"] + output_tokens = beam_outputs["next_beam_tokens"] + output_indices = beam_outputs["next_beam_indices"] + input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) + + + + # finalize + + print("input_ids", input_ids.size()) + print("output_scores", output_scores.size()) + print("self.sequence_length", self.sequence_length) + print("max_length", max_length) + sequence_output = constrained_beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length` + self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size]) + + # check sequence_scores + self.parent.assertFalse((sequence_scores > 0).any().item()) + + # first batch has to finish with eos_token + self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id) + + # other batches cannot finish with eos token + self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id) + self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id) + + # test that the constraint is indeed fulfilled + for output in sequences: + for constraint in constraints: + forced_token_ids = constraint.token_ids + self.parent.assertEqual(self._check_sequence_inside_sequence(output, forced_token_ids), True) + + # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned + + + # constrained_beam_scorer.num_beam_hyps_to_keep = self.num_beams + constrained_beam_scorer = self.prepare_constrained_beam_scorer( + num_beam_hyps_to_keep=self.num_beams, + length_penalty=1.0, + do_early_stopping=False + ) + + sequence_output = constrained_beam_scorer.finalize( + input_ids, + output_scores, + output_tokens, + output_indices, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + max_length=max_length, + ) + sequences = sequence_output["sequences"] + sequence_scores = sequence_output["sequence_scores"] + + self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) + self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) + + + def _check_sequence_inside_sequence( + self, tensor_1, tensor_2 + ): + # set to same device. we don't care what device. + tensor_1, tensor_2 = tensor_1.cpu(), tensor_2.cpu() + + in_order = tensor_1.size(0) <= tensor_2.size(0) + longer = tensor_2 if in_order else tensor_1 + shorter = tensor_1 if in_order else tensor_2 + + flag = False + chunk_size = shorter.size(0) + for chunk_idx in range(longer.size(0) - chunk_size + 1): + subseq = longer[chunk_idx : chunk_idx+chunk_size] + if torch.equal(subseq, shorter): + flag = True + break + + return flag + + + + @require_torch class BeamSearchTest(unittest.TestCase): def setUp(self): @@ -247,4 +524,22 @@ def test_beam_scorer_update(self): def test_beam_scorer_finalize(self): inputs = self.beam_search_tester.prepare_inputs() - self.beam_search_tester.check_beam_scores_finalize(*inputs) \ No newline at end of file + self.beam_search_tester.check_beam_scores_finalize(*inputs) + + +@require_torch +class ConstrainedBeamSearchTest(unittest.TestCase): + def setUp(self): + self.constrained_beam_search_tester = ConstrainedBeamSearchTester(self) + + def test_constrained_beam_hypotheses(self): + inputs = self.constrained_beam_search_tester.prepare_inputs() + self.constrained_beam_search_tester.check_beam_hypotheses(*inputs) + + def test_constrained_beam_scorer_update(self): + inputs = self.constrained_beam_search_tester.prepare_inputs() + self.constrained_beam_search_tester.check_constrained_beam_scorer_update(*inputs) + + def test_constrained_beam_scorer_finalize(self): + inputs = self.constrained_beam_search_tester.prepare_inputs() + self.constrained_beam_search_tester.check_constrained_beam_scorer_finalize(*inputs) \ No newline at end of file diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 86d66512c233..67879e637492 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -37,7 +37,7 @@ VisionEncoderDecoderModel, top_k_top_p_filtering, ) - from transformers.generation_beam_search import BeamSearchScorer + from transformers.generation_beam_search import BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation_logits_process import ( ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor, @@ -63,6 +63,8 @@ SampleDecoderOnlyOutput, SampleEncoderDecoderOutput, ) + from transformers.generation_beam_constraints import PhrasalConstraint + class GenerationTesterMixin: @@ -190,6 +192,26 @@ def _get_diverse_beam_scorer_and_kwargs(batch_size, max_length, num_return_seque ) return beam_kwargs, beam_scorer + + @staticmethod + def _get_constrained_beam_scorer_and_kwargs(batch_size, max_length, constraints, num_return_sequences=1): + beam_kwargs = { + "early_stopping": False, + "length_penalty": 2.0, + "num_beams": num_return_sequences * 4, + "num_return_sequences": num_return_sequences, + } + beam_scorer = ConstrainedBeamSearchScorer( + batch_size=batch_size, + constraints=constraints, + num_beams=beam_kwargs["num_beams"], + device=torch_device, + length_penalty=beam_kwargs["length_penalty"], + do_early_stopping=beam_kwargs["early_stopping"], + num_beam_hyps_to_keep=num_return_sequences, + ) + return beam_kwargs, beam_scorer + @staticmethod def _get_encoder_outputs( model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1 @@ -526,6 +548,69 @@ def _group_beam_search_generate( ) return output_generate, output_group_beam_search + def _constrained_beam_search_generate( + self, + model, + input_ids, + attention_mask, + max_length, + constrained_beam_scorer, + constraints, + beam_kwargs, + logits_processor, + logits_process_kwargs, + output_scores=False, + output_attentions=False, + output_hidden_states=False, + return_dict_in_generate=False, + ): + output_generate = model.generate( + input_ids, + attention_mask=attention_mask, + do_sample=False, + max_length=max_length, + output_scores=output_scores, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict_in_generate=return_dict_in_generate, + remove_invalid_values=True, + constraints=constraints, + **beam_kwargs, + **logits_process_kwargs, + ) + + # group_beam_search does not automatically interleave `batch_size` dim for `num_beams` + kwargs = {} + if model.config.is_encoder_decoder: + encoder_outputs, input_ids_clone, attention_mask_clone = self._get_encoder_outputs( + model, + input_ids, + attention_mask, + num_interleave=constrained_beam_scorer.num_beams, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + kwargs["encoder_outputs"] = encoder_outputs + input_ids_clone = input_ids_clone.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + else: + attention_mask_clone = attention_mask.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + input_ids_clone = input_ids.repeat_interleave(constrained_beam_scorer.num_beams, dim=0) + + with torch.no_grad(): + output_group_beam_search = model.constrained_beam_search( + input_ids_clone, + constrained_beam_scorer, + max_length=max_length, + attention_mask=attention_mask_clone, + logits_processor=logits_processor, + output_scores=output_scores, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict_in_generate=return_dict_in_generate, + **kwargs, + ) + return output_generate, output_group_beam_search + def test_greedy_generate(self): # check `generate()` and `greedy_search()` are equal for model_class in self.all_generative_model_classes: @@ -1085,6 +1170,153 @@ def test_group_beam_search_generate_dict_output(self): output, input_ids, model.config, num_return_sequences=num_return_sequences * beam_scorer.num_beams ) + def test_constrained_beam_search_generate(self): + for model_class in self.all_generative_model_classes: + config, input_ids, attention_mask, max_length = self._get_input_ids_and_config() + + # It is important set set the eos_token_id to None to ensure that no sequences + # shorter than `max_length` can be generated which could lead to flaky circle ci + # failures if the top `num_return_sequences` beams are all shorter than the longest beam + config.eos_token_id = None + config.forced_eos_token_id = None + + model = model_class(config).to(torch_device).eval() + max_length = 10 + + logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( + input_ids.shape[-1], + config.eos_token_id, + config.forced_bos_token_id, + config.forced_eos_token_id, + max_length, + ) + + # check `generate()` and `group_beam_search()` are equal + # Sample constraints + min_id = torch.min(input_ids) + 3 + max_id = torch.max(input_ids) + force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0] + constraints = [ + PhrasalConstraint(force_tokens), + ] + + beam_kwargs, beam_scorer = self._get_constrained_beam_scorer_and_kwargs( + input_ids.shape[0], max_length, constraints, num_return_sequences=1 + ) + output_generate, output_beam_search = self._constrained_beam_search_generate( + model=model, + input_ids=input_ids, + attention_mask=attention_mask, + max_length=max_length, + constrained_beam_scorer=beam_scorer, + constraints=constraints, + beam_kwargs=beam_kwargs, + logits_processor=logits_processor, + logits_process_kwargs=logits_process_kwargs, + ) + self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) + for generation_output in output_generate: + self._check_sequence_inside_sequence(force_tokens, generation_output) + + # check `generate()` and `constrained_beam_search()` are equal for `num_return_sequences` + # Sample constraints + force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0] + constraints = [ + PhrasalConstraint(force_tokens), + ] + + num_return_sequences = 2 + max_length = 10 + + beam_kwargs, beam_scorer = self._get_constrained_beam_scorer_and_kwargs( + input_ids.shape[0], max_length, constraints, num_return_sequences=num_return_sequences + ) + + output_generate, output_beam_search = self._constrained_beam_search_generate( + model=model, + input_ids=input_ids, + attention_mask=attention_mask, + max_length=max_length, + constrained_beam_scorer=beam_scorer, + constraints=constraints, + beam_kwargs=beam_kwargs, + logits_processor=logits_processor, + logits_process_kwargs=logits_process_kwargs, + ) + self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) + + for generation_output in output_generate: + self._check_sequence_inside_sequence(force_tokens, generation_output) + + def test_constrained_beam_search_generate_dict_output(self): + for model_class in self.all_generative_model_classes: + config, input_ids, attention_mask, max_length = self._get_input_ids_and_config() + + # disable cache + config.use_cache = False + + # It is important set set the eos_token_id to None to ensure that no sequences + # shorter than `max_length` can be generated which could lead to flaky circle ci + # failures if the top `num_return_sequences` beams are all shorter than the longest beam + config.eos_token_id = None + config.forced_eos_token_id = None + + model = model_class(config).to(torch_device).eval() + if model.config.is_encoder_decoder: + max_length = 4 + + logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( + input_ids.shape[-1], + config.eos_token_id, + config.forced_bos_token_id, + config.forced_eos_token_id, + max_length, + ) + + # Sample constraints + min_id = torch.min(input_ids) + 3 + max_id = torch.max(input_ids) + force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0] + constraints = [ + PhrasalConstraint(force_tokens), + ] + + beam_kwargs, beam_scorer = self._get_constrained_beam_scorer_and_kwargs( + input_ids.shape[0], max_length, constraints, num_return_sequences=1 + ) + output_generate, output_beam_search = self._constrained_beam_search_generate( + model=model, + input_ids=input_ids, + attention_mask=attention_mask, + max_length=max_length, + constrained_beam_scorer=beam_scorer, + constraints=constraints, + beam_kwargs=beam_kwargs, + logits_processor=logits_processor, + logits_process_kwargs=logits_process_kwargs, + output_scores=True, + output_hidden_states=True, + output_attentions=True, + return_dict_in_generate=True, + ) + + if model.config.is_encoder_decoder: + self.assertIsInstance(output_beam_search, BeamSearchEncoderDecoderOutput) + self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput) + else: + self.assertIsInstance(output_beam_search, BeamSearchDecoderOnlyOutput) + self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput) + + self.assertListEqual(output_generate.sequences.tolist(), output_beam_search.sequences.tolist()) + self.assertTrue( + torch.allclose(output_generate["sequences_scores"], output_beam_search["sequences_scores"], atol=1e-3) + ) + self.assertTrue(output_generate["sequences_scores"].shape == (output_generate["sequences"].shape[0],)) + self.assertTrue((output_generate["sequences_scores"] < 0).all().item()) + + for output in (output_beam_search, output_generate): + self._check_outputs(output, input_ids, model.config, num_return_sequences=beam_scorer.num_beams) + def test_generate_with_head_masking(self): """Test designed for encoder-decoder models to ensure the attention head masking is used.""" attention_names = ["encoder_attentions", "decoder_attentions", "cross_attentions"] @@ -1254,6 +1486,25 @@ def _check_encoder_hidden_states_for_generate(self, hidden_states, batch_size, c [encoder_expected_shape] * len(hidden_states), ) + def _check_sequence_inside_sequence( + self, tensor_1, tensor_2 + ): + # set to same device. we don't care what device. + tensor_1, tensor_2 = tensor_1.cpu(), tensor_2.cpu() + + in_order = tensor_1.size(0) <= tensor_2.size(0) + longer = tensor_2 if in_order else tensor_1 + shorter = tensor_1 if in_order else tensor_2 + + flag = False + chunk_size = shorter.size(0) + for chunk_idx in range(longer.size(0) - chunk_size + 1): + subseq = longer[chunk_idx : chunk_idx+chunk_size] + if torch.equal(subseq, shorter): + flag = True + break + + assert flag @require_torch class UtilsFunctionsTest(unittest.TestCase): @@ -1896,7 +2147,7 @@ def test_generate_encoder_outputs_attention_mask(self): encoder = model.get_encoder() - encoder_outputs = encoder(inp/gut_values) + encoder_outputs = encoder(input_values) output_sequences_no_mask = model.generate(encoder_outputs=encoder_outputs).cpu() output_sequences_with_mask = model.generate(encoder_outputs=encoder_outputs, attention_mask=attention_mask) From 10f0679c80156d5bbf972f25aa372884af5b96e2 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 06:19:27 +0000 Subject: [PATCH 13/34] shifted docs --- src/transformers/generation_beam_constraints.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 1b7b4afe6fc1..f8ebaf8f7a2c 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -7,10 +7,14 @@ import torch from .file_utils import add_start_docstrings -from .utils.logging import get_logger +PHRASAL_CONSTRAINT_DOCSTRING = r""" + [`Constraint`] enforcing that an ordered sequence of tokens is generated. -logger = get_logger(__name__) + Args: + token_ids (`torch.LongTensor`): + The sequence of tokens that must be generated by the output. +""" class Constraint(ABC): @@ -143,6 +147,7 @@ def __init__(self, token_id: Union[int, torch.LongTensor]): "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`.") self.token_id = token_id + self.token_ids = token_id # for compatibility reasons self.completed = False def advance(self): @@ -174,13 +179,7 @@ def copy(self, stateful=False): return constraint class PhrasalConstraint(Constraint): - r""" - [`Constraint`] enforcing that an ordered sequence of tokens is generated. - - Args: - token_ids (`torch.Tensor`): - The sequence of tokens that must be generated by the output. - """ + @add_start_docstrings(PHRASAL_CONSTRAINT_DOCSTRING) def __init__(self, token_ids: Union[List[int], torch.LongTensor]): super(Constraint, self).__init__() From db9e9641e00c0a49a8a7516930558a3e26bc9ab1 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 07:16:19 +0000 Subject: [PATCH 14/34] passing all tests realistically ran locally --- docs/source/internal/generation_utils.mdx | 21 +- src/transformers/__init__.py | 11 +- .../generation_beam_constraints.py | 226 ++++++++++-------- src/transformers/generation_beam_search.py | 162 +++++++------ src/transformers/generation_utils.py | 37 ++- src/transformers/utils/dummy_pt_objects.py | 35 +++ tests/test_generation_beam_search.py | 62 +++-- tests/test_generation_utils.py | 15 +- 8 files changed, 329 insertions(+), 240 deletions(-) diff --git a/docs/source/internal/generation_utils.mdx b/docs/source/internal/generation_utils.mdx index 88e5e9e31551..c157247b97e3 100644 --- a/docs/source/internal/generation_utils.mdx +++ b/docs/source/internal/generation_utils.mdx @@ -16,8 +16,9 @@ This page lists all the utility functions used by [`~generation_utils.Generation [`~generation_utils.GenerationMixin.greedy_search`], [`~generation_utils.GenerationMixin.sample`], [`~generation_utils.GenerationMixin.beam_search`], -[`~generation_utils.GenerationMixin.beam_sample`], and -[`~generation_utils.GenerationMixin.group_beam_search`]. +[`~generation_utils.GenerationMixin.beam_sample`], +[`~generation_utils.GenerationMixin.group_beam_search`], and +[`~generation_utils.GenerationMixin.constrained_beam_search`]. Most of those are only useful if you are studying the code of the generate methods in the library. @@ -190,6 +191,18 @@ A [`StoppingCriteria`] can be used to change when to stop generation (other than [[autodoc]] MaxTimeCriteria - __call__ +## Constraints + +A [`Constraint`] can be used to force the generation to include specific tokens or sequences in the output. + +[[autodoc]] Constraint + +[[autodoc]] TokenConstraint + +[[autodoc]] PhrasalConstraint + +[[autodoc]] ConstraintListState + ## BeamSearch [[autodoc]] BeamScorer @@ -200,6 +213,10 @@ A [`StoppingCriteria`] can be used to change when to stop generation (other than - process - finalize +[[autodoc]] ConstrainedBeamSearchScorer + - process + - finalize + ## Utilities [[autodoc]] top_k_top_p_filtering diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 4c48536e27d3..32efd5c8e3ea 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -608,7 +608,13 @@ "TextDatasetForNextSentencePrediction", ] _import_structure["deepspeed"] = [] - _import_structure["generation_beam_search"] = ["BeamScorer", "BeamSearchScorer"] + _import_structure["generation_beam_constraints"] = [ + "Constraint", + "ConstraintListState", + "PhrasalConstraint", + "TokenConstraint", + ] + _import_structure["generation_beam_search"] = ["BeamScorer", "BeamSearchScorer", "ConstrainedBeamSearchScorer"] _import_structure["generation_logits_process"] = [ "ForcedBOSTokenLogitsProcessor", "ForcedEOSTokenLogitsProcessor", @@ -2723,7 +2729,8 @@ TextDataset, TextDatasetForNextSentencePrediction, ) - from .generation_beam_search import BeamScorer, BeamSearchScorer + from .generation_beam_constraints import Constraint, ConstraintListState, PhrasalConstraint, TokenConstraint + from .generation_beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from .generation_logits_process import ( ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor, diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index f8ebaf8f7a2c..8876cdea2465 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -1,44 +1,51 @@ -from abc import ABC - -from itertools import chain -from collections import Counter -from typing import List, Optional, Set, Tuple, Union +from abc import ABC, abstractmethod +from typing import List, Optional, Union import torch from .file_utils import add_start_docstrings -PHRASAL_CONSTRAINT_DOCSTRING = r""" - [`Constraint`] enforcing that an ordered sequence of tokens is generated. + +TOKEN_CONSTRAINT_DOCSTRING = r""" + [`Constraint`] enforcing that a specific token is included in the output. Args: token_ids (`torch.LongTensor`): The sequence of tokens that must be generated by the output. """ +PHRASAL_CONSTRAINT_DOCSTRING = r""" + [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. + + Args: + token_id (`int`): + The id of the token that must be generated by the output. +""" + class Constraint(ABC): r"""Abstract base class for all constraints that can be applied during generation. It must define how the constraint can be satisfied. All classes that inherit Constraint must follow the requirement that - + ``` completed = False while(not completed): _, completed = constraint.update(constraint.advance()) ``` - - will always terminate (halt). + + will always terminate (halt). """ + def __init__(self): # test for the above condition self.test() def test(self): - ''' + """ Tests whether this constraint has been properly defined. - ''' + """ counter = 0 completed = False while not completed: @@ -52,20 +59,21 @@ def test(self): if counter > 10000: raise Exception("update() does not fulfill the constraint.") - assert self.remaining() == 0 - + assert self.remaining() == 0 + + @abstractmethod def advance(self): - ''' - When called, returns the token that would take this constraint - one step closer to being fulfilled. + """ + When called, returns the token that would take this constraint one step closer to being fulfilled. returns: token_ids(`torch.tensor`): Must be a tensor of a list of indexable tokens, not some integer. - ''' + """ raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) + @abstractmethod def does_advance(self, token_id: int): """ Reads in a token and returns whether it creates progress. @@ -74,14 +82,15 @@ def does_advance(self, token_id: int): f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) + @abstractmethod def update(self, token_id: int): """ - Reads in a token and returns booleans that indicate the progress made by it. - This function will update the state of this object unlikes `does_advance(self, token_id: int)`. + Reads in a token and returns booleans that indicate the progress made by it. This function will update the + state of this object unlikes `does_advance(self, token_id: int)`. - This isn't to test whether a certain token will advance the progress; it's to update its state - as if it has been generated. This becomes important if token_id != desired token - (refer to else statement in PhrasalConstraint) + This isn't to test whether a certain token will advance the progress; it's to update its state as if it has + been generated. This becomes important if token_id != desired token (refer to else statement in + PhrasalConstraint) Args: token_id(`int`): @@ -97,33 +106,36 @@ def update(self, token_id: int): raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) - + + @abstractmethod def reset(self): """ - Resets the state of this constraint to its initialization. - We would call this in cases where the fulfillment of a constraint is abrupted by an unwanted token. + Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of + a constraint is abrupted by an unwanted token. """ raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) + @abstractmethod def remaining(self): - ''' + """ Returns the number of remaining steps of `advance()` in order to complete this constraint. - ''' + """ raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) + @abstractmethod def copy(self, stateful=False): - ''' + """ Creates a new instance of this constraint. Args: stateful(`boolean`): Whether to not only copy the constraint for new instance, but also its state. Returns: constraint(`Constraint`): The same constraint as the one being called from. - ''' + """ raise NotImplementedError( f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." ) @@ -137,22 +149,28 @@ class TokenConstraint(Constraint): token_id (`int`): The token that must be generated by the output. """ - def __init__(self, token_id: Union[int, torch.LongTensor]): + + @add_start_docstrings(TOKEN_CONSTRAINT_DOCSTRING) + def __init__(self, token_id: int): super(Constraint, self).__init__() if not (isinstance(token_id, int) or isinstance(token_id, torch.LongTensor)) or token_id < 0: - raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integer, but is {token_id}") + raise ValueError( + f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integer, but is {token_id}" + ) else: if isinstance(token_id, torch.LongTensor) and token_id.size(0) > 1: - raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." - "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`.") + raise ValueError( + f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." + "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`." + ) self.token_id = token_id - self.token_ids = token_id # for compatibility reasons + self.token_ids = token_id # for compatibility reasons self.completed = False def advance(self): return self.token_id - + def does_advance(self, token_id: Union[int, torch.LongTensor]): return token_id == self.token_id @@ -162,22 +180,23 @@ def update(self, token_id: Union[int, torch.LongTensor]): if self.does_advance(token_id): self.completed = True - return True, True, True # stepped, completed, reset + return True, True, True # stepped, completed, reset else: return False, False, False - + def reset(self): self.completed = False def remaining(self): return 0 if self.completed else 1 - + def copy(self, stateful=False): constraint = TokenConstraint(self.token_id) if stateful: constraint.completed = self.completed return constraint + class PhrasalConstraint(Constraint): @add_start_docstrings(PHRASAL_CONSTRAINT_DOCSTRING) def __init__(self, token_ids: Union[List[int], torch.LongTensor]): @@ -186,21 +205,25 @@ def __init__(self, token_ids: Union[List[int], torch.LongTensor]): is_int_list = isinstance(token_ids, List) and isinstance(token_ids[0], int) is_long_tensor = isinstance(token_ids, torch.LongTensor) and len(token_ids.size()) == 1 if not (is_int_list or is_long_tensor) or torch.any(token_ids < 0): - raise ValueError(f"`token_ids` has to be a single list of positive integers or a `torch.LongTensor` but is {token_ids}") + raise ValueError( + f"`token_ids` has to be a single list of positive integers or a `torch.LongTensor` but is {token_ids}" + ) else: if (is_int_list or is_long_tensor) and token_ids.size(0) == 1: - raise ValueError(f"`token_ids` has to be list of positive integers or a `torch.LongTensor` but is {token_ids}" - "For single token constraints, refer to `TokenConstraint`.") + raise ValueError( + f"`token_ids` has to be list of positive integers or a `torch.LongTensor` but is {token_ids}" + "For single token constraints, refer to `TokenConstraint`." + ) self.token_ids = token_ids self.seqlen = self.token_ids.size(0) - self.fulfilled_idx = -1 # the index of the currently fulfilled step + self.fulfilled_idx = -1 # the index of the currently fulfilled step self.completed = False def advance(self): return self.token_ids[self.fulfilled_idx + 1] - + def does_advance(self, token_id: int): if self.completed: return False @@ -247,7 +270,7 @@ def __init__(self, constraints: List[Constraint]): self.constraints = constraints # max # of steps required to fulfill a given constraint - self.max_seqlen = max([c.seqlen for c in constraints]) + self.max_seqlen = max([c.seqlen for c in constraints]) self.n_constraints = len(constraints) self.completed = False @@ -257,48 +280,47 @@ def init_state(self): self.complete_constraints = [] self.inprogress_constraint = None self.pending_constraints = [constraint.copy(stateful=False) for constraint in self.constraints] - + def get_bank(self): add = 0 if self.inprogress_constraint: # extra points for having a constraint mid-fulfilled add += self.max_seqlen - self.inprogress_constraint.remaining() - return (len(self.complete_constraints)*self.max_seqlen) + add + return (len(self.complete_constraints) * self.max_seqlen) + add def advance(self): - '''The list of tokens to generate such that we can make progress. + """The list of tokens to generate such that we can make progress. By "list" we don't mean the list of token that will fully fulfill a constraint. - Given constraints c_i = {t_ij | j == # of tokens}, - If we're not in the middle of progressing through a specific constraint c_i, we return: + Given constraints c_i = {t_ij | j == # of tokens}, If we're not in the middle of progressing through a specific + constraint c_i, we return: [t_k1 for k in indices of unfulfilled constraints] If we are in the middle of a constraint, then we return: [t_ij], where i == index of the inprogress constraint, j == the next step for the constraint. - Though we don't care which constraint is fulfilled first, - if we are in the progress of fulfilling a constraint, that's the only one we'll return. - ''' + Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint, + that's the only one we'll return. + """ if self.inprogress_constraint is None: token_list = [] - for constraint in self.pending_constraints: # "pending" == "unfulfilled yet" + for constraint in self.pending_constraints: # "pending" == "unfulfilled yet" advance = constraint.advance() token_list.append(advance) else: token_list = [self.inprogress_constraint.advance()] - + if len(token_list) == 0: return None else: return torch.stack(token_list) def reset(self, token_ids: Optional[torch.LongTensor]): - ''' - token_ids: the tokens generated thus far to reset the state of the - progress through constraints. - ''' + """ + token_ids: the tokens generated thus far to reset the state of the progress through constraints. + """ self.init_state() if token_ids is not None and token_ids.size(0) > 0: @@ -308,14 +330,16 @@ def reset(self, token_ids: Optional[torch.LongTensor]): break return self - + def add(self, token_id: Union[int, torch.LongTensor]): complete, stepped = False, False if isinstance(token_id, torch.LongTensor): if (token_id.size(0)) > 1: - raise ValueError(f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." - "It must have length 1.") + raise ValueError( + f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." + "It must have length 1." + ) else: token_id = token_id[0] @@ -325,47 +349,49 @@ def add(self, token_id: Union[int, torch.LongTensor]): return complete, stepped if self.inprogress_constraint is not None: - ''' - In the middle of fulfilling a constraint. - If the `token_id` *does* makes an incremental progress to current job, simply update the state - ''' + """ + In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current + job, simply update the state + """ stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: - ''' + """ 1. If the next token breaks the progress, then we must restart. e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books". - But that doesn't mean we self.init_state(), since we only reset the state - for this particular constraint, not the full list of constraints. - ''' + But that doesn't mean we self.init_state(), since we only reset the state for this particular + constraint, not the full list of constraints. + """ self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False)) - self.inprogress_constraint = None + self.inprogress_constraint = None if complete: - ''' - 2. If the next token completes the constraint, move it to completed list, set - inprogress to None. If there are no pending constraints either, then this - full list of constraints is complete. - ''' + """ + 2. If the next token completes the constraint, move it to completed list, set + inprogress to None. If there are no pending constraints either, then this full list of constraints + is complete. + """ self.complete_constraints.append(self.inprogress_constraint) self.inprogress_constraint = None if len(self.pending_constraints) == 0: # we're done! self.completed = True - + else: - ''' - Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards - any of our list of constraints? - ''' + """ + Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list + of constraints? + """ for cidx, pending_constraint in enumerate(self.pending_constraints): if pending_constraint.does_advance(token_id): stepped, complete, reset = pending_constraint.update(token_id) if not stepped: - raise Exception("constraint.update(token_id) is not yielding incremental progress, " - "even though constraint.does_advance(token_id) is true.") + raise Exception( + "constraint.update(token_id) is not yielding incremental progress, " + "even though constraint.does_advance(token_id) is true." + ) if complete: self.complete_constraints.append(pending_constraint) @@ -373,38 +399,36 @@ def add(self, token_id: Union[int, torch.LongTensor]): if not complete and stepped: self.inprogress_constraint = pending_constraint - + if complete or stepped: - ''' + """ If we made any progress at all, then it's at least not a "pending constraint". - ''' - self.pending_constraints = self.pending_constraints[:cidx] + self.pending_constraints[cidx+1:] + """ + self.pending_constraints = ( + self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :] + ) if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: - ''' - If there's no longer any pending after this and no inprogress either, then we must be complete. - ''' - self.completed = True + """ + If there's no longer any pending after this and no inprogress either, then we must be + complete. + """ + self.completed = True - break # prevent accidentally stepping through multiple constraints with just one token. + break # prevent accidentally stepping through multiple constraints with just one token. return complete, stepped - + def copy(self, stateful=True): - new_state = ConstraintListState(self.constraints) # we actually never though self.constraints objects + new_state = ConstraintListState(self.constraints) # we actually never though self.constraints objects # throughout this process. So it's at initialization state. if stateful: new_state.complete_constraints = [ - constraint.copy(stateful=True) - for constraint in self.complete_constraints + constraint.copy(stateful=True) for constraint in self.complete_constraints ] if self.inprogress_constraint is not None: new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True) - new_state.pending_constraints = [ - constraint.copy() - for constraint in self.pending_constraints - ] + new_state.pending_constraints = [constraint.copy() for constraint in self.pending_constraints] return new_state - diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index f20616b54b8c..26c2b5d1d23a 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -18,11 +18,11 @@ from collections import UserDict from typing import List, Optional, Tuple -import torch import numpy as np +import torch -from .generation_beam_constraints import Constraint, ConstraintListState from .file_utils import add_start_docstrings +from .generation_beam_constraints import Constraint, ConstraintListState PROCESS_INPUTS_DOCSTRING = r""" @@ -85,6 +85,41 @@ """ +CONSTRAINT_PROCESS_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + next_scores (`torch.FloatTensor` of shape `(batch_size, 2 * num_beams)`): + Current scores of the top `2 * num_beams` non-finished beam hypotheses. + next_tokens (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + `input_ids` of the tokens corresponding to the top `2 * num_beams` non-finished beam hypotheses. + next_indices (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + Beam indices indicating to which beam hypothesis the `next_tokens` correspond. + scores_for_all_vocab (`torch.FloatTensor` of shape `(batch_size * num_beams, sequence_length)`): + The scores of all tokens in the vocabulary for each of the beam hypotheses. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`int`, *optional*): + The id of the *end-of-sequence* token. + + Return: + `UserDict`: A dictionary composed of the fields as defined above: + + - **next_beam_scores** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Updated scores of all + non-finished beams. + - **next_beam_tokens** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Next tokens to be added + to the non-finished beam_hypotheses. + - **next_beam_indices** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Beam indices + indicating to which beam the next tokens shall be added. + +""" + + class BeamScorer(ABC): """ Abstract base class for all beam scorers that are used for [`~PreTrainedModel.beam_search`] and @@ -365,8 +400,8 @@ class ConstrainedBeamSearchScorer(BeamScorer): num_beams (`int`): Number of beams for beam search. constraints (`List[Constraint]`): - A list of positive constraints represented as `Constraint` objects that must be fulfilled in - the generation output. For more information, the documentation of [`Constraint`] should be read. + A list of positive constraints represented as `Constraint` objects that must be fulfilled in the generation + output. For more information, the documentation of [`Constraint`] should be read. device (`torch.device`): Defines the device type (*e.g.*, `"cpu"` or `"cuda"`) on which this instance of `BeamSearchScorer` will be allocated. @@ -416,8 +451,6 @@ def __init__( ] self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device) - - if not isinstance(num_beams, int) or num_beams <= 1: raise ValueError( f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1, one should make use of `greedy_search` instead." @@ -441,19 +474,14 @@ def is_done(self) -> bool: return self._done.all() def make_constraint_states(self, n): - return [ - ConstraintListState([ - constraint.copy() - for constraint in self.constraints - ]) - for _ in range(n) - ] + return [ConstraintListState([constraint.copy() for constraint in self.constraints]) for _ in range(n)] def check_completes_constraints(self, sequence): new_state = self.make_constraint_states(1)[0] new_state = new_state.reset(sequence) return new_state.completed + @add_start_docstrings(CONSTRAINT_PROCESS_INPUTS_DOCSTRING) def process( self, input_ids: torch.LongTensor, @@ -496,7 +524,6 @@ def process( next_beam_indices[batch_idx, :] = 0 continue - # next tokens for this sentence. beam_idx = 0 for beam_token_rank, (next_token, next_score, next_index) in enumerate( @@ -505,7 +532,7 @@ def process( batch_beam_idx = batch_idx * self.group_size + next_index # add to generated hypotheses if end of sentence if (eos_token_id is not None) and (next_token.item() == eos_token_id): - + # if beam_token does not belong to top num_beams tokens, it should not be added is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size if is_beam_token_worse_than_top_num_beams: @@ -541,7 +568,6 @@ def process( next_beam_tokens[batch_idx] = new_tokens next_beam_indices[batch_idx] = new_indices - if beam_idx < self.group_size: raise ValueError( f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id: {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected." @@ -568,50 +594,47 @@ def step_sentence_constraint( sent_beam_scores: torch.FloatTensor, sent_beam_tokens: torch.LongTensor, sent_beam_indices: torch.LongTensor, - push_progress: bool = False + push_progress: bool = False, ): # from transformers import GPT2Tokenizer # tokenizer = GPT2Tokenizer.from_pretrained("gpt2") - ''' - sent_beam_tokens are the next {num_beams} number of tokens that are under consideration - for this beam (candidate next tokens) + """ + sent_beam_tokens are the next {num_beams} number of tokens that are under consideration for this beam + (candidate next tokens) 1. Adding "advance_tokens" - using ConstraintStateList.advance(), we propose new tokens to be added into this - "candidate list" that will advance us in fulfilling the constraints. + using ConstraintStateList.advance(), we propose new tokens to be added into this "candidate list" that will + advance us in fulfilling the constraints. 2. Selecting best candidates such that we end up with highest probable candidates that fulfill our constraints. - ''' + """ orig_len = sent_beam_indices.size(0) device = sent_beam_indices.get_device() # initialize states - topk_contraint_states = self.make_constraint_states(orig_len) + topk_contraint_states = self.make_constraint_states(orig_len) advance_constraint_states = self.make_constraint_states(orig_len) - sidx, eidx = batch_idx*orig_len, (batch_idx+1) * orig_len + sidx, eidx = batch_idx * orig_len, (batch_idx + 1) * orig_len this_batch_input_ids = input_ids[sidx:eidx] this_batch_token_scores = vocab_scores[sidx:eidx] - full_hypotheses = torch.cat(( - input_ids[sent_beam_indices], - sent_beam_tokens.unsqueeze(-1)), - dim=-1 - ) + full_hypotheses = torch.cat((input_ids[sent_beam_indices], sent_beam_tokens.unsqueeze(-1)), dim=-1) # need to make new hypothesis that advance the constraints track_new = {"new_seqs": [], "new_states": [], "new_indices": [], "new_tokens": [], "new_scores": []} for seq_idx, pre_seq in enumerate(this_batch_input_ids): - ''' + """ pre_seq = ith sequence generated before this step. - - input_ids -> (topk) generic beam search best model next tokens - -> (advance) constraints forcing the next token - either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of hypotheses. - ''' + + input_ids -> (topk) generic beam search best model next tokens + -> (advance) constraints forcing the next token + either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of + hypotheses. + """ topk_state = topk_contraint_states[seq_idx] - topk_state.reset(full_hypotheses[seq_idx]) - + topk_state.reset(full_hypotheses[seq_idx]) + advance_state = advance_constraint_states[seq_idx] advance_state.reset(pre_seq) @@ -621,7 +644,7 @@ def step_sentence_constraint( # since adding each `advance_token` leads to a different hypothesis, create new state instance. new_state = advance_state.copy(stateful=True) new_state.add(advance_token) - + advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).cpu().tolist() if advance_seq not in track_new["new_seqs"]: # prevent duplicates, which are basically bound to happen in this process. @@ -629,29 +652,29 @@ def step_sentence_constraint( track_new["new_indices"].append(seq_idx) track_new["new_tokens"].append(advance_token) track_new["new_scores"].append(this_batch_token_scores[seq_idx].take(advance_token)) - track_new["new_states"].append(new_state) + track_new["new_states"].append(new_state) elif push_progress: - ''' - Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences - that actually fulfill our constraints. For example, let constraints == ["loves pies"] and - - pre_seq_1 = "The child loves pies and" - pre_seq_2 = "The child plays in the playground and" - - Without this step, if `sent_beam_indices` is something like [1,1], then + """ + Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences that + actually fulfill our constraints. For example, let constraints == ["loves pies"] and + + pre_seq_1 = "The child loves pies and" pre_seq_2 = "The child plays in the playground and" + + Without this step, if `sent_beam_indices` is something like [1,1], then 1. `pre_seq_1` won't be added to the list of (topk) hypothesis since it's not in the indices and - 2. it won't be added to the list of (advance) hypothesis since it's completed already. - (this is the else part of `if constraints_completed[seq_idx]`) + 2. it won't be added to the list of (advance) hypothesis since it's completed already. (this is + the else part of `if constraints_completed[seq_idx]`) 3. it ends up simply getting removed from consideration. - - #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, especially - if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam search times, - since completed sequences keep getting removed after all this effort for constrained generation. - - Here, we basically take `pre_seq_1` and to "push" it into the considered list of - hypotheses, by simply appending the next likely token in the vocabulary and adding it to the list of hypotheses. - ''' - new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) # some next probable token + + #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, + especially if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam + search times, since completed sequences keep getting removed after all this effort for constrained + generation. + + Here, we basically take `pre_seq_1` and to "push" it into the considered list of hypotheses, by simply + appending the next likely token in the vocabulary and adding it to the list of hypotheses. + """ + new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) # some next probable token advance_seq = torch.cat((pre_seq, new_token.unsqueeze(0)), -1) advance_state = advance_constraint_states[seq_idx] @@ -664,27 +687,24 @@ def step_sentence_constraint( track_new["new_indices"].append(seq_idx) track_new["new_tokens"].append(new_token) track_new["new_scores"].append(new_score) - track_new["new_states"].append(advance_state) + track_new["new_states"].append(advance_state) if len(track_new["new_indices"]) > 0: new_indices = torch.tensor(track_new["new_indices"]).to(device) new_tokens = torch.stack(track_new["new_tokens"]).to(device) new_scores = torch.stack(track_new["new_scores"]).to(device) - + all_states = topk_contraint_states + track_new["new_states"] all_tokens = torch.cat((sent_beam_tokens, new_tokens), -1) all_scores = torch.cat((sent_beam_scores, new_scores), -1) - all_banks = torch.tensor([one.get_bank() - for one in all_states - ]).to(device) + all_banks = torch.tensor([one.get_bank() for one in all_states]).to(device) zipped = all_banks * 100 + all_scores indices = zipped.sort(descending=True).indices - sorted_banks = all_banks[indices] - ''' - Then we end up with - {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} - ''' + sorted_banks = all_banks[indices] + """ + Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} + """ counter = -1 cur_bank = sorted_banks[0] increments = [] @@ -728,7 +748,7 @@ def finalize( batch_beam_idx = batch_idx * self.num_beams + beam_id final_score = final_beam_scores[batch_beam_idx].item() final_tokens = input_ids[batch_beam_idx] - + completes_constraint = self.check_completes_constraints(final_tokens) if completes_constraint: beam_hyp.add(final_tokens, final_score) @@ -772,7 +792,6 @@ def finalize( ) - class BeamHypotheses: def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool): """ @@ -818,4 +837,3 @@ def is_done(self, best_sum_logprobs: float, cur_len: int) -> bool: cur_score = best_sum_logprobs / cur_len ** self.length_penalty ret = self.worst_score >= cur_score return ret - diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index c3629aff8638..9d0e9b3f7b0d 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -24,6 +24,7 @@ from torch import nn from .file_utils import ModelOutput +from .generation_beam_constraints import Constraint from .generation_beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from .generation_logits_process import ( EncoderNoRepeatNGramLogitsProcessor, @@ -48,10 +49,6 @@ StoppingCriteriaList, validate_stopping_criteria, ) -from .generation_beam_constraints import ( - Constraint, - ConstraintListState -) from .utils import logging @@ -941,7 +938,7 @@ def generate( model's config an error is thrown. This feature is intended for advanced users. constraints (`List[Constraint]`, *optional*): Custom constraints that can be added to the generation to ensure that the output will contain the use - of certain tokens as defined by `Constraint` objects, in the most sensible way possible. + of certain tokens as defined by `Constraint` objects, in the most sensible way possible. output_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more details. @@ -1138,7 +1135,9 @@ def generate( is_greedy_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is False and constraints is None is_sample_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is True and constraints is None is_beam_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is False and constraints is None - is_beam_sample_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is True and constraints is None + is_beam_sample_gen_mode = ( + (num_beams > 1) and (num_beam_groups == 1) and do_sample is True and constraints is None + ) is_group_beam_gen_mode = (num_beams > 1) and (num_beam_groups > 1) and constraints is None if num_beam_groups > num_beams: @@ -1369,7 +1368,6 @@ def generate( **model_kwargs, ) - def greedy_search( self, input_ids: torch.LongTensor, @@ -2856,7 +2854,7 @@ def constrained_beam_search( The sequence used as a prompt for the generation. constrained_beam_scorer (`ConstrainedBeamScorer`): An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and - sorted during generation, while satisfying a list of positive constraints. For more information, the + sorted during generation, while satisfying a list of positive constraints. For more information, the documentation of [`ConstrainedBeamScorer`] should be read. logits_processor (`LogitsProcessorList`, *optional*): An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] @@ -2908,7 +2906,7 @@ def constrained_beam_search( ... LogitsProcessorList, ... MinLengthLogitsProcessor, ... ConstrainedBeamSearchScorer, - ... PhrasalConstraint + ... PhrasalConstraint, ... ) >>> import torch @@ -2932,17 +2930,12 @@ def constrained_beam_search( ... ) ... } - >>> constraints = [ - ... PhrasalConstraint(tokenizer.encode("required phrase")[0]) - ... ] + >>> constraints = [PhrasalConstraint(tokenizer.encode("required phrase")[0])] >>> # instantiate beam scorer >>> beam_scorer = ConstrainedBeamSearchScorer( - ... batch_size=1, - ... num_beams=num_beams, - ... device=model.device, - ... constraints=constraints + ... batch_size=1, num_beams=num_beams, device=model.device, constraints=constraints ... ) >>> # instantiate logits processors @@ -2952,11 +2945,13 @@ def constrained_beam_search( ... ] ... ) - >>> outputs = model.constrained_beam_search(input_ids, beam_scorer, constraints=constraints, logits_processor=logits_processor, **model_kwargs) + >>> outputs = model.constrained_beam_search( + ... input_ids, beam_scorer, constraints=constraints, logits_processor=logits_processor, **model_kwargs + ... ) >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) ```""" - # init values + # init values logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() if max_length is not None: @@ -3018,7 +3013,6 @@ def constrained_beam_search( if this_peer_finished_flag.item() == 0.0: break - model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) outputs = self( @@ -3050,7 +3044,6 @@ def constrained_beam_search( next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores) - # Store scores, attentions and hidden_states when required if return_dict_in_generate: if output_scores: @@ -3073,7 +3066,6 @@ def constrained_beam_search( vocab_size = next_token_scores.shape[-1] next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) - next_token_scores, next_tokens = torch.topk( next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True ) @@ -3145,7 +3137,8 @@ def constrained_beam_search( ) else: return sequence_outputs["sequences"] - + + def top_k_top_p_filtering( logits: torch.FloatTensor, top_k: int = 0, diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 21e94075a8da..1e6f0b4eb886 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -80,6 +80,34 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) +class Constraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConstraintListState(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhrasalConstraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TokenConstraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class BeamScorer(metaclass=DummyObject): _backends = ["torch"] @@ -94,6 +122,13 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) +class ConstrainedBeamSearchScorer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class ForcedBOSTokenLogitsProcessor(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index 139c48ff847e..d6885d5569eb 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -25,8 +25,8 @@ if is_torch_available(): import torch + from transformers.generation_beam_constraints import PhrasalConstraint from transformers.generation_beam_search import BeamHypotheses, BeamSearchScorer, ConstrainedBeamSearchScorer - from transformers.generation_beam_constraints import Constraint, PhrasalConstraint class BeamSearchTester: @@ -284,7 +284,9 @@ def prepare_inputs(self): next_tokens = ids_tensor((self.batch_size, 2 * self.num_beams), self.vocab_size).to(torch_device) next_indices = ids_tensor((self.batch_size, 2 * self.num_beams), self.num_beams).to(torch_device) next_scores, _ = (-floats_tensor((self.batch_size, 2 * self.num_beams)).to(torch_device)).sort(descending=True) - scores_for_all_vocab, _ = (-floats_tensor((self.batch_size * self.num_beams, self.vocab_size)).to(torch_device)).sort(descending=True) + scores_for_all_vocab, _ = ( + -floats_tensor((self.batch_size * self.num_beams, self.vocab_size)).to(torch_device) + ).sort(descending=True) return (input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab) def check_beam_hypotheses(self, input_ids, *args): @@ -324,12 +326,12 @@ def check_beam_hypotheses(self, input_ids, *args): # -20.0 is worse than worst score => should be finished self.parent.assertTrue(beam_hyp.is_done(-20.0, self.sequence_length)) - def check_constrained_beam_scorer_update(self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab): + def check_constrained_beam_scorer_update( + self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab + ): # check too many eos tokens constrained_beam_scorer = self.prepare_constrained_beam_scorer() - fulfilling_sequence = torch.stack( - [constraint.token_ids for constraint in self.constraints] - ).flatten() + fulfilling_sequence = torch.stack([constraint.token_ids for constraint in self.constraints]).flatten() fulfill_len = fulfilling_sequence.size(0) input_ids[:, :fulfill_len] = fulfilling_sequence @@ -337,14 +339,18 @@ def check_constrained_beam_scorer_update(self, input_ids, next_tokens, next_indi tokens[0, :] = self.eos_token_id with self.parent.assertRaises(ValueError): - constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id) + constrained_beam_scorer.process( + input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id + ) # check all batches are done constrained_beam_scorer = self.prepare_constrained_beam_scorer() tokens = next_tokens.clone() - tokens[:, :self.num_beams] = self.eos_token_id - constrained_beam_scorer.process(input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id) + tokens[:, : self.num_beams] = self.eos_token_id + constrained_beam_scorer.process( + input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id + ) # beam scorer should be done self.parent.assertTrue(constrained_beam_scorer.is_done) @@ -385,18 +391,20 @@ def cut_expected_tensor(tensor): input_ids[correct_idx].tolist(), constrained_beam_scorer._beam_hyps[batch_idx].beams[0][-1].tolist() ) - def check_constrained_beam_scorer_finalize(self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab): + def check_constrained_beam_scorer_finalize( + self, input_ids, next_tokens, next_indices, next_scores, scores_for_all_vocab + ): # max_length should be only one more than current input_ids to check that eos is correctly appended max_length = self.sequence_length + 1 # for testing finalize, we do want to have fulfilled constraints - fulfilling_sequence = torch.stack( - [constraint.token_ids for constraint in self.constraints] - ).flatten() + fulfilling_sequence = torch.stack([constraint.token_ids for constraint in self.constraints]).flatten() fulfill_len = fulfilling_sequence.size(0) input_ids[:, :fulfill_len] = fulfilling_sequence - constrained_beam_scorer = self.prepare_constrained_beam_scorer(num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False) + constrained_beam_scorer = self.prepare_constrained_beam_scorer( + num_beam_hyps_to_keep=1, length_penalty=1.0, do_early_stopping=False + ) constraints = constrained_beam_scorer.constraints # update beams and append to input_ids @@ -419,8 +427,6 @@ def check_constrained_beam_scorer_finalize(self, input_ids, next_tokens, next_in output_indices = beam_outputs["next_beam_indices"] input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) - - # finalize print("input_ids", input_ids.size()) @@ -457,18 +463,15 @@ def check_constrained_beam_scorer_finalize(self, input_ids, next_tokens, next_in # test that the constraint is indeed fulfilled for output in sequences: for constraint in constraints: - forced_token_ids = constraint.token_ids + forced_token_ids = constraint.token_ids self.parent.assertEqual(self._check_sequence_inside_sequence(output, forced_token_ids), True) # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned - # constrained_beam_scorer.num_beam_hyps_to_keep = self.num_beams constrained_beam_scorer = self.prepare_constrained_beam_scorer( - num_beam_hyps_to_keep=self.num_beams, - length_penalty=1.0, - do_early_stopping=False - ) + num_beam_hyps_to_keep=self.num_beams, length_penalty=1.0, do_early_stopping=False + ) sequence_output = constrained_beam_scorer.finalize( input_ids, @@ -485,28 +488,23 @@ def check_constrained_beam_scorer_finalize(self, input_ids, next_tokens, next_in self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length]) self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size]) - - def _check_sequence_inside_sequence( - self, tensor_1, tensor_2 - ): + def _check_sequence_inside_sequence(self, tensor_1, tensor_2): # set to same device. we don't care what device. tensor_1, tensor_2 = tensor_1.cpu(), tensor_2.cpu() in_order = tensor_1.size(0) <= tensor_2.size(0) longer = tensor_2 if in_order else tensor_1 shorter = tensor_1 if in_order else tensor_2 - + flag = False chunk_size = shorter.size(0) for chunk_idx in range(longer.size(0) - chunk_size + 1): - subseq = longer[chunk_idx : chunk_idx+chunk_size] + subseq = longer[chunk_idx : chunk_idx + chunk_size] if torch.equal(subseq, shorter): flag = True break - - return flag - + return flag @require_torch @@ -542,4 +540,4 @@ def test_constrained_beam_scorer_update(self): def test_constrained_beam_scorer_finalize(self): inputs = self.constrained_beam_search_tester.prepare_inputs() - self.constrained_beam_search_tester.check_constrained_beam_scorer_finalize(*inputs) \ No newline at end of file + self.constrained_beam_search_tester.check_constrained_beam_scorer_finalize(*inputs) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 67879e637492..f751a31e69c1 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -37,6 +37,7 @@ VisionEncoderDecoderModel, top_k_top_p_filtering, ) + from transformers.generation_beam_constraints import PhrasalConstraint from transformers.generation_beam_search import BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation_logits_process import ( ForcedBOSTokenLogitsProcessor, @@ -63,8 +64,6 @@ SampleDecoderOnlyOutput, SampleEncoderDecoderOutput, ) - from transformers.generation_beam_constraints import PhrasalConstraint - class GenerationTesterMixin: @@ -192,7 +191,6 @@ def _get_diverse_beam_scorer_and_kwargs(batch_size, max_length, num_return_seque ) return beam_kwargs, beam_scorer - @staticmethod def _get_constrained_beam_scorer_and_kwargs(batch_size, max_length, constraints, num_return_sequences=1): beam_kwargs = { @@ -1486,26 +1484,25 @@ def _check_encoder_hidden_states_for_generate(self, hidden_states, batch_size, c [encoder_expected_shape] * len(hidden_states), ) - def _check_sequence_inside_sequence( - self, tensor_1, tensor_2 - ): + def _check_sequence_inside_sequence(self, tensor_1, tensor_2): # set to same device. we don't care what device. tensor_1, tensor_2 = tensor_1.cpu(), tensor_2.cpu() in_order = tensor_1.size(0) <= tensor_2.size(0) longer = tensor_2 if in_order else tensor_1 shorter = tensor_1 if in_order else tensor_2 - + flag = False chunk_size = shorter.size(0) for chunk_idx in range(longer.size(0) - chunk_size + 1): - subseq = longer[chunk_idx : chunk_idx+chunk_size] + subseq = longer[chunk_idx : chunk_idx + chunk_size] if torch.equal(subseq, shorter): flag = True break - + assert flag + @require_torch class UtilsFunctionsTest(unittest.TestCase): From 8242282a41bd05b3a54b4052967004615d2e6a4c Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 07:37:21 +0000 Subject: [PATCH 15/34] removing accidentally included print statements --- tests/test_generation_beam_search.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index d6885d5569eb..6b7b84542ebf 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -428,11 +428,6 @@ def check_constrained_beam_scorer_finalize( input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1) # finalize - - print("input_ids", input_ids.size()) - print("output_scores", output_scores.size()) - print("self.sequence_length", self.sequence_length) - print("max_length", max_length) sequence_output = constrained_beam_scorer.finalize( input_ids, output_scores, From 88945d514bfcfab667f07050e60bc693f75a8816 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 08:09:03 +0000 Subject: [PATCH 16/34] fixed source of error in initial PR test --- src/transformers/generation_beam_constraints.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 8876cdea2465..897f8dd5ffa2 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -334,14 +334,14 @@ def reset(self, token_ids: Optional[torch.LongTensor]): def add(self, token_id: Union[int, torch.LongTensor]): complete, stepped = False, False - if isinstance(token_id, torch.LongTensor): - if (token_id.size(0)) > 1: - raise ValueError( - f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." - "It must have length 1." - ) - else: - token_id = token_id[0] + # if isinstance(token_id, torch.LongTensor): + # if (token_id.size(0)) > 1: + # raise ValueError( + # f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." + # "It must have length 1." + # ) + # else: + # token_id = token_id[0] if self.completed: complete = True From 73f3acd25f5fc4c48f2f0069a42e954bb8afe3a2 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 10:20:00 +0000 Subject: [PATCH 17/34] fixing the get_device() vs device trap --- src/transformers/generation_beam_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 26c2b5d1d23a..e32d7d63f44f 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -610,7 +610,7 @@ def step_sentence_constraint( that fulfill our constraints. """ orig_len = sent_beam_indices.size(0) - device = sent_beam_indices.get_device() + device = sent_beam_indices.device # initialize states topk_contraint_states = self.make_constraint_states(orig_len) From 42efa234e09200a2c9fcf2696933330b6f4c7a95 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 10:25:34 +0000 Subject: [PATCH 18/34] fixed documentation docstrings about constrained_beam_search --- src/transformers/generation_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 9d0e9b3f7b0d..31440b50db60 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -2852,10 +2852,10 @@ def constrained_beam_search( Parameters: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. - constrained_beam_scorer (`ConstrainedBeamScorer`): - An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and + constrained_beam_scorer (`ConstrainedBeamSearchScorer`): + A derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and sorted during generation, while satisfying a list of positive constraints. For more information, the - documentation of [`ConstrainedBeamScorer`] should be read. + documentation of [`ConstrainedBeamSearchScorer`] should be read. logits_processor (`LogitsProcessorList`, *optional*): An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] used to modify the prediction scores of the language modeling head applied at each generation step. From fb2195aa1b101a543c70bcc69ce9e513206d45b2 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Mon, 31 Jan 2022 14:11:59 +0000 Subject: [PATCH 19/34] fixed tests having failing for Speech2TextModel's floating point inputs --- tests/test_generation_utils.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index f751a31e69c1..274ed554a7fa 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -802,6 +802,11 @@ def test_beam_search_generate(self): logits_process_kwargs=logits_process_kwargs, logits_processor=logits_processor, ) + + print("input_ids", input_ids) + print("output_generate", output_generate) + assert False + self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) # check `generate()` and `beam_search()` are equal for `num_return_sequences` @@ -1179,7 +1184,7 @@ def test_constrained_beam_search_generate(self): config.forced_eos_token_id = None model = model_class(config).to(torch_device).eval() - max_length = 10 + max_length = 20 logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( input_ids.shape[-1], @@ -1189,10 +1194,16 @@ def test_constrained_beam_search_generate(self): max_length, ) - # check `generate()` and `group_beam_search()` are equal + # check `generate()` and `constrained_beam_search()` are equal # Sample constraints - min_id = torch.min(input_ids) + 3 - max_id = torch.max(input_ids) + if not input_ids.dtype == torch.float32: + min_id = torch.min(input_ids) + 3 + max_id = torch.max(input_ids) + else: + # otherwise this throws an error for Speech2TextModel since its inputs are floating points + min_id = 3 + max_id = 100 + force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0] constraints = [ PhrasalConstraint(force_tokens), @@ -1224,7 +1235,7 @@ def test_constrained_beam_search_generate(self): ] num_return_sequences = 2 - max_length = 10 + max_length = 20 beam_kwargs, beam_scorer = self._get_constrained_beam_scorer_and_kwargs( input_ids.shape[0], max_length, constraints, num_return_sequences=num_return_sequences @@ -1261,7 +1272,7 @@ def test_constrained_beam_search_generate_dict_output(self): model = model_class(config).to(torch_device).eval() if model.config.is_encoder_decoder: - max_length = 4 + max_length = 20 logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs( input_ids.shape[-1], @@ -1272,8 +1283,13 @@ def test_constrained_beam_search_generate_dict_output(self): ) # Sample constraints - min_id = torch.min(input_ids) + 3 - max_id = torch.max(input_ids) + if not input_ids.dtype == torch.float32: + min_id = torch.min(input_ids) + 3 + max_id = torch.max(input_ids) + else: + # otherwise this throws an error for Speech2TextModel since its inputs are floating points + min_id = 3 + max_id = 100 force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0] constraints = [ PhrasalConstraint(force_tokens), From f522031be6ce71fd39f8fb4c2612fbc7c231b578 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 31 Jan 2022 15:21:29 +0100 Subject: [PATCH 20/34] fix cuda long tensor --- src/transformers/generation_beam_constraints.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 897f8dd5ffa2..8840a810adea 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -153,12 +153,12 @@ class TokenConstraint(Constraint): @add_start_docstrings(TOKEN_CONSTRAINT_DOCSTRING) def __init__(self, token_id: int): super(Constraint, self).__init__() - if not (isinstance(token_id, int) or isinstance(token_id, torch.LongTensor)) or token_id < 0: + if not (isinstance(token_id, int) or isinstance(token_id, (torch.LongTensor, torch.cuda.LongTensor))) or token_id < 0: raise ValueError( f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integer, but is {token_id}" ) else: - if isinstance(token_id, torch.LongTensor) and token_id.size(0) > 1: + if isinstance(token_id, (torch.LongTensor, torch.cuda.LongTensor)) and token_id.size(0) > 1: raise ValueError( f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`." @@ -203,7 +203,7 @@ def __init__(self, token_ids: Union[List[int], torch.LongTensor]): super(Constraint, self).__init__() is_int_list = isinstance(token_ids, List) and isinstance(token_ids[0], int) - is_long_tensor = isinstance(token_ids, torch.LongTensor) and len(token_ids.size()) == 1 + is_long_tensor = isinstance(token_ids, (torch.LongTensor, torch.cuda.LongTensor)) and len(token_ids.size()) == 1 if not (is_int_list or is_long_tensor) or torch.any(token_ids < 0): raise ValueError( f"`token_ids` has to be a single list of positive integers or a `torch.LongTensor` but is {token_ids}" From 12ac97d9c4f206e2e8d487c1fddf2ebfa13a47cd Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Thu, 3 Feb 2022 02:23:53 +0000 Subject: [PATCH 21/34] added examples and testing for them and founx & fixed a bug in beam_search and constrained_beam_search --- docs/source/internal/generation_utils.mdx | 2 - src/transformers/__init__.py | 3 +- .../generation_beam_constraints.py | 96 +++----------- src/transformers/generation_beam_search.py | 26 +++- src/transformers/generation_utils.py | 14 ++- src/transformers/utils/dummy_pt_objects.py | 7 -- tests/test_generation_utils.py | 119 ++++++++++++++++++ 7 files changed, 172 insertions(+), 95 deletions(-) diff --git a/docs/source/internal/generation_utils.mdx b/docs/source/internal/generation_utils.mdx index c157247b97e3..9eb4abe06d34 100644 --- a/docs/source/internal/generation_utils.mdx +++ b/docs/source/internal/generation_utils.mdx @@ -197,8 +197,6 @@ A [`Constraint`] can be used to force the generation to include specific tokens [[autodoc]] Constraint -[[autodoc]] TokenConstraint - [[autodoc]] PhrasalConstraint [[autodoc]] ConstraintListState diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 32efd5c8e3ea..e715e6f894e6 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -612,7 +612,6 @@ "Constraint", "ConstraintListState", "PhrasalConstraint", - "TokenConstraint", ] _import_structure["generation_beam_search"] = ["BeamScorer", "BeamSearchScorer", "ConstrainedBeamSearchScorer"] _import_structure["generation_logits_process"] = [ @@ -2729,7 +2728,7 @@ TextDataset, TextDatasetForNextSentencePrediction, ) - from .generation_beam_constraints import Constraint, ConstraintListState, PhrasalConstraint, TokenConstraint + from .generation_beam_constraints import Constraint, ConstraintListState, PhrasalConstraint from .generation_beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from .generation_logits_process import ( ForcedBOSTokenLogitsProcessor, diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 897f8dd5ffa2..d647589ba94a 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -141,79 +141,22 @@ def copy(self, stateful=False): ) -class TokenConstraint(Constraint): - r""" - [`Constraint`] enforcing that a specific token is generated. - - Args: - token_id (`int`): - The token that must be generated by the output. - """ - - @add_start_docstrings(TOKEN_CONSTRAINT_DOCSTRING) - def __init__(self, token_id: int): - super(Constraint, self).__init__() - if not (isinstance(token_id, int) or isinstance(token_id, torch.LongTensor)) or token_id < 0: - raise ValueError( - f"`token_id` has to be a positive integer or a `torch.LongTensor` with one positive integer, but is {token_id}" - ) - else: - if isinstance(token_id, torch.LongTensor) and token_id.size(0) > 1: - raise ValueError( - f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." - "For sequential constraints for multiple tokens, refer to `PhrasalConstraint`." - ) - - self.token_id = token_id - self.token_ids = token_id # for compatibility reasons - self.completed = False - - def advance(self): - return self.token_id - - def does_advance(self, token_id: Union[int, torch.LongTensor]): - return token_id == self.token_id - - def update(self, token_id: Union[int, torch.LongTensor]): - if not isinstance(token_id, int) or token_id < 0: - raise ValueError(f"`token_id` has to be a positive integer, but is {token_id}") - - if self.does_advance(token_id): - self.completed = True - return True, True, True # stepped, completed, reset - else: - return False, False, False - - def reset(self): - self.completed = False - - def remaining(self): - return 0 if self.completed else 1 - - def copy(self, stateful=False): - constraint = TokenConstraint(self.token_id) - if stateful: - constraint.completed = self.completed - return constraint - - class PhrasalConstraint(Constraint): @add_start_docstrings(PHRASAL_CONSTRAINT_DOCSTRING) def __init__(self, token_ids: Union[List[int], torch.LongTensor]): super(Constraint, self).__init__() is_int_list = isinstance(token_ids, List) and isinstance(token_ids[0], int) - is_long_tensor = isinstance(token_ids, torch.LongTensor) and len(token_ids.size()) == 1 - if not (is_int_list or is_long_tensor) or torch.any(token_ids < 0): - raise ValueError( - f"`token_ids` has to be a single list of positive integers or a `torch.LongTensor` but is {token_ids}" - ) - else: - if (is_int_list or is_long_tensor) and token_ids.size(0) == 1: - raise ValueError( - f"`token_ids` has to be list of positive integers or a `torch.LongTensor` but is {token_ids}" - "For single token constraints, refer to `TokenConstraint`." - ) + is_tensor = isinstance(token_ids, torch.Tensor) + is_int_tensor = ( + is_tensor and token_ids.dtype in [torch.int16, torch.int32, torch.int64] and len(token_ids.size()) == 1 + ) + not_positive = torch.any(token_ids < 0) if is_tensor else len([t for t in token_ids if t < 0]) > 0 + if isinstance(token_ids, int) or not (is_int_list or is_int_tensor) or not_positive: + raise ValueError(f"`token_ids` has to be a single list or tensor of positive integers but is {token_ids}") + + if not is_tensor: + token_ids = torch.tensor(token_ids) self.token_ids = token_ids @@ -227,7 +170,8 @@ def advance(self): def does_advance(self, token_id: int): if self.completed: return False - return token_id.cpu() == self.token_ids[self.fulfilled_idx + 1] + # move to cpu to guarantee no device issues. + return token_id.cpu() == self.token_ids[self.fulfilled_idx + 1].cpu() def update(self, token_id: int): stepped = False @@ -270,7 +214,7 @@ def __init__(self, constraints: List[Constraint]): self.constraints = constraints # max # of steps required to fulfill a given constraint - self.max_seqlen = max([c.seqlen for c in constraints]) + self.max_seqlen = max([c.seqlen for c in constraints if isinstance(c, PhrasalConstraint)]) self.n_constraints = len(constraints) self.completed = False @@ -325,8 +269,11 @@ def reset(self, token_ids: Optional[torch.LongTensor]): if token_ids is not None and token_ids.size(0) > 0: for token in token_ids: + # completes or steps **one** constraint complete, stepped = self.add(token) - if complete: + + # the entire list of constraints are fulfilled + if self.completed: break return self @@ -334,15 +281,6 @@ def reset(self, token_ids: Optional[torch.LongTensor]): def add(self, token_id: Union[int, torch.LongTensor]): complete, stepped = False, False - # if isinstance(token_id, torch.LongTensor): - # if (token_id.size(0)) > 1: - # raise ValueError( - # f"`token_id` has to be a positive integer or a `torch.LongTensor` with one integer, but is {token_id}." - # "It must have length 1." - # ) - # else: - # token_id = token_id[0] - if self.completed: complete = True stepped = False diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index e32d7d63f44f..c886d7eb56b3 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -367,7 +367,8 @@ def finalize( best_scores[i * self.num_beam_hyps_to_keep + j] = best_score # prepare for adding eos - sent_max_len = min(sent_lengths.max().item() + 1, max_length) + sent_lengths_max = sent_lengths.max().item() + 1 + sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) # shorter batches are padded if needed if sent_lengths.min().item() != sent_lengths.max().item(): @@ -377,8 +378,9 @@ def finalize( # fill with hypotheses and eos_token_id if the latter fits in for i, hypo in enumerate(best): decoded[i, : sent_lengths[i]] = hypo - if sent_lengths[i] < max_length: + if sent_lengths[i] < sent_max_len: decoded[i, sent_lengths[i]] = eos_token_id + return UserDict( { "sequences": decoded, @@ -744,6 +746,8 @@ def finalize( # all open beam hypotheses are added to the beam hypothesis # beam hypothesis class automatically keeps the best beams + + ids_collect = [] for beam_id in range(self.num_beams): batch_beam_idx = batch_idx * self.num_beams + beam_id final_score = final_beam_scores[batch_beam_idx].item() @@ -752,6 +756,19 @@ def finalize( completes_constraint = self.check_completes_constraints(final_tokens) if completes_constraint: beam_hyp.add(final_tokens, final_score) + ids_collect.append(beam_id) + + # due to overly complex constraints or other factors, sometimes we can't gaurantee a successful + # generation. In these cases we simply return the highest scoring outputs. + if len(ids_collect) < self.num_beam_hyps_to_keep: + for beam_id in range(self.num_beams): + if beam_id not in ids_collect: + batch_beam_idx = batch_idx * self.num_beams + beam_id + final_score = final_beam_scores[batch_beam_idx].item() + final_tokens = input_ids[batch_beam_idx] + beam_hyp.add(final_tokens, final_score) + if len(ids_collect) >= self.num_beam_hyps_to_keep: + break # select the best hypotheses sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) @@ -772,7 +789,8 @@ def finalize( best_scores[i * self.num_beam_hyps_to_keep + j] = best_score # prepare for adding eos - sent_max_len = min(sent_lengths.max().item() + 1, max_length) + sent_lengths_max = sent_lengths.max().item() + 1 + sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) # shorter batches are padded if needed if sent_lengths.min().item() != sent_lengths.max().item(): @@ -782,7 +800,7 @@ def finalize( # fill with hypotheses and eos_token_id if the latter fits in for i, hypo in enumerate(best): decoded[i, : sent_lengths[i]] = hypo - if sent_lengths[i] < max_length: + if sent_lengths[i] < sent_max_len: decoded[i, sent_lengths[i]] = eos_token_id return UserDict( { diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 31440b50db60..1598f9be6094 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -1340,6 +1340,15 @@ def generate( if stopping_criteria.max_length is None: raise ValueError("`max_length` needs to be a stopping_criteria for now.") + if num_beams <= 1: + raise ValueError("`num_beams` needs to be greater than 1 for constrained genertation.") + + if do_sample: + raise ValueError("`do_sample` needs to be false for constrained generation.") + + if num_beam_groups is not None and num_beam_groups > 1: + raise ValueError("`num_beam_groups` not supported yet for constrained generation.") + # 10. prepare beam search scorer constrained_beam_scorer = ConstrainedBeamSearchScorer( constraints=constraints, @@ -2930,7 +2939,9 @@ def constrained_beam_search( ... ) ... } - >>> constraints = [PhrasalConstraint(tokenizer.encode("required phrase")[0])] + >>> constraint_str = "sind" + >>> constraint_token_ids = tokenizer.encode(constraint_str)[:-1] # slice to remove eos token + >>> constraints = [PhrasalConstraint(token_ids=constraint_token_ids)] >>> # instantiate beam scorer @@ -2950,6 +2961,7 @@ def constrained_beam_search( ... ) >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True)) + # => ['Wie alter sind Sie?'] ```""" # init values logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 1e6f0b4eb886..e2354d86075a 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -101,13 +101,6 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) -class TokenConstraint(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - class BeamScorer(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 274ed554a7fa..9a2a7449a372 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -27,6 +27,8 @@ import torch from transformers import ( + AutoModelForSeq2SeqLM, + AutoTokenizer, BartForConditionalGeneration, BartTokenizer, GPT2LMHeadModel, @@ -2311,3 +2313,120 @@ def test_transition_scores_group_beam_search_encoder_decoder(self): transition_scores_sum = transition_scores.sum(-1) self.assertTrue(torch.allclose(transition_scores_sum, outputs.sequences_scores, atol=1e-3)) + + def test_constrained_beam_search(self): + model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device) + tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + + force_tokens = tokenizer.encode(" scared", return_tensors="pt").to(torch_device)[0] + force_tokens_2 = tokenizer.encode(" big weapons", return_tensors="pt").to(torch_device)[0] + + constraints = [ + PhrasalConstraint(force_tokens), + PhrasalConstraint(force_tokens_2), + ] + + starting_text = ["The soldiers were not prepared and"] + + input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device) + + outputs = model.generate( + input_ids, + constraints=constraints, + num_beams=10, + num_return_sequences=1, + no_repeat_ngram_size=1, + max_length=30, + remove_invalid_values=True, + ) + + generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True) + + self.assertListEqual( + generated_text, + [ + "The soldiers were not prepared and didn't know how big the big weapons would be, so they scared them off. They had no idea what to do", + ], + ) + + def test_beam_search_example_integration(self): + tokenizer = AutoTokenizer.from_pretrained("t5-base") + model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") + + encoder_input_str = "translate English to German: How old are you?" + encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids + + # lets run beam search using 3 beams + num_beams = 3 + # define decoder start token ids + input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long) + input_ids = input_ids * model.config.decoder_start_token_id + + # add encoder_outputs to model keyword arguments + model_kwargs = { + "encoder_outputs": model.get_encoder()( + encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True + ) + } + + # instantiate beam scorer + beam_scorer = BeamSearchScorer( + batch_size=1, + num_beams=num_beams, + device=model.device, + ) + + # instantiate logits processors + logits_processor = LogitsProcessorList( + [ + MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id), + ] + ) + + outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) + outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) + + self.assertListEqual(outputs, ["Wie alt bist du?"]) + + def test_constrained_beam_search_example_integration(self): + tokenizer = AutoTokenizer.from_pretrained("t5-base") + model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") + + encoder_input_str = "translate English to German: How old are you?" + encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids + + # lets run beam search using 5 beams + num_beams = 5 + # define decoder start token ids + input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long) + input_ids = input_ids * model.config.decoder_start_token_id + + # add encoder_outputs to model keyword arguments + model_kwargs = { + "encoder_outputs": model.get_encoder()( + encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True + ) + } + + constraint_str = "sind" + constraint_token_ids = tokenizer.encode(constraint_str)[:-1] # remove eos token + constraints = [PhrasalConstraint(token_ids=constraint_token_ids)] + + # instantiate beam scorer + beam_scorer = ConstrainedBeamSearchScorer( + batch_size=1, num_beams=num_beams, device=model.device, constraints=constraints + ) + + # instantiate logits processors + logits_processor = LogitsProcessorList( + [ + MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id), + ] + ) + + outputs = model.constrained_beam_search( + input_ids, beam_scorer, constraints=constraints, logits_processor=logits_processor, **model_kwargs + ) + outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) + + self.assertListEqual(outputs, ["Wie alter sind Sie?"]) From 2169a9ff29f2dedc6581a2e11de7667356353a6d Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Thu, 3 Feb 2022 02:51:13 +0000 Subject: [PATCH 22/34] deleted accidentally added test halting code with assert False --- src/transformers/generation_beam_search.py | 4 ++-- tests/test_generation_utils.py | 4 ---- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index c886d7eb56b3..c65a2a10afef 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -757,7 +757,7 @@ def finalize( if completes_constraint: beam_hyp.add(final_tokens, final_score) ids_collect.append(beam_id) - + # due to overly complex constraints or other factors, sometimes we can't gaurantee a successful # generation. In these cases we simply return the highest scoring outputs. if len(ids_collect) < self.num_beam_hyps_to_keep: @@ -769,7 +769,7 @@ def finalize( beam_hyp.add(final_tokens, final_score) if len(ids_collect) >= self.num_beam_hyps_to_keep: break - + # select the best hypotheses sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) best = [] diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 9a2a7449a372..6121e7374cc9 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -805,10 +805,6 @@ def test_beam_search_generate(self): logits_processor=logits_processor, ) - print("input_ids", input_ids) - print("output_generate", output_generate) - assert False - self.assertListEqual(output_generate.tolist(), output_beam_search.tolist()) # check `generate()` and `beam_search()` are equal for `num_return_sequences` From 77660bdaff0070feea0285817ae23a71f53c6457 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Thu, 3 Feb 2022 02:53:24 +0000 Subject: [PATCH 23/34] code reformat --- src/transformers/generation_beam_search.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index c65a2a10afef..c886d7eb56b3 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -757,7 +757,7 @@ def finalize( if completes_constraint: beam_hyp.add(final_tokens, final_score) ids_collect.append(beam_id) - + # due to overly complex constraints or other factors, sometimes we can't gaurantee a successful # generation. In these cases we simply return the highest scoring outputs. if len(ids_collect) < self.num_beam_hyps_to_keep: @@ -769,7 +769,7 @@ def finalize( beam_hyp.add(final_tokens, final_score) if len(ids_collect) >= self.num_beam_hyps_to_keep: break - + # select the best hypotheses sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) best = [] From b21aae0d88d8b8e22d088b255ed2bc64242fc417 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Fri, 4 Feb 2022 11:16:03 +0900 Subject: [PATCH 24/34] Update tests/test_generation_utils.py Co-authored-by: Patrick von Platen --- tests/test_generation_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 6121e7374cc9..563e42de7182 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2310,6 +2310,7 @@ def test_transition_scores_group_beam_search_encoder_decoder(self): self.assertTrue(torch.allclose(transition_scores_sum, outputs.sequences_scores, atol=1e-3)) + @slow def test_constrained_beam_search(self): model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device) tokenizer = GPT2Tokenizer.from_pretrained("gpt2") From 00506217368191a5a11b604f7564237256e87d81 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Fri, 4 Feb 2022 11:16:19 +0900 Subject: [PATCH 25/34] Update tests/test_generation_utils.py Co-authored-by: Patrick von Platen --- tests/test_generation_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 563e42de7182..823d45befa02 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2310,6 +2310,7 @@ def test_transition_scores_group_beam_search_encoder_decoder(self): self.assertTrue(torch.allclose(transition_scores_sum, outputs.sequences_scores, atol=1e-3)) + @slow @slow def test_constrained_beam_search(self): model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device) From 3e35647fcba7f71fac2ade57900e0b0e507d3efd Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Fri, 4 Feb 2022 11:16:28 +0900 Subject: [PATCH 26/34] Update tests/test_generation_utils.py Co-authored-by: Patrick von Platen --- tests/test_generation_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 823d45befa02..2316eb696cfb 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2347,6 +2347,7 @@ def test_constrained_beam_search(self): ], ) + @slow def test_beam_search_example_integration(self): tokenizer = AutoTokenizer.from_pretrained("t5-base") model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") From edd86815d46bef31bf5a1d7375e085b7c8687724 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Fri, 4 Feb 2022 11:16:35 +0900 Subject: [PATCH 27/34] Update tests/test_generation_utils.py Co-authored-by: Patrick von Platen --- tests/test_generation_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 2316eb696cfb..2a80dc263c31 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2387,6 +2387,7 @@ def test_beam_search_example_integration(self): self.assertListEqual(outputs, ["Wie alt bist du?"]) + @slow def test_constrained_beam_search_example_integration(self): tokenizer = AutoTokenizer.from_pretrained("t5-base") model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") From e1f6419857c82e4b0c5a58d450e06100f07a1ec2 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 7 Feb 2022 19:41:28 +0100 Subject: [PATCH 28/34] Update tests/test_generation_utils.py --- tests/test_generation_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 2a80dc263c31..ae44877f1fca 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2310,7 +2310,6 @@ def test_transition_scores_group_beam_search_encoder_decoder(self): self.assertTrue(torch.allclose(transition_scores_sum, outputs.sequences_scores, atol=1e-3)) - @slow @slow def test_constrained_beam_search(self): model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device) From ba7a3105719bbeaa5d61a2add09d7d06a1419313 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Tue, 8 Feb 2022 06:01:00 +0000 Subject: [PATCH 29/34] fixing based on comments on PR --- .../generation_beam_constraints.py | 92 +++++++++---------- src/transformers/generation_beam_search.py | 66 +++++++------ tests/test_generation_beam_search.py | 5 - tests/test_generation_utils.py | 2 +- 4 files changed, 73 insertions(+), 92 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index d647589ba94a..dcd7f0fbc3f6 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -6,22 +6,6 @@ from .file_utils import add_start_docstrings -TOKEN_CONSTRAINT_DOCSTRING = r""" - [`Constraint`] enforcing that a specific token is included in the output. - - Args: - token_ids (`torch.LongTensor`): - The sequence of tokens that must be generated by the output. -""" - -PHRASAL_CONSTRAINT_DOCSTRING = r""" - [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. - - Args: - token_id (`int`): - The id of the token that must be generated by the output. -""" - class Constraint(ABC): r"""Abstract base class for all constraints that can be applied during generation. @@ -29,7 +13,7 @@ class Constraint(ABC): All classes that inherit Constraint must follow the requirement that - ``` + ```py completed = False while(not completed): _, completed = constraint.update(constraint.advance()) @@ -52,14 +36,17 @@ def test(self): if counter == 1: self.reset() advance = self.advance() - assert self.does_advance(advance) + if not self.does_advance(advance): + raise Exception("Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.") + stepped, completed, reset = self.update(advance) counter += 1 if counter > 10000: raise Exception("update() does not fulfill the constraint.") - assert self.remaining() == 0 + if self.remaining() != 0: + raise Exception("Custom Constraint is not defined correctly.") @abstractmethod def advance(self): @@ -96,11 +83,11 @@ def update(self, token_id: int): token_id(`int`): The id of a newly generated token in the beam search. returns: - stepped(`boolean`): + stepped(`bool`): Whether this constraint has become one step closer to being fulfuilled. - completed(`boolean`): + completed(`bool`): Whether this constraint has been completely fulfilled by this token being generated. - reset (`boolean`): + reset (`bool`): Whether this constraint has reset its progress by this token being generated. """ raise NotImplementedError( @@ -132,7 +119,7 @@ def copy(self, stateful=False): Creates a new instance of this constraint. Args: - stateful(`boolean`): Whether to not only copy the constraint for new instance, but also its state. + stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state. Returns: constraint(`Constraint`): The same constraint as the one being called from. """ @@ -142,7 +129,14 @@ def copy(self, stateful=False): class PhrasalConstraint(Constraint): - @add_start_docstrings(PHRASAL_CONSTRAINT_DOCSTRING) + r""" + [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. + + Args: + token_id (`int`): + The id of the token that must be generated by the output. + """ + def __init__(self, token_ids: Union[List[int], torch.LongTensor]): super(Constraint, self).__init__() @@ -208,8 +202,10 @@ def copy(self, stateful=False): return new_constraint -# For beam scorers to track its progress through a list of constraints. class ConstraintListState: + r""" + A class for beam scorers to track its progress through a list of constraints. + """ def __init__(self, constraints: List[Constraint]): self.constraints = constraints @@ -287,28 +283,25 @@ def add(self, token_id: Union[int, torch.LongTensor]): return complete, stepped if self.inprogress_constraint is not None: - """ - In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current - job, simply update the state - """ + # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current + # job, simply update the state + stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: - """ - 1. If the next token breaks the progress, then we must restart. - e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books". + # 1. If the next token breaks the progress, then we must restart. + # e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books". + + # But that doesn't mean we self.init_state(), since we only reset the state for this particular + # constraint, not the full list of constraints. - But that doesn't mean we self.init_state(), since we only reset the state for this particular - constraint, not the full list of constraints. - """ self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False)) self.inprogress_constraint = None if complete: - """ - 2. If the next token completes the constraint, move it to completed list, set - inprogress to None. If there are no pending constraints either, then this full list of constraints - is complete. - """ + # 2. If the next token completes the constraint, move it to completed list, set + # inprogress to None. If there are no pending constraints either, then this full list of constraints + # is complete. + self.complete_constraints.append(self.inprogress_constraint) self.inprogress_constraint = None @@ -317,10 +310,9 @@ def add(self, token_id: Union[int, torch.LongTensor]): self.completed = True else: - """ - Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list - of constraints? - """ + # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list + # of constraints? + for cidx, pending_constraint in enumerate(self.pending_constraints): if pending_constraint.does_advance(token_id): stepped, complete, reset = pending_constraint.update(token_id) @@ -339,18 +331,16 @@ def add(self, token_id: Union[int, torch.LongTensor]): self.inprogress_constraint = pending_constraint if complete or stepped: - """ - If we made any progress at all, then it's at least not a "pending constraint". - """ + # If we made any progress at all, then it's at least not a "pending constraint". + self.pending_constraints = ( self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :] ) if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: - """ - If there's no longer any pending after this and no inprogress either, then we must be - complete. - """ + # If there's no longer any pending after this and no inprogress either, then we must be + # complete. + self.completed = True break # prevent accidentally stepping through multiple constraints with just one token. diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index c886d7eb56b3..822584265978 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -367,18 +367,16 @@ def finalize( best_scores[i * self.num_beam_hyps_to_keep + j] = best_score # prepare for adding eos - sent_lengths_max = sent_lengths.max().item() + 1 - sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max + sent_max_len = min(sent_lengths.max().item() + 1, max_length) decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) # shorter batches are padded if needed if sent_lengths.min().item() != sent_lengths.max().item(): assert pad_token_id is not None, "`pad_token_id` has to be defined" decoded.fill_(pad_token_id) - # fill with hypotheses and eos_token_id if the latter fits in for i, hypo in enumerate(best): decoded[i, : sent_lengths[i]] = hypo - if sent_lengths[i] < sent_max_len: + if sent_lengths[i] < max_length: decoded[i, sent_lengths[i]] = eos_token_id return UserDict( @@ -626,14 +624,13 @@ def step_sentence_constraint( # need to make new hypothesis that advance the constraints track_new = {"new_seqs": [], "new_states": [], "new_indices": [], "new_tokens": [], "new_scores": []} for seq_idx, pre_seq in enumerate(this_batch_input_ids): - """ - pre_seq = ith sequence generated before this step. - - input_ids -> (topk) generic beam search best model next tokens - -> (advance) constraints forcing the next token - either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of - hypotheses. - """ + # pre_seq = ith sequence generated before this step. + + # input_ids -> (topk) generic beam search best model next tokens + # -> (advance) constraints forcing the next token + # either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of + # hypotheses. + topk_state = topk_contraint_states[seq_idx] topk_state.reset(full_hypotheses[seq_idx]) @@ -656,26 +653,25 @@ def step_sentence_constraint( track_new["new_scores"].append(this_batch_token_scores[seq_idx].take(advance_token)) track_new["new_states"].append(new_state) elif push_progress: - """ - Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences that - actually fulfill our constraints. For example, let constraints == ["loves pies"] and - - pre_seq_1 = "The child loves pies and" pre_seq_2 = "The child plays in the playground and" - - Without this step, if `sent_beam_indices` is something like [1,1], then - 1. `pre_seq_1` won't be added to the list of (topk) hypothesis since it's not in the indices and - 2. it won't be added to the list of (advance) hypothesis since it's completed already. (this is - the else part of `if constraints_completed[seq_idx]`) - 3. it ends up simply getting removed from consideration. - - #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, - especially if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam - search times, since completed sequences keep getting removed after all this effort for constrained - generation. - - Here, we basically take `pre_seq_1` and to "push" it into the considered list of hypotheses, by simply - appending the next likely token in the vocabulary and adding it to the list of hypotheses. - """ + # Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences that + # actually fulfill our constraints. For example, let constraints == ["loves pies"] and + + # pre_seq_1 = "The child loves pies and" pre_seq_2 = "The child plays in the playground and" + + # Without this step, if `sent_beam_indices` is something like [1,1], then + # 1. `pre_seq_1` won't be added to the list of (topk) hypothesis since it's not in the indices and + # 2. it won't be added to the list of (advance) hypothesis since it's completed already. (this is + # the else part of `if constraints_completed[seq_idx]`) + # 3. it ends up simply getting removed from consideration. + + # #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, + # especially if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam + # search times, since completed sequences keep getting removed after all this effort for constrained + # generation. + + # Here, we basically take `pre_seq_1` and to "push" it into the considered list of hypotheses, by simply + # appending the next likely token in the vocabulary and adding it to the list of hypotheses. + new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) # some next probable token advance_seq = torch.cat((pre_seq, new_token.unsqueeze(0)), -1) @@ -704,9 +700,9 @@ def step_sentence_constraint( zipped = all_banks * 100 + all_scores indices = zipped.sort(descending=True).indices sorted_banks = all_banks[indices] - """ - Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} - """ + + # Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} + counter = -1 cur_bank = sorted_banks[0] increments = [] diff --git a/tests/test_generation_beam_search.py b/tests/test_generation_beam_search.py index 6b7b84542ebf..125b82393b35 100644 --- a/tests/test_generation_beam_search.py +++ b/tests/test_generation_beam_search.py @@ -414,11 +414,6 @@ def check_constrained_beam_scorer_finalize( # make sure corresponding score is as good as possible to surely be picked first next_scores[0, 0] = 0.0 - # # because constrainted beam search can't possibly fulfill the constraints in one pass - # # especially if constraints involves more than one token, so we repeat this several times. - # # doesn't *really* matter that we're not adjusting the scores & tokens TBH. - # repeat = 10 - # for _ in range(repeat): beam_outputs = constrained_beam_scorer.process( input_ids, next_scores, tokens, next_indices, scores_for_all_vocab, eos_token_id=self.eos_token_id ) diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index ae44877f1fca..ba11e6db0942 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -1514,7 +1514,7 @@ def _check_sequence_inside_sequence(self, tensor_1, tensor_2): flag = True break - assert flag + self.assertTrue(flag) @require_torch From 77a18ae82fa5b9a7ed7dc15ae3ab000acf942b26 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Tue, 8 Feb 2022 07:57:16 +0000 Subject: [PATCH 30/34] took out the testing code that should but work fails without the beam search moditification ; style changes --- .../generation_beam_constraints.py | 24 +++++------ src/transformers/generation_beam_search.py | 8 ++-- tests/test_generation_utils.py | 40 ------------------- 3 files changed, 16 insertions(+), 56 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index dcd7f0fbc3f6..57b86d328414 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -3,9 +3,6 @@ import torch -from .file_utils import add_start_docstrings - - class Constraint(ABC): r"""Abstract base class for all constraints that can be applied during generation. @@ -15,7 +12,7 @@ class Constraint(ABC): ```py completed = False - while(not completed): + while not completed: _, completed = constraint.update(constraint.advance()) ``` @@ -37,7 +34,9 @@ def test(self): self.reset() advance = self.advance() if not self.does_advance(advance): - raise Exception("Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.") + raise Exception( + "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true." + ) stepped, completed, reset = self.update(advance) counter += 1 @@ -130,11 +129,11 @@ def copy(self, stateful=False): class PhrasalConstraint(Constraint): r""" - [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. + [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. - Args: - token_id (`int`): - The id of the token that must be generated by the output. + Args: + token_id (`int`): + The id of the token that must be generated by the output. """ def __init__(self, token_ids: Union[List[int], torch.LongTensor]): @@ -206,6 +205,7 @@ class ConstraintListState: r""" A class for beam scorers to track its progress through a list of constraints. """ + def __init__(self, constraints: List[Constraint]): self.constraints = constraints @@ -285,7 +285,7 @@ def add(self, token_id: Union[int, torch.LongTensor]): if self.inprogress_constraint is not None: # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current # job, simply update the state - + stepped, complete, reset = self.inprogress_constraint.update(token_id) if reset: # 1. If the next token breaks the progress, then we must restart. @@ -301,7 +301,7 @@ def add(self, token_id: Union[int, torch.LongTensor]): # 2. If the next token completes the constraint, move it to completed list, set # inprogress to None. If there are no pending constraints either, then this full list of constraints # is complete. - + self.complete_constraints.append(self.inprogress_constraint) self.inprogress_constraint = None @@ -312,7 +312,7 @@ def add(self, token_id: Union[int, torch.LongTensor]): else: # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list # of constraints? - + for cidx, pending_constraint in enumerate(self.pending_constraints): if pending_constraint.does_advance(token_id): stepped, complete, reset = pending_constraint.update(token_id) diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 822584265978..9f85fcf65d5a 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -367,7 +367,7 @@ def finalize( best_scores[i * self.num_beam_hyps_to_keep + j] = best_score # prepare for adding eos - sent_max_len = min(sent_lengths.max().item() + 1, max_length) + sent_max_len = min(sent_lengths.max().item() + 1, max_length) decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) # shorter batches are padded if needed if sent_lengths.min().item() != sent_lengths.max().item(): @@ -376,7 +376,7 @@ def finalize( # fill with hypotheses and eos_token_id if the latter fits in for i, hypo in enumerate(best): decoded[i, : sent_lengths[i]] = hypo - if sent_lengths[i] < max_length: + if sent_lengths[i] < max_length: decoded[i, sent_lengths[i]] = eos_token_id return UserDict( @@ -700,9 +700,9 @@ def step_sentence_constraint( zipped = all_banks * 100 + all_scores indices = zipped.sort(descending=True).indices sorted_banks = all_banks[indices] - + # Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} - + counter = -1 cur_bank = sorted_banks[0] increments = [] diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index ba11e6db0942..dbe7c2539787 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -2346,46 +2346,6 @@ def test_constrained_beam_search(self): ], ) - @slow - def test_beam_search_example_integration(self): - tokenizer = AutoTokenizer.from_pretrained("t5-base") - model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") - - encoder_input_str = "translate English to German: How old are you?" - encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids - - # lets run beam search using 3 beams - num_beams = 3 - # define decoder start token ids - input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long) - input_ids = input_ids * model.config.decoder_start_token_id - - # add encoder_outputs to model keyword arguments - model_kwargs = { - "encoder_outputs": model.get_encoder()( - encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True - ) - } - - # instantiate beam scorer - beam_scorer = BeamSearchScorer( - batch_size=1, - num_beams=num_beams, - device=model.device, - ) - - # instantiate logits processors - logits_processor = LogitsProcessorList( - [ - MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id), - ] - ) - - outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs) - outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) - - self.assertListEqual(outputs, ["Wie alt bist du?"]) - @slow def test_constrained_beam_search_example_integration(self): tokenizer = AutoTokenizer.from_pretrained("t5-base") From 7a78633ed95ac04fdf874f166f7d87599987a4dc Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Wed, 9 Feb 2022 05:14:37 +0000 Subject: [PATCH 31/34] fixing comments issues --- src/transformers/generation_beam_search.py | 90 +++++++++++----------- src/transformers/generation_utils.py | 1 - 2 files changed, 43 insertions(+), 48 deletions(-) diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py index 9f85fcf65d5a..d0384b7624e3 100644 --- a/src/transformers/generation_beam_search.py +++ b/src/transformers/generation_beam_search.py @@ -85,41 +85,6 @@ """ -CONSTRAINT_PROCESS_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): - Indices of input sequence tokens in the vocabulary. - - Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See - [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. - - [What are input IDs?](../glossary#input-ids) - next_scores (`torch.FloatTensor` of shape `(batch_size, 2 * num_beams)`): - Current scores of the top `2 * num_beams` non-finished beam hypotheses. - next_tokens (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): - `input_ids` of the tokens corresponding to the top `2 * num_beams` non-finished beam hypotheses. - next_indices (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): - Beam indices indicating to which beam hypothesis the `next_tokens` correspond. - scores_for_all_vocab (`torch.FloatTensor` of shape `(batch_size * num_beams, sequence_length)`): - The scores of all tokens in the vocabulary for each of the beam hypotheses. - pad_token_id (`int`, *optional*): - The id of the *padding* token. - eos_token_id (`int`, *optional*): - The id of the *end-of-sequence* token. - - Return: - `UserDict`: A dictionary composed of the fields as defined above: - - - **next_beam_scores** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Updated scores of all - non-finished beams. - - **next_beam_tokens** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Next tokens to be added - to the non-finished beam_hypotheses. - - **next_beam_indices** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Beam indices - indicating to which beam the next tokens shall be added. - -""" - - class BeamScorer(ABC): """ Abstract base class for all beam scorers that are used for [`~PreTrainedModel.beam_search`] and @@ -481,7 +446,6 @@ def check_completes_constraints(self, sequence): new_state = new_state.reset(sequence) return new_state.completed - @add_start_docstrings(CONSTRAINT_PROCESS_INPUTS_DOCSTRING) def process( self, input_ids: torch.LongTensor, @@ -492,6 +456,41 @@ def process( pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, ) -> Tuple[torch.Tensor]: + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + next_scores (`torch.FloatTensor` of shape `(batch_size, 2 * num_beams)`): + Current scores of the top `2 * num_beams` non-finished beam hypotheses. + next_tokens (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + `input_ids` of the tokens corresponding to the top `2 * num_beams` non-finished beam hypotheses. + next_indices (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + Beam indices indicating to which beam hypothesis the `next_tokens` correspond. + scores_for_all_vocab (`torch.FloatTensor` of shape `(batch_size * num_beams, sequence_length)`): + The scores of all tokens in the vocabulary for each of the beam hypotheses. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`int`, *optional*): + The id of the *end-of-sequence* token. + + Return: + `UserDict`: A dictionary composed of the fields as defined above: + + - **next_beam_scores** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Updated scores of + all + non-finished beams. + - **next_beam_tokens** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Next tokens to be + added + to the non-finished beam_hypotheses. + - **next_beam_indices** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Beam indices + indicating to which beam the next tokens shall be added. + """ + cur_len = input_ids.shape[-1] batch_size = len(self._beam_hyps) if not (batch_size == (input_ids.shape[0] // self.group_size)): @@ -596,19 +595,16 @@ def step_sentence_constraint( sent_beam_indices: torch.LongTensor, push_progress: bool = False, ): - # from transformers import GPT2Tokenizer - # tokenizer = GPT2Tokenizer.from_pretrained("gpt2") - """ - sent_beam_tokens are the next {num_beams} number of tokens that are under consideration for this beam - (candidate next tokens) + # sent_beam_tokens are the next {num_beams} number of tokens that are under consideration for this beam + # (candidate next tokens) - 1. Adding "advance_tokens" - using ConstraintStateList.advance(), we propose new tokens to be added into this "candidate list" that will - advance us in fulfilling the constraints. + # 1. Adding "advance_tokens" + # using ConstraintStateList.advance(), we propose new tokens to be added into this "candidate list" that will + # advance us in fulfilling the constraints. + + # 2. Selecting best candidates such that we end up with highest probable candidates + # that fulfill our constraints. - 2. Selecting best candidates such that we end up with highest probable candidates - that fulfill our constraints. - """ orig_len = sent_beam_indices.size(0) device = sent_beam_indices.device diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index 7919b2c1a646..2ec0dfd116c1 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -851,7 +851,6 @@ def generate( post](https://huggingface.co/blog/how-to-generate). Parameters: - inputs (`torch.Tensor` of shape `(batch_size, sequence_length)`, `(batch_size, sequence_length, feature_dim)` or `(batch_size, num_channels, height, width)`, *optional*): The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the From bbd9e88872cc47c0e848dd32ab5ffd639dcec92a Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Wed, 9 Feb 2022 05:20:05 +0000 Subject: [PATCH 32/34] docstrings for ConstraintListState --- src/transformers/generation_beam_constraints.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 57b86d328414..64f2a222b327 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -204,6 +204,10 @@ def copy(self, stateful=False): class ConstraintListState: r""" A class for beam scorers to track its progress through a list of constraints. + + Args: + constraints (`List[Constraint]`): + A list of [`Constraint`] objects that must be fulfilled by the beam scorer. """ def __init__(self, constraints: List[Constraint]): From 17ab4740aa2710ffc2d6e86a9f0e1c7369d2b233 Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Wed, 9 Feb 2022 05:22:05 +0000 Subject: [PATCH 33/34] typo in PhrsalConstraint docstring --- src/transformers/generation_beam_constraints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 64f2a222b327..93ec55cf0da5 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -132,7 +132,7 @@ class PhrasalConstraint(Constraint): [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. Args: - token_id (`int`): + token_ids (`List[int]`): The id of the token that must be generated by the output. """ From 88e938dc0771c8ff74ad88df148f644283ec614f Mon Sep 17 00:00:00 2001 From: Chan Woo Kim Date: Wed, 9 Feb 2022 09:04:06 +0000 Subject: [PATCH 34/34] docstrings improvements --- src/transformers/generation_beam_constraints.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/transformers/generation_beam_constraints.py b/src/transformers/generation_beam_constraints.py index 93ec55cf0da5..6410d069289a 100644 --- a/src/transformers/generation_beam_constraints.py +++ b/src/transformers/generation_beam_constraints.py @@ -52,7 +52,7 @@ def advance(self): """ When called, returns the token that would take this constraint one step closer to being fulfilled. - returns: + Return: token_ids(`torch.tensor`): Must be a tensor of a list of indexable tokens, not some integer. """ raise NotImplementedError( @@ -81,7 +81,7 @@ def update(self, token_id: int): Args: token_id(`int`): The id of a newly generated token in the beam search. - returns: + Return: stepped(`bool`): Whether this constraint has become one step closer to being fulfuilled. completed(`bool`): @@ -119,7 +119,8 @@ def copy(self, stateful=False): Args: stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state. - Returns: + + Return: constraint(`Constraint`): The same constraint as the one being called from. """ raise NotImplementedError( @@ -237,13 +238,13 @@ def advance(self): """The list of tokens to generate such that we can make progress. By "list" we don't mean the list of token that will fully fulfill a constraint. - Given constraints c_i = {t_ij | j == # of tokens}, If we're not in the middle of progressing through a specific - constraint c_i, we return: + Given constraints `c_i = {t_ij | j == # of tokens}`, If we're not in the middle of progressing through a + specific constraint `c_i`, we return: - [t_k1 for k in indices of unfulfilled constraints] + `[t_k1 for k in indices of unfulfilled constraints]` If we are in the middle of a constraint, then we return: - [t_ij], where i == index of the inprogress constraint, j == the next step for the constraint. + `[t_ij]`, where `i` is the index of the inprogress constraint, `j` is the next step for the constraint. Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint, that's the only one we'll return.