【transformer模型】将huggingface pipeline中的模型保存到本地并加载的方法_huggingface pipeline加载本地模型

第一次下载模型

from transformers import pipeline
classifier = pipeline(task="sentiment-analysis")
preds = classifier("cat haters!")
preds = [{"score": round(pred["score"], 4), "label": pred["label"]} for pred in preds]
preds
1
2
3
4
5

可以看到输出：
[{‘score’: 0.9984, ‘label’: ‘NEGATIVE’}]

每次下载模型需要一点时间，因此想保存到本地使用。

保存模型到本地

save_path = "/Users/***/Desktop/csfiles/transformer/sentiment_analysis_model"
classifier.save_pretrained(save_path)
print(f"模型已保存到本地路径：{save_path}")
1
2
3

注意这里save_path写全，可以用pwd命令看一下自己所在的当前目录。

保存后可以看到生成了一个新的文件夹：

加载已保存的模型

loaded_model = pipeline(task='sentiment-analysis', model = save_path)
1

注意# task=‘sentiment-analysis’ 这里不要胡乱命名，就用下载的模型本来的名字。

利用加载的模型进行预测

text = 'cat good, human bad'
preds = loaded_model(text)
preds = [{"score": round(pred["score"], 4), "label": pred["label"]} for pred in preds]
print(preds)
1
2
3
4

可以看到输出：
[{‘score’: 0.9965, ‘label’: ‘POSITIVE’}]
猫好，人坏，POSITIVE～