INNER CODE UNIT · Python

y_pred

mlwithme/BertWithPretrained · Tasks/TaskForChineseNER.py:206

    prob, y_pred = torch.max(logits, dim=-1)  # prob, y_pred: [batch_size,src_len-1]
    token_ids = token_ids[1:].transpose(0, 1)  # [ batch_size,src_len-1]， 去掉[cls]
    assert y_pred.shape == token_ids.shape
    labels = []
    probs = []
    for sample in zip(y_pred, token_ids, prob):
        tmp_label, tmp_prob = [], []
        for item in zip(*sample):
            if item[1] == SEP_IDX:  # 忽略最后一个[SEP]字符
                break
            tmp_label.append(label_list[item[0]])
            tmp_prob.append(item[2].item())
        labels.append(tmp_label)
        probs.append(tmp_prob)
    return labels, probs


def pretty_print(sentences, labels, entities):

View source record →