1--AutoModel的使用
官方文档
AutoModel 用于加载模型;
1-1--简单Demo
测试代码:
from transformers import AutoTokenizer, AutoModelif __name__ == "__main__":checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"tokenlizer = AutoTokenizer.from_pretrained(checkpoint) raw_input = ["I love kobe bryant.", "Me too."]inputs = tokenlizer(raw_input, padding = "longest", truncation = True, max_length = 512, return_tensors = "pt")# 加载指定的模型model = AutoModel.from_pretrained(checkpoint)print("model: \n", model)outputs = model(**inputs)print("last_hidden_state: \n", outputs.last_hidden_state.shape) # 打印最后一个隐层的输出维度# [2 7 768] batch_size为2,7个token,每个token的维度为768
输出结果:
last_hidden_state: torch.Size([2, 7, 768])# 最后一个隐层的输出
# batchsize为2,表示两个句子
# 7表示token数,每一个句子有7个token
# 768表示特征大小,每一个token的维度为768
测试代码:
from transformers import AutoTokenizer, AutoModelForSequenceClassificationif __name__ == "__main__":checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"tokenlizer = AutoTokenizer.from_pretrained(checkpoint) raw_input = ["I love kobe bryant.", "Me too."]inputs = tokenlizer(raw_input, padding = "longest", truncation = True, max_length = 512, return_tensors = "pt")model2 = AutoModelForSequenceClassification.from_pretrained(checkpoint) # 二分类任务print(model2)outputs2 = model2(**inputs)print(outputs2.logits.shape)
运行结果:
torch.Size([2, 2])
# 两个句子,每个句子二分类的概率