commit
549cdeb054
@ -165,6 +165,13 @@ cd ChatGLM2-6B
|
||||
git clone https://huggingface.co/THUDM/chatglm2-6b
|
||||
```
|
||||
|
||||
如果你从 Hugging Face Hub 上下载 checkpoint 的速度较慢,可以只下载模型实现
|
||||
```Shell
|
||||
GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/THUDM/chatglm2-6b
|
||||
```
|
||||
然后从[这里](https://cloud.tsinghua.edu.cn/d/674208019e314311ab5c/)手动下载模型参数文件,并将下载的文件替换到本地的 `chatglm2-6b` 目录下。
|
||||
|
||||
|
||||
将模型下载到本地之后,将以上代码中的 `THUDM/chatglm2-6b` 替换为你本地的 `chatglm2-6b` 文件夹的路径,即可从本地加载模型。
|
||||
|
||||
模型的实现仍然处在变动中。如果希望固定使用的模型实现以保证兼容性,可以在 `from_pretrained` 的调用中增加 `revision="v1.0"` 参数。`v1.0` 是当前最新的版本号,完整的版本列表参见 [Change Log](https://huggingface.co/THUDM/chatglm2-6b#change-log)。
|
||||
|
4
api.py
4
api.py
@ -52,5 +52,9 @@ async def create_item(request: Request):
|
||||
if __name__ == '__main__':
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
# 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量
|
||||
# model_path = "THUDM/chatglm2-6b"
|
||||
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
# model = load_model_on_gpus(model_path, num_gpus=2)
|
||||
model.eval()
|
||||
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)
|
||||
|
20
cli_demo.py
20
cli_demo.py
@ -3,9 +3,14 @@ import platform
|
||||
import signal
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
import readline
|
||||
from utils import load_model_on_gpus
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
# 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量
|
||||
# model_path = "THUDM/chatglm2-6b"
|
||||
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
# model = load_model_on_gpus(model_path, num_gpus=2)
|
||||
model = model.eval()
|
||||
|
||||
os_name = platform.system()
|
||||
@ -17,7 +22,7 @@ def build_prompt(history):
|
||||
prompt = "欢迎使用 ChatGLM2-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序"
|
||||
for query, response in history:
|
||||
prompt += f"\n\n用户:{query}"
|
||||
prompt += f"\n\nChatGLM-6B:{response}"
|
||||
prompt += f"\n\nChatGLM2-6B:{response}"
|
||||
return prompt
|
||||
|
||||
|
||||
@ -39,7 +44,8 @@ def main():
|
||||
os.system(clear_command)
|
||||
print("欢迎使用 ChatGLM2-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序")
|
||||
continue
|
||||
count = 0
|
||||
print("\nChatGLM:", end="")
|
||||
current_length = 0
|
||||
for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history,
|
||||
past_key_values=past_key_values,
|
||||
return_past_key_values=True):
|
||||
@ -47,13 +53,9 @@ def main():
|
||||
stop_stream = False
|
||||
break
|
||||
else:
|
||||
count += 1
|
||||
if count % 8 == 0:
|
||||
os.system(clear_command)
|
||||
print(build_prompt(history), flush=True)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
os.system(clear_command)
|
||||
print(build_prompt(history), flush=True)
|
||||
print(response[current_length:], end="", flush=True)
|
||||
current_length = len(response)
|
||||
print("")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
@ -158,6 +158,10 @@ async def predict(query: str, history: List[List[str]], model_id: str):
|
||||
if __name__ == "__main__":
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
# 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量
|
||||
# model_path = "THUDM/chatglm2-6b"
|
||||
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
# model = load_model_on_gpus(model_path, num_gpus=2)
|
||||
model.eval()
|
||||
|
||||
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)
|
||||
|
59
utils.py
Normal file
59
utils.py
Normal file
@ -0,0 +1,59 @@
|
||||
import os
|
||||
from typing import Dict, Tuple, Union, Optional
|
||||
|
||||
from torch.nn import Module
|
||||
from transformers import AutoModel
|
||||
|
||||
|
||||
def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
|
||||
# transformer.word_embeddings 占用1层
|
||||
# transformer.final_layernorm 和 lm_head 占用1层
|
||||
# transformer.layers 占用 28 层
|
||||
# 总共30层分配到num_gpus张卡上
|
||||
num_trans_layers = 28
|
||||
per_gpu_layers = 30 / num_gpus
|
||||
|
||||
# bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError
|
||||
# windows下 model.device 会被设置成 transformer.word_embeddings.device
|
||||
# linux下 model.device 会被设置成 lm_head.device
|
||||
# 在调用chat或者stream_chat时,input_ids会被放到model.device上
|
||||
# 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
|
||||
# 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
|
||||
# 本文件来源于https://github.com/THUDM/ChatGLM-6B/blob/main/utils.py
|
||||
# 仅此处做少许修改以支持ChatGLM2
|
||||
device_map = {
|
||||
'transformer.embedding.word_embeddings': 0,
|
||||
'transformer.encoder.final_layernorm': 0,
|
||||
'transformer.output_layer': 0,
|
||||
'transformer.rotary_pos_emb': 0,
|
||||
'lm_head': 0
|
||||
}
|
||||
|
||||
used = 2
|
||||
gpu_target = 0
|
||||
for i in range(num_trans_layers):
|
||||
if used >= per_gpu_layers:
|
||||
gpu_target += 1
|
||||
used = 0
|
||||
assert gpu_target < num_gpus
|
||||
device_map[f'transformer.encoder.layers.{i}'] = gpu_target
|
||||
used += 1
|
||||
|
||||
return device_map
|
||||
|
||||
|
||||
def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
|
||||
device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module:
|
||||
if num_gpus < 2 and device_map is None:
|
||||
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda()
|
||||
else:
|
||||
from accelerate import dispatch_model
|
||||
|
||||
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half()
|
||||
|
||||
if device_map is None:
|
||||
device_map = auto_configure_device_map(num_gpus)
|
||||
|
||||
model = dispatch_model(model, device_map=device_map)
|
||||
|
||||
return model
|
@ -1,9 +1,14 @@
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
import gradio as gr
|
||||
import mdtex2html
|
||||
from utils import load_model_on_gpus
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
# 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量
|
||||
# model_path = "THUDM/chatglm2-6b"
|
||||
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
# model = load_model_on_gpus(model_path, num_gpus=2)
|
||||
model = model.eval()
|
||||
|
||||
"""Override Chatbot.postprocess"""
|
||||
|
@ -14,6 +14,10 @@ st.set_page_config(
|
||||
def get_model():
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
|
||||
# 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量
|
||||
# model_path = "THUDM/chatglm2-6b"
|
||||
# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
# model = load_model_on_gpus(model_path, num_gpus=2)
|
||||
model = model.eval()
|
||||
return tokenizer, model
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user