I am in iterator hell.

for tok, predicted_toks, log_prob, prob in list(zip(
    token_list(target_ids[0]),
    [
        zip(topk_log_probs, topk_probs, token_list(topk_tokens))
        for topk_log_probs, topk_probs, topk_tokens
        in zip(
            greedy_log_probs[0].tolist(),
            greedy_probs[0].tolist(),
            greedy_tokens[0],
        )
    ],
    target_log_probs[0].tolist(),
    target_probs[0].tolist(),
))[max_length - stride if begin_loc > 0 else 0:]:
    tokens.append({
        'token': tok,
        'predicted_tokens': [{
            'token': tok,
            'log_prob': log_prob,
            'prob': prob,
        } for log_prob, prob, tok in predicted_toks],
        'log_prob': log_prob,
        'prob': prob,
    })