def normalized_scores_with_batch_vocab(
        model_decoder, features, target_ids, k, bsz, beam_size,
        pad_idx, top_k=0, vocab_size_meter=None, start_idx=None,
        end_idx=None, **kwargs):
    """
        Get normalized probabilities (or log probs) from a net's output
        w.r.t. vocab consisting of target IDs in the batch
    """
    if model_decoder.adaptive_softmax is None:
        weight = model_decoder.output_projection.weight
        vocab_ids = torch.unique(
            torch.cat(
                (torch.unique(target_ids), torch.arange(top_k, device=target_ids.device))
            )
        )
        id_map = dict(zip(vocab_ids.tolist(), range(len(vocab_ids))))
        mapped_target_ids = target_ids.cpu().apply_(
            lambda x, id_map=id_map: id_map[x]
        ).to(target_ids.device)
        expanded_target_ids = mapped_target_ids[:, :].repeat(1, k).view(bsz*beam_size*k, -1)
        if start_idx is not None and end_idx is not None:
            expanded_target_ids = expanded_target_ids[start_idx:end_idx, :]
        logits = F.linear(features, weight[vocab_ids, :])
        log_softmax = F.log_softmax(logits, dim=-1, dtype=torch.float32)
        intermed_scores = torch.gather(
            log_softmax[:, :-1, :],
            2,
            expanded_target_ids[:, 1:].unsqueeze(2),
        ).squeeze()
        not_padding = expanded_target_ids[:, 1:] != pad_idx
        intermed_scores *= not_padding.float()
        return intermed_scores
    else:
        raise ValueError("adaptive softmax doesn't work with " +
                         "`normalized_scores_with_batch_vocab()`")
Get normalized probabilities (or log probs) from a net's output
w.r.t. vocab consisting of target IDs in the batch
examples/fast_noisy_channel/noisy_channel_sequence_generator.py