I am trying to obtain output embeddings by the transformer, along with some generated text summaries. So far, I have tried the following.
config = {'num_return_sequences': 3, 'num_beams': 4, ...}
translated = model.generate(tokenized_input, **config, output_scores=True, return_dict_in_generate=True, output_hidden_states=True)
d_model = model.config.d_model
# Embeddings of the final output layer, for every generation step
output_embeddings_per_step = [step[0] for step in translated['decoder_hidden_states']]
# zip(*output_embeddings_per_step) : creates (pos1_embedding, pos2_embedding, ...) for each sample in sampled sentences
current_output_embeddings = [torch.cat(sample, dim=0).reshape(-1).split(d_model) for sample in zip(*output_embeddings_per_step)]
# len(current_output_embeddings)=4, whereas len(translated['sequences'])=3
However, as I have num_return_sequences < num_beams. I got embeddings for more sequences than intended (incl. those not picked as well). Is there a way to identify which of the sequences as picked by num_return_sequences?
Thank you.