使用Transformers、ChatGLM3项目、创建FastAPI应用等方式部署调用ChatGLM3-6B模型

Transformers部署调用

环境初始化

升级pip

python 复制代码

python -m pip install --upgrade pip

更换pypi源

python 复制代码

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

创建虚拟环境并使用

python 复制代码

conda create -n chatglm3-6b python=3.10
conda activate chatglm3-6b

安装PyTorch

注意：需根据GPU显卡驱动版本和CUDA版本选择对应匹配的PyTorch版本。

PyTorch版本列表

python 复制代码

pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu121

安装依赖

python 复制代码

pip install modelscope
pip install transformers

模型下载

使用魔塔的modelscope库中的snapshot_download函数下载模型

python 复制代码

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import os

# 指定：模型名称、下载路径、版本号
model_dir = snapshot_download('ZhipuAI/chatglm3-6b', cache_dir='/root/models', revision='master')

模型加载

python 复制代码

# 导入transformer库中的AutoTokenizer和AutoModelForCausalLM来加载分词器和对话模型
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 使用模型下载到的本地路径以加载
model_dir = './chatglm3-6b'
# 从本地模型加载分词器，trust_remote_code=True：允许从网络上下载模型权重和相关的代码
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
# 使用AutoModelForCausalLM类，从本地加载模型
model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True)

# 将模型移动到GPU上进行加速，否则使用CPU
if torch.cuda.is_available:
    device = torch.device("cuda")
else:
     device = torch.device("cpu")
model.to(device)
# 使用模型的评估模式来产生对话
model.eval()

# 第一轮对话
response, history = model.chat(tokenizer, "你好", history=[])
print(response + "\n")
# 第二轮对话
response, history = model.chat(tokenizer, "你是谁", history=history)
print(response + "\n")

模型GPU显存占用

默认情况下，模型以半精度（float16）加载，模型权重大概需要13GB不到的显存。

执行nvidia-smi，查看显卡显存占用情况

获取当前模型实际占用的GPU显存

python 复制代码

memory_bytes = model.get_memory_footprint()
# 转换为GB
memory_gb = memory_bytes / (1024 ** 3)  
print(f"{memory_gb :.2f}GB") # 11.63GB

注意：与实际进程占用有差异，差值为预留给PyTorch的显存

基于ChatGLM3项目部署调用

ChatGLM3项目

使用开源项目ChatGLM3进行ChatGLM3-6B部署
ChatGLM3项目提供了OpenAI / ZhipuAI格式的开源模型API部署代码，这里使用OpenAI格式的API部署。

下载ChatGLM3仓库

python 复制代码

git clone https://github.com/THUDM/ChatGLM3

cd ChatGLM3

使用pip安装依赖

注意：当环境已安装torch时，需修改requirements.txt文件，移除torch

python 复制代码

pip install -r requirements.txt

配置模型、嵌入模型路径

编辑ChatGLM3/openai_api_demo/api_server.py文件，需要设置AI大模型、嵌入模型的地址

python 复制代码

# 设置AI大模型的路径
MODEL_PATH = os.environ.get('MODEL_PATH', '/root/models/chatglm3-6b')

# 设置嵌入模型的路径
EMBEDDING_PATH = os.environ.get('EMBEDDING_PATH', '/root/models/bge-large-zh')

部署API

启动部署API Sever

python 复制代码

python  ChatGLM3/openai_api_demo/api_server.py

启动日志如下：

python 复制代码

(ChatGLM3) root@master:~/work/ChatGLM3/openai_api_demo# python api_server.py
Setting eos_token is not supported, use the default one.
Setting pad_token is not supported, use the default one.
Setting unk_token is not supported, use the default one.
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:13<00:00,  2.69s/it]
INFO:     Started server process [437763]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

Curl测试

python 复制代码

!curl -X POST "http://192.168.5.210:8001/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{\"model\": \"chatglm3-6b\", \"messages\": [{\"role\": \"system\", \"content\": \"You are ChatGLM3, a large language model trained by Zhipu.AI\"}, {\"role\": \"user\", \"content\": \"你好，给我讲一个故事，大概100字\"}], \"stream\": false, \"max_tokens\": 100, \"temperature\": 0.8, \"top_p\": 0.8}"

python代码测试

安装openai库

python 复制代码

pip install openai

编写python代码调用API

python 复制代码

from openai import OpenAI

base_url = "http://127.0.0.1:8000/v1/"
client = OpenAI(api_key="EMPTY", base_url=base_url)

def simple_chat(use_stream=True):
    messages = [{"role":"system","content":"你是一个乐于助人的助手"},{"role":"user","content":"你是谁？"}]
    
    response = client.chat.completions.create(
        model="chatglm3-6b",
        messages=messages,
        stream=use_stream,
        max_tokens=256,
        temperature=0.8,
        presence_penalty=1.1,
        top_p=0.8)
    
    if response:
        if use_stream:
            for chunk in response:
                print(chunk.choices[0].delta.content)
        else:
            content = response.choices[0].message.content
            print(content)
    else:
        print("Error:", response.status_code)

if __name__ == '__main__':
    simple_chat(use_stream=False)

创建FastAPI应用部署调用

基于ChatGLM3项目提供的API部署代码可以很方便的进行ChatGLM3-6B模型的API部署。
但是其提供的API部署代码可能不适用于其他AI模型，因此可以创建一个更通用的FastAPI应用，用于进行各种模型的API部署

安装依赖库

python 复制代码

pip install fastapi uvicorn requests

创建FastAPI应用

创建api.py文件，编写如下代码

python 复制代码

from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModelForCausalLM
import uvicorn
import json
import datetime
import torch

def torch_gc():
	 """
    清理GPU显存
    """
    if torch.cuda.is_available():  # 检查CUDA是否可用
        with torch.cuda.device("cuda:0"):  # 指定CUDA设备
            torch.cuda.empty_cache()  # 清空CUDA缓存
            torch.cuda.ipc_collect()  # 收集CUDA内存碎片

# 创建FastAPI应用
app = FastAPI()

@app.post("/")
async def create_item(request: Request):
	"""
    处理POST请求
    """
    global model, tokenizer   # 声明全局变量以便在函数内部使用模型和分词器
    json_post_raw = await request.json()  # 获取请求JSON数据
    json_post = json.dumps(json_post_raw)  # 将JSON数据转换为字符串
    json_post_list = json.loads(json_post)  # 将字符串转换为Python对象
    prompt = json_post_list.get('prompt')  # 获取参数prompt
    history = json_post_list.get('history')  # 获取参数history
    max_length = json_post_list.get('max_length')  # 获取参数max_length
    top_p = json_post_list.get('top_p')  # 获取参数top_p
    temperature = json_post_list.get('temperature')  # 获取参数temperature
    
    # 调用模型进行对话生成
    response, history = model.chat(
        tokenizer,
        prompt,
        history=history,
        max_length=max_length if max_length else 2048,  # 如果未提供最大长度，默认使用2048
        top_p=top_p if top_p else 0.7,  # 如果未提供top_p参数，默认使用0.7
        temperature=temperature if temperature else 0.6  # 如果未提供温度参数，默认使用0.6
    )

    now = datetime.datetime.now()  # 获取当前时间
    time = now.strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间为字符串

    # 构建响应JSON
    answer = {
        "response": response,
        "history": history,
        "status": 200,
        "time": time
    }

    # 构建日志信息
    log = f"[{time}] prompt: {prompt}, response: {repr(response)}"
    print(log)  # 打印日志
    torch_gc()  # 执行GPU内存清理
    return answer  # 返回响应


if __name__ == '__main__':
	model_path = "./chatglm3-6b"
    # 加载预训练的分词器和模型
    tokenizer = AutoTokenizer.from_pretrained(model_path , trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_path , trust_remote_code=True).to(torch.bfloat16).cuda()
    model.eval()  # 设置模型为评估模式
    # 启动FastAPI应用
    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

启动api服务

python 复制代码

(ChatGLM3) root@master:~/work/# python api.py 
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:02<00:00,  3.01it/s]
INFO:     Started server process [439374]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
[2024-05-11 08:22:10] prompt: 你好, response: '你好👋！我是人工智能助手 ChatGLM3-6B，很高兴见到你，欢迎问我任何问题。'
INFO:     192.168.5.210:55094 - "POST / HTTP/1.1" 200 OK

python代码测试

requests库代码调用如下

python 复制代码

import requests
import json

def get_completion(prompt):
    headers = {'Content-Type': 'application/json'}
    data = {"prompt": prompt, "history": []}
    response = requests.post(url='http://127.0.0.1:8000', headers=headers, data=json.dumps(data))
    return response.json()['response']

if __name__ == '__main__':
    print(get_completion('你好'))