def test_different_per_node_beam_size(self):
# per_node_beam_size = 1
beam_search = BeamSearch(self.end_index,
beam_size=3,
per_node_beam_size=1)
self._check_results(beam_search=beam_search)
# per_node_beam_size = 2
beam_search = BeamSearch(self.end_index,
beam_size=3,
per_node_beam_size=2)
self._check_results(beam_search=beam_search)
def test_greedy_decode_matches_beam_search(self):
beam_search = BeamSearch(
self.model._end_index, max_steps=self.model._max_decoding_steps, beam_size=1
)
training_tensors = self.dataset.as_tensor_dict()
# Get greedy predictions from _forward_loop method of model.
state = self.model._encode(training_tensors["source_tokens"])
state = self.model._init_decoder_state(state)
output_dict_greedy = self.model._forward_loop(state)
output_dict_greedy = self.model.decode(output_dict_greedy)
# Get greedy predictions from beam search (beam size = 1).
state = self.model._encode(training_tensors["source_tokens"])
state = self.model._init_decoder_state(state)
batch_size = state["source_mask"].size()[0]
start_predictions = state["source_mask"].new_full(
(batch_size,), fill_value=self.model._start_index
)
all_top_k_predictions, _ = beam_search.search(
start_predictions, state, self.model.take_step
)
output_dict_beam_search = {"predictions": all_top_k_predictions}
output_dict_beam_search = self.model.decode(output_dict_beam_search)
# Predictions from model._forward_loop and beam_search should match.
assert output_dict_greedy["predicted_tokens"] == output_dict_beam_search["predicted_tokens"]
def test_stochastic_beam_search(self):
initial_predictions = torch.tensor([0] * 5)
batch_size = 5
beam_size = 3
take_step = take_step_with_timestep
gumbel_sampler = GumbelSampler()
top_k, log_probs = BeamSearch(self.end_index,
beam_size=beam_size,
max_steps=10,
sampler=gumbel_sampler).search(
initial_predictions, {}, take_step)
# top_p should be shape `(batch_size, beam_size, max_predicted_length)`.
assert list(top_k.size())[:-1] == [batch_size, beam_size]
assert ((0 <= top_k) & (top_k <= 5)).all()
# log_probs should be shape `(batch_size, beam_size, max_predicted_length)`.
assert list(log_probs.size()) == [batch_size, beam_size]
# Check to make sure that once the end index is predicted, all subsequent tokens
# must be the end index. This has been tested on toy examples in which
for batch in top_k:
for beam in batch:
reached_end = False
for token in beam:
if token == self.end_index:
reached_end = True
if reached_end:
assert token == self.end_index
def test_greedy_search(self):
beam_search = BeamSearch(self.end_index, beam_size=1)
expected_top_k = np.array([[1, 2, 3, 4, 5]])
expected_log_probs = np.log(np.array([0.4])) # pylint: disable=assignment-from-no-return
self._check_results(expected_top_k=expected_top_k,
expected_log_probs=expected_log_probs,
beam_search=beam_search)
def __init__(
self,
vocab: Vocabulary,
embed: TextFieldEmbedder,
encoder_size: int,
decoder_size: int,
num_layers: int,
beam_size: int,
max_decoding_steps: int,
use_bleu: bool = True,
initializer: InitializerApplicator = InitializerApplicator()
) -> None:
super().__init__(vocab)
self.START, self.END = self.vocab.get_token_index(
START_SYMBOL), self.vocab.get_token_index(END_SYMBOL)
self.OOV = self.vocab.get_token_index(self.vocab._oov_token) # pylint: disable=protected-access
self.PAD = self.vocab.get_token_index(self.vocab._padding_token) # pylint: disable=protected-access
self.COPY = self.vocab.get_token_index("@@COPY@@")
self.KEEP = self.vocab.get_token_index("@@KEEP@@")
self.DROP = self.vocab.get_token_index("@@DROP@@")
self.SYMBOL = (self.START, self.END, self.PAD, self.KEEP, self.DROP)
self.vocab_size = vocab.get_vocab_size()
self.EMB = embed
self.emb_size = self.EMB.token_embedder_tokens.output_dim
self.encoder_size, self.decoder_size = encoder_size, decoder_size
self.FACT_ENCODER = FeedForward(3 * self.emb_size, 1, encoder_size,
nn.Tanh())
self.ATTN = AdditiveAttention(encoder_size + decoder_size,
encoder_size)
self.COPY_ATTN = AdditiveAttention(decoder_size, encoder_size)
module = nn.LSTM(self.emb_size,
encoder_size // 2,
num_layers,
bidirectional=True,
batch_first=True)
self.BUFFER = PytorchSeq2SeqWrapper(
module) # BiLSTM to encode draft text
self.STREAM = nn.LSTMCell(2 * encoder_size,
decoder_size) # Store revised text
self.BEAM = BeamSearch(self.END,
max_steps=max_decoding_steps,
beam_size=beam_size)
self.U = nn.Sequential(nn.Linear(2 * encoder_size, decoder_size),
nn.Tanh())
self.ADD = nn.Sequential(nn.Linear(self.emb_size, encoder_size),
nn.Tanh())
self.P = nn.Sequential(
nn.Linear(encoder_size + decoder_size, decoder_size), nn.Tanh())
self.W = nn.Linear(decoder_size, self.vocab_size)
self.G = nn.Sequential(nn.Linear(decoder_size, 1), nn.Sigmoid())
initializer(self)
self._bleu = BLEU(
exclude_indices=set(self.SYMBOL)) if use_bleu else None
def __init__(self,
source_embedder: TextFieldEmbedder,
target_embedder: TextFieldEmbedder,
max_steps: int,
encoder: Encoder,
decoder: Decoder,
hidden_size: int,
vocab: Vocabulary,
teacher_force_ratio: float,
regularizer: RegularizerApplicator = None) -> None:
super().__init__(vocab, regularizer)
# TODO: Workon BeamSearch, try to switch to OpenNMT BeamSearch but implement our own beamsearch first
self.max_steps = max_steps
self.hidden_size = hidden_size
self.source_embedder = source_embedder
self.target_embedder = target_embedder
self.encoder = encoder
self.decoder = decoder
self.teacher_force_ratio = teacher_force_ratio
self.decoder.add_vocab(self.vocab)
self.padding_idx = self.vocab.get_token_index(DEFAULT_PADDING_TOKEN)
self.start_idx = self.vocab.get_token_index(START_SYMBOL)
self.end_idx = self.vocab.get_token_index(END_SYMBOL)
self.unk_idx = self.vocab.get_token_index(DEFAULT_OOV_TOKEN)
self.beam = BeamSearch(self.end_idx,
max_steps=self.max_steps,
beam_size=5)
self.criterion = CrossEntropyLoss(ignore_index=self.padding_idx)
def generate(self, code: torch.Tensor, label: torch.Tensor,
max_length: int, beam_size: int,
lp_alpha: float) -> torch.Tensor:
start_index = self.vocab.get_token_index('<s>')
end_index = self.vocab.get_token_index('</s>')
beam_search = BeamSearch(end_index=end_index,
max_steps=max_length,
beam_size=beam_size,
per_node_beam_size=3)
batch_size = code.shape[0]
start_predictions = (
torch.empty(batch_size).to(label).fill_(start_index))
zero_state = code.new_zeros(batch_size,
self._generator._module.hidden_size)
start_state = {
'h': zero_state,
'c': zero_state,
'code': code,
'label': label,
'length': label.new_ones(batch_size),
'length_alpha': (code.new_empty(batch_size).fill_(lp_alpha))
}
all_predictions, last_log_probs = beam_search.search(
start_predictions=start_predictions,
start_state=start_state,
step=self.beam_search_step)
return all_predictions
def __init__(self,
vocab: Vocabulary,
token_embedder: TextFieldEmbedder,
document_encoder: Seq2VecEncoder,
utterance_encoder: Seq2VecEncoder,
context_encoder: Seq2SeqEncoder,
beam_size: int = None,
max_decoding_steps: int = 50,
scheduled_sampling_ratio: float = 0.,
use_bleu: bool = True) -> None:
super(MultiTurnHred, self).__init__(vocab)
self._scheduled_sampling_ratio = scheduled_sampling_ratio
# We need the start symbol to provide as the input at the first timestep of decoding, and
# end symbol as a way to indicate the end of the decoded sequence.
self._start_index = self.vocab.get_token_index(START_SYMBOL)
self._end_index = self.vocab.get_token_index(END_SYMBOL)
if use_bleu:
pad_index = self.vocab.get_token_index(self.vocab._padding_token) # pylint: disable=protected-access
self._bleu = BLEU(exclude_indices={pad_index, self._end_index, self._start_index})
else:
self._bleu = None
# At prediction time, we use a beam search to find the most likely sequence of target tokens.
self._beam_size = beam_size or 1
self._max_decoding_steps = max_decoding_steps
self._beam_search = BeamSearch(self._end_index, max_steps=max_decoding_steps, beam_size=self._beam_size)
# At prediction time, we use a beam search to find the most likely sequence of target tokens.
self._max_decoding_steps = max_decoding_steps
# Dense embedding of word level tokens.
self._token_embedder = token_embedder
# Document word level encoder.
self._document_encoder = document_encoder
# Dialogue word level encoder.
self._utterance_encoder = utterance_encoder
# Sentence level encoder.
self._context_encoder = context_encoder
num_classes = self.vocab.get_vocab_size()
document_output_dim = self._document_encoder.get_output_dim()
utterance_output_dim = self._utterance_encoder.get_output_dim()
context_output_dim = self._context_encoder.get_output_dim()
decoder_output_dim = utterance_output_dim
decoder_input_dim = token_embedder.get_output_dim() + document_output_dim + context_output_dim
# We'll use an LSTM cell as the recurrent cell that produces a hidden state
# for the decoder at each time step.
# TODO (pradeep): Do not hardcode decoder cell type.
self._decoder_cell = GRUCell(decoder_input_dim, decoder_output_dim)
# We project the hidden state from the decoder into the output vocabulary space
# in order to get log probabilities of each target token, at each time step.
self._output_projection_layer = Linear(decoder_output_dim, num_classes)
def test_greedy_search(self):
beam_search = BeamSearch(self.end_index, beam_size=1)
expected_top_k = np.array([[1, 2, 3, 4, 5]])
expected_log_probs = np.log(np.array([0.4]))
self._check_results(expected_top_k=expected_top_k,
expected_log_probs=expected_log_probs,
beam_search=beam_search)
def test_beam_search_matches_greedy(self):
model = self.trained_model
state = model._states["xintent"]
beam_search = BeamSearch(model._end_index,
max_steps=model._max_decoding_steps,
beam_size=1)
final_encoder_output = self.get_sample_encoded_output()
batch_size = final_encoder_output.size()[0]
start_predictions = final_encoder_output.new_full(
(batch_size, ), fill_value=model._start_index, dtype=torch.long)
start_state = {"decoder_hidden": final_encoder_output}
greedy_prediction = model.greedy_predict(
final_encoder_output=final_encoder_output,
target_embedder=state.embedder,
decoder_cell=state.decoder_cell,
output_projection_layer=state.output_projection_layer,
)
greedy_tokens = model.decode_all(greedy_prediction)
(beam_predictions, _) = beam_search.search(start_predictions,
start_state,
state.take_step)
beam_prediction = beam_predictions[0]
beam_tokens = model.decode_all(beam_prediction)
assert beam_tokens == greedy_tokens
def __init__(
self,
vocab: Vocabulary,
decoder_net: DecoderNet,
max_decoding_steps: int,
target_embedder: Embedding,
target_namespace: str = "tokens",
tie_output_embedding: bool = False,
scheduled_sampling_ratio: float = 0,
label_smoothing_ratio: Optional[float] = None,
beam_size: int = 4,
tensor_based_metric: Metric = None,
token_based_metric: Metric = None,
) -> None:
super().__init__(target_embedder)
self._vocab = vocab
# Decodes the sequence of encoded hidden states into e new sequence of hidden states.
self._decoder_net = decoder_net
self._max_decoding_steps = max_decoding_steps
self._target_namespace = target_namespace
self._label_smoothing_ratio = label_smoothing_ratio
# At prediction time, we use a beam search to find the most likely sequence of target tokens.
# We need the start symbol to provide as the input at the first timestep of decoding, and
# end symbol as a way to indicate the end of the decoded sequence.
self._start_index = self._vocab.get_token_index(
START_SYMBOL, self._target_namespace)
self._end_index = self._vocab.get_token_index(END_SYMBOL,
self._target_namespace)
self._beam_search = BeamSearch(self._end_index,
max_steps=max_decoding_steps,
beam_size=beam_size)
target_vocab_size = self._vocab.get_vocab_size(self._target_namespace)
if self.target_embedder.get_output_dim(
) != self._decoder_net.target_embedding_dim:
raise ConfigurationError(
"Target Embedder output_dim doesn't match decoder module's input."
)
# We project the hidden state from the decoder into the output vocabulary space
# in order to get log probabilities of each target token, at each time step.
self._output_projection_layer = Linear(
self._decoder_net.get_output_dim(), target_vocab_size)
if tie_output_embedding:
if self._output_projection_layer.weight.shape != self.target_embedder.weight.shape:
raise ConfigurationError(
"Can't tie embeddings with output linear layer, due to shape mismatch"
)
self._output_projection_layer.weight = self.target_embedder.weight
# These metrics will be updated during training and validation
self._tensor_based_metric = tensor_based_metric
self._token_based_metric = token_based_metric
self._scheduled_sampling_ratio = scheduled_sampling_ratio
def __init__(
self,
word_embedder: TextFieldEmbedder,
attribute_embedder: Embedding,
content_encoder: Seq2SeqEncoder,
vocab: Vocabulary,
max_decoding_steps: int = 20,
beam_size: int = None,
scheduled_sampling_ratio: float = 0.,
) -> None:
super().__init__(vocab)
self.scheduled_sampling_ratio = scheduled_sampling_ratio
# We need the start symbol to provide as the input at the first timestep of decoding, and
# end symbol as a way to indicate the end of the decoded sequence.
self.start_index = self.vocab.get_token_index(START_SYMBOL, 'tokens')
self.end_index = self.vocab.get_token_index(END_SYMBOL, 'tokens')
# TODO: not sure if we need this
self.bleu = None
# At prediction time, we use a beam search to find the most likely sequence of target tokens.
beam_size = beam_size or 1
self.max_decoding_steps = max_decoding_steps
self.beam_search = BeamSearch(self.end_index,
max_steps=max_decoding_steps,
beam_size=beam_size)
# Dense embedding of source and target vocab tokens and attribute.
self.word_embedder = word_embedder
self.attribute_embedder = attribute_embedder
# Encodes the sequence of source embeddings into a sequence of hidden states.
self.content_encoder = content_encoder
num_classes = self.vocab.get_vocab_size('tokens')
# TODO: not sure if we need this
self.attention = None
# Dense embedding of vocab words in the target space.
embedding_dim = word_embedder.get_output_dim()
self.target_embedder = Embedding(num_classes, embedding_dim)
# Decoder output dim needs to be the same as the encoder output dim since we initialize the
# hidden state of the decoder with the final hidden state of the encoder.
self.encoder_output_dim = self.content_encoder.get_output_dim(
) + embedding_dim
self.decoder_output_dim = self.encoder_output_dim
self.decoder_input_dim = embedding_dim
self.decoder_cell = LSTMCell(self.decoder_input_dim,
self.decoder_output_dim)
self.output_projection_layer = Linear(self.decoder_output_dim,
num_classes)
def test_catch_bad_config(self):
"""
If `per_node_beam_size` (which defaults to `beam_size`) is larger than
the size of the target vocabulary, `BeamSearch.search` should raise
a ConfigurationError.
"""
beam_search = BeamSearch(self.end_index, beam_size=20)
with pytest.raises(ConfigurationError):
self._check_results(beam_search=beam_search)
def __init__(
self,
model_name: str,
vocab: Vocabulary,
indexer: PretrainedTransformerIndexer = None,
max_decoding_steps: int = 140,
beam_size: int = 4,
encoder: Seq2SeqEncoder = None,
):
"""
# Parameters
model_name : `str`, required
Name of the pre-trained BART model to use. Available options can be found in
`transformers.models.bart.modeling_bart.BART_PRETRAINED_MODEL_ARCHIVE_MAP`.
vocab : `Vocabulary`, required
Vocabulary containing source and target vocabularies.
indexer : `PretrainedTransformerIndexer`, optional (default = `None`)
Indexer to be used for converting decoded sequences of ids to to sequences of tokens.
max_decoding_steps : `int`, optional (default = `128`)
Number of decoding steps during beam search.
beam_size : `int`, optional (default = `5`)
Number of beams to use in beam search. The default is from the BART paper.
encoder : `Seq2SeqEncoder`, optional (default = `None`)
Encoder to used in BART. By default, the original BART encoder is used.
"""
super().__init__(vocab)
self.bart = BartForConditionalGeneration.from_pretrained(model_name)
self._indexer = indexer or PretrainedTransformerIndexer(
model_name, namespace="tokens")
self._start_id = self.bart.config.bos_token_id # CLS
self._decoder_start_id = self.bart.config.decoder_start_token_id or self._start_id
self._end_id = self.bart.config.eos_token_id # SEP
self._pad_id = self.bart.config.pad_token_id # PAD
self._max_decoding_steps = max_decoding_steps
self._beam_search = BeamSearch(self._end_id,
max_steps=max_decoding_steps,
beam_size=beam_size or 1)
self._rouge = ROUGE(
exclude_indices={self._start_id, self._pad_id, self._end_id})
self._bleu = BLEU(
exclude_indices={self._start_id, self._pad_id, self._end_id})
# Replace bart encoder with given encoder. We need to extract the two embedding layers so that
# we can use them in the encoder wrapper
if encoder is not None:
assert (encoder.get_input_dim() == encoder.get_output_dim() ==
self.bart.config.hidden_size)
self.bart.model.encoder = _BartEncoderWrapper(
encoder,
self.bart.model.encoder.embed_tokens,
self.bart.model.encoder.embed_positions,
)
def test_early_stopping(self):
"""
Checks case where beam search will reach `max_steps` before finding end tokens.
"""
beam_search = BeamSearch(self.end_index, beam_size=3, max_steps=3)
expected_top_k = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
expected_log_probs = np.log(np.array([0.4, 0.3, 0.2])) # pylint: disable=assignment-from-no-return
self._check_results(expected_top_k=expected_top_k,
expected_log_probs=expected_log_probs,
beam_search=beam_search)
def __init__(self,
vocab: Vocabulary,
source_embedder: TextFieldEmbedder,
embedding_dropout: float,
encoder: Seq2VecEncoder,
max_decoding_steps: int,
beam_size: int = 10,
target_names: List[str] = None,
target_namespace: str = "tokens",
target_embedding_dim: int = None) -> None:
target_names = target_names or ["xintent", "xreact", "oreact"]
super(Event2Mind, self).__init__(vocab)
# Note: The original tweaks the embeddings for "personx" to be the mean
# across the embeddings for "he", "she", "him" and "her". Similarly for
# "personx's" and so forth. We could consider that here as a well.
self._source_embedder = source_embedder
self._embedding_dropout = nn.Dropout(embedding_dropout)
self._encoder = encoder
self._max_decoding_steps = max_decoding_steps
self._target_namespace = target_namespace
# We need the start symbol to provide as the input at the first timestep of decoding, and
# end symbol as a way to indicate the end of the decoded sequence.
self._start_index = self.vocab.get_token_index(START_SYMBOL, self._target_namespace)
self._end_index = self.vocab.get_token_index(END_SYMBOL, self._target_namespace)
# Warning: The different decoders share a vocabulary! This may be
# counterintuitive, but consider the case of xreact and oreact. A
# reaction of "happy" could easily apply to both the subject of the
# event and others. This could become less appropriate as more decoders
# are added.
num_classes = self.vocab.get_vocab_size(self._target_namespace)
# Decoder output dim needs to be the same as the encoder output dim since we initialize the
# hidden state of the decoder with that of the final hidden states of the encoder.
self._decoder_output_dim = self._encoder.get_output_dim()
target_embedding_dim = target_embedding_dim or self._source_embedder.get_output_dim()
self._states: Dict[str, StateDecoder] = {}
for name in target_names:
self._states[name] = StateDecoder(
name,
self,
num_classes,
target_embedding_dim,
self._decoder_output_dim
)
self._beam_search = BeamSearch(
self._end_index,
beam_size=beam_size,
max_steps=max_decoding_steps
)
def __init__(
self,
vocab: Vocabulary,
decoder_net: DecoderNet,
max_decoding_steps: int,
target_embedder: Embedding,
target_namespace: str = "tokens",
tie_output_embedding: bool = False,
scheduled_sampling_ratio: float = 0,
label_smoothing_ratio: Optional[float] = None,
beam_size: int = 4,
tensor_based_metric: Metric = None,
token_based_metric: Metric = None,
) -> None:
super().__init__(target_embedder)
self._vocab = vocab
self._decoder_net = decoder_net
self._max_decoding_steps = max_decoding_steps
self._target_namespace = target_namespace
self._label_smoothing_ratio = label_smoothing_ratio
self._start_index = self._vocab.get_token_index(
START_SYMBOL, self._target_namespace)
self._end_index = self._vocab.get_token_index(END_SYMBOL,
self._target_namespace)
self._beam_search = BeamSearch(self._end_index,
max_steps=max_decoding_steps,
beam_size=beam_size)
target_vocab_size = self._vocab.get_vocab_size(self._target_namespace)
if self.target_embedder.get_output_dim(
) != self._decoder_net.target_embedding_dim:
raise ConfigurationError(
"Target Embedder output_dim doesn't match decoder module's input."
)
self._output_projection_layer = Linear(
self._decoder_net.get_output_dim(), target_vocab_size)
if tie_output_embedding:
if self._output_projection_layer.weight.shape != self.target_embedder.weight.shape:
raise ConfigurationError(
"Can't tie embeddings with output linear layer, due to shape mismatch"
)
self._output_projection_layer.weight = self.target_embedder.weight
self._tensor_based_metric = tensor_based_metric
self._token_based_metric = token_based_metric
self._scheduled_sampling_ratio = scheduled_sampling_ratio
def setUp(self):
super(BeamSearchTest, self).setUp()
self.end_index = transition_probabilities.size()[0] - 1
self.beam_search = BeamSearch(self.end_index,
max_steps=10,
beam_size=3)
# This is what the top k should look like for each item in the batch.
self.expected_top_k = np.array([[1, 2, 3, 4, 5], [2, 3, 4, 5, 5],
[3, 4, 5, 5, 5]])
# This is what the log probs should look like for each item in the batch.
self.expected_log_probs = np.log(np.array([0.4, 0.3, 0.2])) # pylint: disable=assignment-from-no-return
def test_k_val(self, k_val):
with pytest.raises(ValueError):
initial_predictions = torch.tensor([0] * 5)
take_step = take_step_with_timestep
beam_size = 3
k_sampler = TopKSampler(k=k_val, with_replacement=True)
top_k, log_probs = BeamSearch(self.end_index,
beam_size=beam_size,
max_steps=10,
sampler=k_sampler).search(
initial_predictions, {},
take_step)
def test_empty_sequences(self):
initial_predictions = torch.LongTensor(
[self.end_index - 1, self.end_index - 1])
beam_search = BeamSearch(self.end_index, beam_size=1)
with pytest.warns(RuntimeWarning, match="Empty sequences predicted"):
predictions, log_probs = beam_search.search(
initial_predictions, {}, take_step_with_timestep)
# predictions hould have shape `(batch_size, beam_size, max_predicted_length)`.
assert list(predictions.size()) == [2, 1, 1]
# log probs hould have shape `(batch_size, beam_size)`.
assert list(log_probs.size()) == [2, 1]
assert (predictions == self.end_index).all()
assert (log_probs == 0).all()
def setup_method(self):
super().setup_method()
self.end_index = transition_probabilities.size()[0] - 1
self.beam_search = BeamSearch(self.end_index,
max_steps=10,
beam_size=3)
# This is what the top k should look like for each item in the batch.
self.expected_top_k = np.array([[1, 2, 3, 4, 5], [2, 3, 4, 5, 5],
[3, 4, 5, 5, 5]])
# This is what the log probs should look like for each item in the batch.
self.expected_log_probs = np.log(np.array([0.4, 0.3, 0.2]))
def __init__(self,
vocab: Vocabulary,
source_embedder: TextFieldEmbedder,
encoder: Seq2SeqEncoder,
attention: Attention,
max_decoding_steps: int,
beam_size: int = None,
target_namespace: str = "tokens",
target_embedding_dim: int = None,
scheduled_sampling_ratio: float = 0.,
projection_dim: int = None,
use_coverage: bool = False,
coverage_loss_weight: float = None) -> None:
super(PointerGeneratorNetwork, self).__init__(vocab)
self._target_namespace = target_namespace
self._start_index = self.vocab.get_token_index(START_SYMBOL, target_namespace)
self._end_index = self.vocab.get_token_index(END_SYMBOL, target_namespace)
self._unk_index = self.vocab.get_token_index(DEFAULT_OOV_TOKEN, target_namespace)
self._vocab_size = self.vocab.get_vocab_size(target_namespace)
# Encoder
self._source_embedder = source_embedder
self._encoder = encoder
self._encoder_output_dim = self._encoder.get_output_dim()
# Decoder
self._target_embedding_dim = target_embedding_dim or source_embedder.get_output_dim()
self._num_classes = self.vocab.get_vocab_size(target_namespace)
self._target_embedder = Embedding(self._num_classes, self._target_embedding_dim)
self._decoder_input_dim = self._encoder_output_dim + self._target_embedding_dim
self._decoder_output_dim = self._encoder_output_dim
self._decoder_cell = LSTMCell(self._decoder_input_dim, self._decoder_output_dim)
self._projection_dim = projection_dim or self._source_embedder.get_output_dim()
self._hidden_projection_layer = Linear(self._decoder_output_dim, self._projection_dim)
self._output_projection_layer = Linear(self._projection_dim, self._num_classes)
self._p_gen_layer = Linear(self._decoder_output_dim * 3 + self._decoder_input_dim, 1)
self._attention = attention
self._use_coverage = use_coverage
self._coverage_loss_weight = coverage_loss_weight
self._eps = 1e-31
# Decoding
self._scheduled_sampling_ratio = scheduled_sampling_ratio
self._max_decoding_steps = max_decoding_steps
self._beam_search = BeamSearch(self._end_index, max_steps=max_decoding_steps, beam_size=beam_size or 1)
def __init__(self,
vocab: Vocabulary,
text_field_embedder: TextFieldEmbedder,
sentence_encoder: Seq2VecEncoder,
claim_encoder: Seq2SeqEncoder,
attention: Attention,
max_steps: int = 100,
beam_size: int = 5,
beta: float = 1.0) -> None:
super(Seq2SeqClaimRank, self).__init__(vocab)
self.text_field_embedder = text_field_embedder
self.sentence_encoder = sentence_encoder
self.claim_encoder = TimeDistributed(claim_encoder) # Handles additional sequence dim
self.claim_encoder_dim = claim_encoder.get_output_dim()
self.attention = attention
self.decoder_embedding_dim = text_field_embedder.get_output_dim()
self.max_steps = max_steps
self.beam_size = beam_size
self.beta = beta
# self.target_embedder = torch.nn.Embedding(vocab.get_vocab_size(), decoder_embedding_dim)
# Since we are using the sentence encoding as the initial hidden state to the decoder, the
# decoder hidden dim must match the sentence encoder hidden dim.
self.decoder_output_dim = sentence_encoder.get_output_dim()
self.decoder_0_cell = torch.nn.LSTMCell(self.decoder_embedding_dim + self.claim_encoder_dim,
self.decoder_output_dim)
self.decoder_1_cell = torch.nn.LSTMCell(self.decoder_output_dim,
self.decoder_output_dim)
# When projecting out we will use attention to combine claim embeddings into a single
# context embedding, this will be concatenated with the decoder cell output before being
# fed to the projection layer. Hence the expected input size is:
# decoder output dim + claim encoder output dim
projection_input_dim = self.decoder_output_dim + self.claim_encoder_dim
self.output_projection_layer = torch.nn.Linear(projection_input_dim,
vocab.get_vocab_size())
self._start_index = self.vocab.get_token_index('<s>')
self._end_index = self.vocab.get_token_index('</s>')
self.beam_search = BeamSearch(self._end_index, max_steps=max_steps, beam_size=beam_size)
pad_index = vocab.get_token_index(vocab._padding_token)
self.bleu = BLEU(exclude_indices={pad_index, self._start_index, self._end_index})
self.avg_reconstruction_loss = Average()
self.avg_claim_scoring_loss = Average()
def __init__(
self,
vocabulary: Vocabulary,
image_feature_size: int,
embedding_size: int,
hidden_size: int,
attention_projection_size: int,
max_caption_length: int = 20,
beam_size: int = 1,
) -> None:
super().__init__()
self._vocabulary = vocabulary
self.image_feature_size = image_feature_size
self.embedding_size = embedding_size
self.hidden_size = hidden_size
self.attention_projection_size = attention_projection_size
# Short hand variable names for convenience
vocab_size = vocabulary.get_vocab_size()
self._pad_index = vocabulary.get_token_index("@@UNKNOWN@@")
self._boundary_index = vocabulary.get_token_index("@@BOUNDARY@@")
self._embedding_layer = nn.Embedding(
vocab_size, embedding_size, padding_idx=self._pad_index
)
self._updown_cell = UpDownCell(
image_feature_size, embedding_size, hidden_size, attention_projection_size
)
self._output_layer = nn.Linear(hidden_size, vocab_size)
self._log_softmax = nn.LogSoftmax(dim=1)
# We use beam search to find the most likely caption during inference.
self._beam_size = beam_size
self._beam_search = BeamSearch(
self._boundary_index,
max_steps=max_caption_length,
beam_size=beam_size,
per_node_beam_size=beam_size // 2,
)
self._max_caption_length = max_caption_length
def __init__(
self,
model_name: str,
vocab: Vocabulary,
indexer: PretrainedTransformerIndexer = None,
max_decoding_steps: int = 140,
beam_size: int = 4,
encoder: Seq2SeqEncoder = None,
):
super().__init__(vocab)
self.bart = BartForConditionalGeneration.from_pretrained(model_name)
self._indexer = indexer or PretrainedTransformerIndexer(
model_name, namespace="tokens")
self._start_id = self.bart.config.bos_token_id # CLS
self._decoder_start_id = self.bart.config.decoder_start_token_id or self._start_id
self._end_id = self.bart.config.eos_token_id # SEP
self._pad_id = self.bart.config.pad_token_id # PAD
self._max_decoding_steps = max_decoding_steps
self._beam_search = BeamSearch(self._end_id,
max_steps=max_decoding_steps,
beam_size=beam_size or 1)
self._rouge = ROUGE(
exclude_indices={self._start_id, self._pad_id, self._end_id})
self._bleu = BLEU(
exclude_indices={self._start_id, self._pad_id, self._end_id})
# Replace bart encoder with given encoder. We need to extract the two embedding layers so that
# we can use them in the encoder wrapper
if encoder is not None:
assert (encoder.get_input_dim() == encoder.get_output_dim() ==
self.bart.config.hidden_size)
self.bart.model.encoder = _BartEncoderWrapper(
encoder,
self.bart.model.encoder.embed_tokens,
self.bart.model.encoder.embed_positions,
)
def __init__(self,
vocab: Vocabulary,
pretrained_model_path,
beam_size=5,
max_decoding_steps=140,
indexer=None):
super().__init__(vocab)
self.plm = MT5ForConditionalGeneration.from_pretrained(pretrained_model_path)
self._indexer = indexer or PretrainedTransformerIndexer(pretrained_model_path, namespace="tokens")
##
self._start_id = self.plm.config.decoder_start_token_id
##
self._end_id = self.plm.config.eos_token_id #
self._decoder_start_id = self.plm.config.decoder_start_token_id
self._end_id = self.plm.config.eos_token_id #
self._pad_id = self.plm.config.pad_token_id #
self._beam_search = BeamSearch(
self._end_id, max_steps=max_decoding_steps, beam_size=beam_size or 1
)
self._rouge = ROUGE(exclude_indices={self._start_id, self._pad_id, self._end_id})
self._bleu = BLEU(exclude_indices={self._start_id, self._pad_id, self._end_id})
def test_top_k_search(self):
initial_predictions = torch.tensor([0] * 5)
beam_size = 3
take_step = take_step_with_timestep
k_sampler = TopKSampler(k=5, with_replacement=True)
top_k, log_probs = BeamSearch(self.end_index,
beam_size=beam_size,
max_steps=10,
sampler=k_sampler).search(
initial_predictions, {}, take_step)
beam_size = beam_size or 1
batch_size = 5
# top_p should be shape `(batch_size, beam_size, max_predicted_length)`.
assert list(top_k.size())[:-1] == [batch_size, beam_size]
assert ((0 <= top_k) & (top_k <= 5)).all()
# log_probs should be shape `(batch_size, beam_size, max_predicted_length)`.
assert list(log_probs.size()) == [batch_size, beam_size]
def __init__(self, vocab: Vocabulary, max_timesteps: int = 50, encoder_size: int = 14, encoder_dim: int = 512,
embedding_dim: int = 64, attention_dim: int = 64, decoder_dim: int = 64, beam_size: int = 3, teacher_forcing: bool = True) -> None:
super().__init__(vocab)
self._max_timesteps = max_timesteps
self._vocab_size = self.vocab.get_vocab_size()
self._start_index = self.vocab.get_token_index(START_SYMBOL)
self._end_index = self.vocab.get_token_index(END_SYMBOL)
# POSSIBLE CHANGE LATER
self._pad_index = self.vocab.get_token_index('@@PADDING@@')
self._encoder_size = encoder_size
self._encoder_dim = encoder_dim
self._embedding_dim = embedding_dim
self._attention_dim = attention_dim
self._decoder_dim = decoder_dim
self._beam_size = beam_size
self._teacher_forcing = teacher_forcing
self._init_h = nn.Linear(self._encoder_dim, self._decoder_dim)
self._init_c = nn.Linear(self._encoder_dim, self._decoder_dim)
self._resnet = torchvision.models.resnet18()
modules = list(self._resnet.children())[:-2]
self._encoder = nn.Sequential(
*modules,
nn.AdaptiveAvgPool2d(self._encoder_size)
)
self._decoder = ImageCaptioningDecoder(self._vocab_size, self._encoder_dim, self._embedding_dim, self._attention_dim, self._decoder_dim)
self.beam_search = BeamSearch(self._end_index, self._max_timesteps, self._beam_size)
self._bleu = BLEU(exclude_indices={self._start_index, self._end_index, self._pad_index})
self._exprate = Exprate(self._end_index)
def __init__(self,
vocab: Vocabulary,
encoder: Encoder,
decoder: CaptioningDecoder,
max_timesteps: int = 75,
teacher_forcing: bool = True,
scheduled_sampling_ratio: float = 1,
beam_size: int = 10) -> None:
super().__init__(vocab)
self._start_index = self.vocab.get_token_index(START_SYMBOL)
self._end_index = self.vocab.get_token_index(END_SYMBOL)
self._pad_index = self.vocab.get_token_index('@@PADDING@@')
self._max_timesteps = max_timesteps
self._teacher_forcing = teacher_forcing
self._scheduled_sampling_ratio = scheduled_sampling_ratio
self._beam_size = beam_size
self._encoder = encoder
self._decoder = decoder
self._init_h = nn.Linear(self._encoder.get_output_dim(),
self._decoder.get_input_dim())
self._init_c = nn.Linear(self._encoder.get_output_dim(),
self._decoder.get_input_dim())
self.beam_search = BeamSearch(self._end_index, self._max_timesteps,
self._beam_size)
self._bleu = BLEU(exclude_indices={
self._start_index, self._end_index, self._pad_index
})
self._exprate = Exprate(self._end_index, self.vocab)
self._attention_weights = None
def __init__(self,
vocab: Vocabulary,
source_embedder: TextFieldEmbedder,
target_embedder: Embedding,
encoder: Seq2SeqEncoder,
max_decoding_steps: int,
decoding_dim: int,
feedforward_hidden_dim: int,
num_layers: int,
num_attention_heads: int,
use_positional_encoding: bool = True,
positional_encoding_max_steps: int = 5000,
dropout_prob: float = 0.1,
residual_dropout_prob: float = 0.2,
attention_dropout_prob: float = 0.2,
beam_size: int = 1,
target_namespace: str = "tokens",
label_smoothing_ratio: Optional[float] = None,
initializer: Optional[InitializerApplicator] = None) -> None:
super(SequenceTransformer, self).__init__(vocab)
self._target_namespace = target_namespace
self._label_smoothing_ratio = label_smoothing_ratio
self._start_index = self.vocab.get_token_index(START_SYMBOL,
self._target_namespace)
self._end_index = self.vocab.get_token_index(END_SYMBOL,
self._target_namespace)
self._token_based_metric = TokenSequenceAccuracy()
# Beam Search
self._max_decoding_steps = max_decoding_steps
self._beam_search = BeamSearch(self._end_index,
max_steps=max_decoding_steps,
beam_size=beam_size)
# Encoder
self._encoder = encoder
# Vocabulary and embedder
self._source_embedder = source_embedder
self._target_embedder = target_embedder
target_vocab_size = self.vocab.get_vocab_size(self._target_namespace)
assert target_vocab_size == self._target_embedder.num_embeddings
target_embedding_dim = self._target_embedder.get_output_dim()
self._decoding_dim = decoding_dim
# Sequence Decoder Features
self._output_projection_layer = Linear(self._decoding_dim,
target_vocab_size)
self._decoder = Decoder(
num_layers=num_layers,
decoding_dim=decoding_dim,
target_embedding_dim=target_embedding_dim,
feedforward_hidden_dim=feedforward_hidden_dim,
num_attention_heads=num_attention_heads,
use_positional_encoding=use_positional_encoding,
positional_encoding_max_steps=positional_encoding_max_steps,
dropout_prob=dropout_prob,
residual_dropout_prob=residual_dropout_prob,
attention_dropout_prob=attention_dropout_prob)
# Parameter checks and cleanup
if self._target_embedder.get_output_dim(
) != self._decoder.target_embedding_dim:
raise ConfigurationError(
"Target Embedder output_dim doesn't match decoder module's input."
)
#
if self._encoder.get_output_dim() != self._decoder.get_output_dim():
raise ConfigurationError(
f"Encoder output dimension {self._encoder.get_output_dim()} should be"
f" equal to decoder dimension {self._self_attention.get_output_dim()}."
)
if initializer:
initializer(self)
# Print the model
print(self)
