huggingface pipeline使用清华智谱的模型THUDM/chatglm3-6b

hehui0921

已于 2024-02-29 15:50:36 修改

阅读量283

点赞数 3

分类专栏： huggingface 文章标签： java python 前端

于 2024-02-12 12:54:04 首次发布

本文链接：https://blog.csdn.net/oHeHui1/article/details/136100664

版权

huggingface 专栏收录该内容

33 篇文章 0 订阅

订阅专栏

以下代码成功运在CPU机器上，需要32G内存：
**Google colab 中运行报内存不足 * *
第一次运行会自动下载模型文件，需要比较长的时间。

from transformers import AutoTokenizer, AutoModel
from transformers import AutoModelForSeq2SeqLM
from huggingface_hub.hf_api import HfFolder

HfFolder.save_token('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')
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True, device='cpu')
model = model.eval()
response = model.chat(tokenizer, "你好", history=[])
print(response)

以下使用GPU运行chatglm3-6b：

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True).half().cuda()
model = model.eval()
for i in range(10):
    response, history = model.chat(tokenizer, "来一首唐诗", history=[])
    print(response)