Python 索引到 torch 张量，沿轴具有可变长度索引

我试图根据一个语言模型计算一组标记化单词的单词概率，我需要一些花哨的索引.

我的输入，用玩具示例说明如下:

• token_list:n_words x max_tokenization_length(例如，三个字，其中最大标记化长度为3)
• pxhs:n_words x(max_tokenization_length +1)x|词汇|，(例如，三个单词，四组logit用于3 + 1令牌，维度1000 vocab)
• Next_word_token_ids:构成新词的令牌列表(例如，所有以空格字符开头的标记).

pxhs = torch.rand((3,4,1000))

pad_token_id = tokenizer.pad_token_id
word_token_list = [
    [120, pad_token_id, pad_token_id],
    [131, 132, pad_token_id],
    [140, 141, 142],
]

new_word_token_ids = [0,1,2,3,5]

期望的输出是长度为3的单词概率列表，计算如下:

word 1: pxhs[0, 0, 120] * pxhs[0, 1, new_word_token_ids].sum()
word 2: pxhs[1, 0, 131] * pxhs[1, 1, 132] * pxhs[1, 2, new_word_token_ids].sum()
word 3: pxhs[2, 0, 140] * pxhs[2, 1, 141] * pxhs[2, 2, 142] * pxhs[2, 3, new_word_token_ids].sum()

在实践中，我想通过用新词token id替换第一个pad_token_id来索引，然后什么都不做(这不能作为索引，只是举例说明):

actual_idx = [
    [[120], new_word_token_ids, [None], [None]],
    [[131], [132], new_word_token_ids, [None]],
    [[140], [142], [143], new_word_token_ids],
]

我写了一个非常慢的函数，它可以做到这一点:

all_word_probs = []
for word_tokens, word_probs in zip(token_list, pxhs):
    counter=0
    p_word=1
    while (counter < len(word_tokens) and 
            word_tokens[counter] != tokenizer.pad_token_id):
        p_word = p_word * word_probs[counter, word_tokens[counter]]
        counter+=1
    new_word_prob = word_probs[counter, new_word_tokens].sum()
    p_word = p_word * new_word_prob
    all_word_probs.append(p_word)

我需要更快的东西，提前感谢你的帮助！