reproducing greedy decoding performance

[freshpearYoon]

Hi, I am trying to reproduce greedy decoding performance, and it seems that current code does not provide option for that. Am I right ? How should I run greedy decoding?

[1jia1-oneplusone]

You only need to change a bit code in the sample function.

The raw code belike: (in the if use_cd: branch, next_token_logits should be replaced with cd_logits)

next_token_scores = logits_processor(input_ids, next_token_logits)
next_token_scores = logits_warper(input_ids, next_token_scores)
probs = nn.functional.softmax(next_token_scores, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)

replace it with:

next_token_scores = logits_processor(input_ids, next_token_logits)
next_tokens = torch.argmax(next_token_scores, dim=-1)

The difference is that the previous one use torch.multinomial to randomly choose a token, while the latter just use torch.argmax to select the token with maximum logit.