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,
})