ChatGLM2-6B 运行代码解读（二）

目录

一、cli_demo.py 解读

3.1 transformers

3.2 AutoTokenizer.from_pretrained 函数

3.3 AutoModel.from_pretrained 函数

3.4 model.eval() 函数

3.5 model.stream_chat 函数

---

本文主要对 cli_demo.py 进行介绍。

一、cli_demo.py 解读

终端代码如下所示。

import os
import platform
import signal
from transformers import AutoTokenizer, AutoModel
import readline

tokenizer = AutoTokenizer.from_pretrained("D:\code\ChatGLM-6B-main\THUDM\chatglm2-6b", trust_remote_code=True)
# model = AutoModel.from_pretrained("D:\code\ChatGLM-6B-main\THUDM\chatglm2-6b", trust_remote_code=True).cuda()
model = AutoModel.from_pretrained("D:\code\ChatGLM-6B-main\THUDM\chatglm2-6b", trust_remote_code=True).float()
# 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量
# from utils import load_model_on_gpus
# model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)
# model.eval()方法用于将模型设置为评估模式，用于推理或测试
model = model.eval()

# windows
os_name = platform.system()
clear_command = '