from transformers import BertTokenizer, BertForMaskedLM
import torch


# 指定模型方式
def dm03():
    # 加载模型和分词器
    tokenizer = BertTokenizer.from_pretrained("../models/chinese-bert-wwm")
    model = BertForMaskedLM.from_pretrained("../models/chinese-bert-wwm")
    # 准备数据
    input = tokenizer.encode_plus("我想明天去[MASK]家吃饭", return_tensors='pt')
    print("input-->", input)
    # 把数据送给模型
    output = model(**input)
    # 打印结果
    # print("output-->", output)
    # print("output-->", output.last_hidden_state.shape) # [1, 11, 768]
    # print("output-->", output.pooler_output.shape) # [1, 768]

    # [1, 11, 21128]：1句话，11个token，每个token用一个21128的向量表示
    # 含义是：需要基于MASK的位置，从21128个词表中，还原到预测的结果token（词）
    print("output.logits-->", output.logits, output.logits.shape)
    # [11, 21128] -> output.logits[0]
    # [21128] -> output.logits[0][6]
    # print(output.logits[0][6].shape)
    # item：由张量转标量
    index = torch.argmax(output.logits[0][6]).item()
    print("index-->", index)
    # 把下标index从词表中还原成具体的token
    token = tokenizer.convert_ids_to_tokens([index])
    print("token-->", token)


if __name__ == '__main__':
    dm03()