from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel, AutoModelForMaskedLM, \
AutoModelForQuestionAnswering, AutoModelForSeq2SeqLM, AutoModelForTokenClassification, AutoConfig
import torch
# TODO 文本分类任务
def dm01():
# 1.加载预训练模型(分词器，模型)
tokenizer = AutoTokenizer.from_pretrained("../models/chinese_sentiment")
model = AutoModelForSequenceClassification.from_pretrained("../models/chinese_sentiment")
# 2.准备数据
message = '我爱北京天安门'
# 3.对数据进行预处理，得到张量（文本数值化+数值张量化）
# text：文本
# truncation：如果超过MAX_LENGTH，就截断
# max_length：最长长度
# return_tensors：pt 返回PyTorch的张量，默认不返回张量，返回list列表
inputs = tokenizer.encode(text=message,
padding='max_length',
truncation=True,
max_length=20,
return_tensors='pt')
print("inputs-->", inputs)
# 4.把数据送给模型
output = model(inputs)
# 5.打印结果
print("output-->", output)
print("output-->", output.logits) # 未经过softmax的logits值
print("output-->", torch.argmax(output.logits, dim=-1))
# TODO 特征提取任务
def dm02():
# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("../models/bert-base-chinese")
model = AutoModel.from_pretrained("../models/bert-base-chinese")
# 准备文本
message = ['你是谁啊', '人生该如何起头']
# 编码的方法有三个：
# encode：普通编码，就是张量数据（input_ids）
# encode_plus：加强版的编码，会多一些（input_ids，token_type_ids，attention_mask）
# batch_encode_plus：真正意义上多个样本的编码（后面再用）
# 对文本进行数值化+张量化
# input = tokenizer.encode(message, padding="max_length", max_length=20, truncation=True, return_tensors='pt')
# print('input-->', input)
# input_ids：对数据进行数值化+张量化后的结果
# token_type_ids：句子类别 CLS 第一句话 SEP 第二句话 SEP 填充的0
# 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0
# attention_mask：注意力掩码，1需要掩码
# 参数小结：input_ids不能少，attention_mask最好要，token_type_ids可有可无。
input = tokenizer.encode_plus(message, padding="max_length", max_length=20, truncation=True, return_tensors='pt')
print('input-->', input)
# 把文本送给模型
# 写法一：
# input_ids = input['input_ids']
# token_type_ids = input['token_type_ids']
# attention_mask = input['attention_mask']
# result = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
# 写法二：**含义：等价于写法一的key=value的关键字赋值
result = model(**input)
# 打印结果
# print("result-->", result)
# [1,20,768]：1句话，20个token，每个token使用768维度的向量表示
# 小结：last_hidden_state：词向量，每个token一个向量，就是词向量
print("result.last_hidden_state-->", result.last_hidden_state, result.last_hidden_state.shape)
# [1, 768]：1句话，这句话使用一个768维度的向量表示
# pooler_output：句向量，一句话（一段文本）的向量表示
print("result.pooler_output-->", result.pooler_output, result.pooler_output.shape)
# TODO 完形填空任务
def dm03():
# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("../models/chinese-bert-wwm")
model = AutoModelForMaskedLM.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)
# TODO 阅读理解任务
def dm04():
# 模型和分词器
tokenizer = AutoTokenizer.from_pretrained("../models/chinese_pretrain_mrc_roberta_wwm_ext_large")
model = AutoModelForQuestionAnswering.from_pretrained("../models/chinese_pretrain_mrc_roberta_wwm_ext_large")
# 准备数据
context = '我叫张三 我是一个程序员 我的喜好是打篮球'
questions = ['我是谁？', '我是做什么的？', '我的爱好是什么？']
# 遍历：一个问题一个问题地送给预训练模型
for question in questions:
# 对每一个问题进行编码
input = tokenizer.encode_plus(question, context, return_tensors='pt')
print("input-->", input)
output = model(**input)
# print("output-->", output)
print("output.start_logits-->", output.start_logits.shape)
print("output.end_logits-->", output.end_logits.shape)
start = torch.argmax(output.start_logits)
end = torch.argmax(output.end_logits)
print("start-->", start)
print("end-->", end)
print('1-->', input['input_ids'].shape) # 【1,26】
print('2-->', input['input_ids'][0].shape) # 【26】降低维度
# convert_ids_to_tokens：把idx还原成token
answer = tokenizer.convert_ids_to_tokens(input['input_ids'][0][start: end + 1])
print("question:", question, "answer:", answer)
# TODO 文本摘要任务
def dm05():
# 模型和分词器
tokenizer = AutoTokenizer.from_pretrained("../models/distilbart-cnn-12-6")
model = AutoModelForSeq2SeqLM.from_pretrained("../models/distilbart-cnn-12-6")
# 准备数据
text = "BERT is a transformers model pretrained on a large corpus of English data " \
"in a self-supervised fashion. This means it was pretrained on the raw texts " \
"only, with no humans labelling them in any way (which is why it can use lots " \
"of publicly available data) with an automatic process to generate inputs and " \
"labels from those texts. More precisely, it was pretrained with two objectives:Masked " \
"language modeling (MLM): taking a sentence, the model randomly masks 15% of the " \
"words in the input then run the entire masked sentence through the model and has " \
"to predict the masked words. This is different from traditional recurrent neural " \
"networks (RNNs) that usually see the words one after the other, or from autoregressive " \
"models like GPT which internally mask the future tokens. It allows the model to learn " \
"a bidirectional representation of the sentence.Next sentence prediction (NSP): the models" \
" concatenates two masked sentences as inputs during pretraining. Sometimes they correspond to " \
"sentences that were next to each other in the original text, sometimes not. The model then " \
"has to predict if the two sentences were following each other or not."
# 对文本进行编码
# input = tokenizer.encode(text, return_tensors='pt')
# print("input-->", input)
input = tokenizer([text], return_tensors='pt')
# print("input-->", input)
# 摘要生成
output = model.generate(input.input_ids)
# output就是token的下标索引index，需要还原成token
print("output-->", output)
# 还原，解码
print([tokenizer.decode(index, skip_special_tokens=True, clean_up_tokenization_spaces=False)
for index in output])
# TODO NER任务
class AutoModelTokenClassification:
pass
# TODO NER任务（命名实体识别）token级别的分类
def dm06():
# 模型和分词器
tokenizer = AutoTokenizer.from_pretrained("../models/roberta-base-finetuned-cluener2020-chinese")
model = AutoModelForTokenClassification.from_pretrained("../models/roberta-base-finetuned-cluener2020-chinese")
# 导入config
config = AutoConfig.from_pretrained("../models/roberta-base-finetuned-cluener2020-chinese")
# print("config-->", config)
# 准备数据
inputs = tokenizer.encode_plus('我爱北京天安门', return_tensors='pt')
print("inputs-->", inputs)
# 对数据进行编码
output = model(**inputs)
# print("output-->", output)
logits = output.logits
# [1, 9, 32]：1句话，9个token，每个token的NER类别数是32
print("logits-->", logits.shape)
# 把这一句话还原成原始的token
input_tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
# ['[CLS]', '我', '爱', '北', '京', '天', '安', '门', '[SEP]']
print("input_tokens-->", input_tokens)
outputs = []
# 遍历token和所有的value
for token, value in zip(input_tokens, logits[0]):
# 如果是特殊token（CLS、MASK、SEP），应该过滤掉
if token in tokenizer.all_special_tokens:
continue
# print("token-->", token)
# print("value-->", value)
# 取value中的最大值，也就是32个类别中最有可能的类别
idx = torch.argmax(value).item()
# print("idx-->", idx)
label = config.id2label[idx]
outputs.append((token , label))
# 输出outputs
print("outputs-->", outputs)
if __name__ == '__main__':
# dm01()
# dm02()
# dm03()
dm04()
# dm05()
# dm06()