【大模型】DeepSeek-R1-Distill-Qwen部署及API调用

DeepSeek-R1-Distill-Qwen 是由中国人工智能公司深度求索（DeepSeek）开发的轻量化大语言模型，基于阿里巴巴的Qwen系列模型通过知识蒸馏技术优化而来。当前模型开源后，我们可以将其部署，使用API方式进行本地调用

1.部署环境

本文中的部署基础环境如下所示：

PyTorch  2.5.1
Python  3.12(ubuntu22.04)
Cuda  12.4
GPU RTX 3090(24GB) * 1
CPU 14 vCPU Intel(R) Xeon(R) Platinum 8362 CPU @ 2.80GHz

2.依赖安装

# 升级pip
python -m pip install --upgrade pip
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

pip install requests==2.32.3
pip install fastapi==0.115.8
pip install uvicorn==0.34.0
pip install transformers==4.48.2
pip install huggingface-hub==0.28.1
pip install accelerate==1.3.0
pip install modelscope==1.22.3

3.模型下载

使用 modelscope 中的 snapshot_download 函数下载模型，第一个参数为模型名称，参数 cache_dir 为模型的下载路径，这里放在autodl的数据文件夹中比较合适。

4.API部署代码

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

DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE

def clear_gpu_cache():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

def split_response(text):
    match = re.search(r'(.*?)(.*)', text, re.DOTALL)
    if match:
        return match.group(1).strip(), match.group(2).strip()
    return "", text.strip()

app = FastAPI()

@app.post("/")
async def process_request(request: Request):
    data = await request.json()
    prompt = data.get('prompt', '')
    
    messages = [{"role": "user", "content": prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
    
    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=8192)
    response_text = tokenizer.batch_decode(generated_ids[:, model_inputs.input_ids.shape[-1]:], skip_special_tokens=True)[0]
    
    think_content, answer_content = split_response(response_text)
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    log_entry = f"[{timestamp}], prompt:\"{prompt}\", response:\"{repr(response_text)}\", think:\"{think_content}\", answer:\"{answer_content}\""
    print(log_entry)
    
    clear_gpu_cache()
    
    return {
        "response": response_text,
        "think": think_content,
        "answer": answer_content,
        "status": 200,
        "time": timestamp
    }

if __name__ == '__main__':
    model_path = '/root/autodl-tmp/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B'
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=CUDA_DEVICE, torch_dtype=torch.bfloat16)
    
    uvicorn.run(app, host='0.0.0.0', port=6006, workers=1)

这段代码的主要作用是构建一个基于 FastAPI 的 HTTP API 服务，用于调用预训练的语言生成模型进行对话生成：

①依赖导入与设备配置

• 依赖库：
  引入了 FastAPI、Request（用于处理请求）、Transformers 中的 AutoTokenizer 与 AutoModelForCausalLM（用于加载预训练模型和分词器）、uvicorn（用于启动服务器）、json、datetime、torch 和 re（正则表达式处理）。
• 设备配置：
  定义了 DEVICE 为 “cuda”，并指定 DEVICE_ID 为 “0”。根据是否有设备ID，组合出 CUDA_DEVICE 信息，用于在加载模型时指定使用的 GPU 设备。

②GPU内存管理

• 函数 clear_gpu_cache()：
  当使用 GPU 进行推理时，可能会因内存碎片或缓存问题影响性能。该函数首先检查是否有可用的 CUDA 设备，然后调用 torch.cuda.empty_cache() 清空未使用的缓存，并用 torch.cuda.ipc_collect() 收集内存碎片，从而确保 GPU 内存得到有效释放。

③文本分割

• 函数 split_response(text)：
  该函数使用正则表达式将模型生成的文本分为两部分：
  • ... 标签内的内容被认为是“思考过程”；
  • 标签之后的内容作为最终的“回答”。
    如果文本中没有符合该格式的部分，则返回空字符串作为思考内容，而将整个文本视作回答。

④FastAPI 应用与请求处理

• 创建 FastAPI 实例：
  通过 app = FastAPI() 创建 API 应用实例。
• POST 请求端点 /：
  定义了一个 POST 接口 process_request，用于处理客户端请求：
  1. 请求数据获取：
     • 使用 await request.json() 获取请求中的 JSON 数据，并从中提取 prompt（提示内容），如果不存在则默认为空字符串。
  2. 构造对话消息：
     • 将提取到的 prompt 以对话消息格式（角色为 “user”）构造为一个列表 messages。
  3. 模型输入准备：
     • 使用 tokenizer.apply_chat_template 方法将消息转换为模型所需的输入文本。
     • 利用分词器将输入文本转换为张量（tensor），并将其移动到模型所在的设备上（即 GPU）。
  4. 文本生成：
     • 调用模型的 generate 方法生成文本，设置 max_new_tokens 为 8192，允许生成较长的文本。
     • 使用 tokenizer.batch_decode 方法将生成的 token 转换回可读文本，并跳过特殊标记。这里注意只解码模型生成部分（通过对生成结果进行切片）。
  5. 文本分割与日志记录：
     • 调用 split_response 函数，将生成文本分割为思考内容和回答内容。
     • 获取当前时间，并以格式化的字符串形式记录。
     • 构建日志条目，包含时间戳、请求提示、完整响应、思考内容以及回答内容，并打印到控制台。
  6. GPU 内存清理：
     • 调用 clear_gpu_cache() 清理 GPU 内存，确保内存资源得到释放。
  7. 返回响应：
     • 构建 JSON 格式的响应，包含生成的完整响应文本、分割后的“think”和“answer”内容、状态码（200）和时间戳。

⑤模型加载与服务启动

• 加载模型和分词器：
  在 if __name__ == '__main__': 代码块中：

  • 从指定路径 '/root/autodl-tmp/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B' 加载预训练的分词器和模型。
  • 使用 device_map=CUDA_DEVICE 指定使用 GPU，并设置 torch_dtype=torch.bfloat16 以优化内存使用。

• 启动服务器：

  • 使用 uvicorn.run 启动 FastAPI 应用，绑定主机地址 0.0.0.0 和端口 6006。这里设置了单个 worker 进程，适用于资源有限或简单部署场景。

5.调用测试

首先需要启动API：

python api.py

然后使用python的request库进行调用：

import requests
import json

def get_completion(prompt, url='http://127.0.0.1:6006'):
    headers = {'Content-Type': 'application/json'}
    data = {"prompt": prompt}
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data), timeout=10)
        response.raise_for_status()  # 检查 HTTP 状态码
        return response.json()  # 返回完整 JSON 结果
    except requests.exceptions.RequestException as e:
        return {"error": str(e)}

def pretty_print_result(result):
    # 使用 json.dumps 格式化输出
    formatted = json.dumps(result, indent=4, ensure_ascii=False)
    print("Structured Output:")
    print(formatted)

if __name__ == '__main__':
    prompt = '说一下9.11和9.8那个更大，你愿意要9.11吨黄金还是9.8吨'
    result = get_completion(prompt)
    pretty_print_result(result)