【0基础】使用Mistral-7B模型搭建简易的聊天API（进阶篇）

OrikasaYuki

已于 2024-08-08 21:29:31 修改

阅读量603

点赞数 11

分类专栏： 0基础AI对话机器人教程文章标签： python transformer 语言模型

于 2024-08-08 21:26:20 首次发布

本文链接：https://blog.csdn.net/qq_36548916/article/details/141034633

版权

0基础AI对话机器人教程专栏收录该内容

3 篇文章 0 订阅

订阅专栏

准备环节

前言的前言

接着基础篇的前言继续闲谈：UCloud新出的Compshare 算力平台，说实话，抛开广告成分而言，我仍然认为是最佳的选择，毕竟说实话，对像我这样的学生党来说，便宜大部分时候就是好用，更何况这个便宜并没有使质量产生损失。如果你和我一样预算有限，我十分推荐你尝试一下，因为有独立外网IP、预装好环境的镜像等等，对新手朋友来说体验只会比阿里云等更好。首推的4090等GPU平台正好适合用于这种小模型推理，价格每月也只要一千出头（也可以日付，才50r）（这个价格的最低配甚至是16核64G，200GB系统盘，共享100M带宽）。学生还有额外的八折优惠。（此外，像我这样发布相关文章还可以直接抵扣。）

前言

本文是在基础篇之上的改进，如果你还没有看过基础篇，请移步：

【0基础】使用Mistral-7B模型搭建简易的聊天API（基础篇）-CSDN博客

MistralAI团队去年发布mistral 7b模型至今，mistral7b就一直是同参数量内表现最亮眼的模型，逻辑能力、代码能力都非常强大，官方还发布了为聊天对话而微调的版本，也没有为个人用户自行下载部署设置任何额外的门槛（甚至当时在推特上直接发出磁力链接的时候我都被惊到了）。但是总感觉该模型在国内关注度不是很高（也可能是我相关网页刷少了），如果查找相关部署教程，找到的大多是使用mistral官方的软件包进行部署的说明。因此今天想发点文章给在部署这类模型的领域里0基础的朋友们分享一下，即如何使用Mistral-7B Instruct v0.3模型搭建一个简易的聊天API。之后还会写一些本文用到的内容的单独的短文，都会是基础向。

上一篇基础篇只支持一问一答的形式，本篇进阶篇会带来实时更新生成情况、多轮对话等更复杂的功能实操。

模型下载

嘿，当你想看这部分内容的时候我就知道你把前言跳过了。

这篇文章是进阶篇，最基本的模型下载部署都在基础篇里，快去看（

代码实现

书接上回

首先我们回顾一下基础篇的最终成品：

import torch
from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM
from flask import request
import flask
import json

model_addr = "./model"  # 模型路径（到文件夹为止）
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_addr,
                                          trust_remote_code=True,
                                          use_safetensors=True,
                                          quantization_config=nf4_config)
model = AutoModelForCausalLM.from_pretrained(model_addr,
                                             trust_remote_code=True,
                                             use_safetensors=True,
                                             quantization_config=nf4_config)

model = model.eval()
system_prompt = "你是一个AI助理，你的功能是与用户对话交流。"
B_INST, E_INST = "[INST]", "[/INST]"
max_tokens = 200


def generate_ans(que):
    prompt = f"{system_prompt}{B_INST}{que.strip()}\n{E_INST}"
    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
    attention_mask = torch.ones_like(inputs).to("cuda")

    res = model.generate(inputs, max_new_tokens=max_tokens,
                         attention_mask=attention_mask,
                         pad_token_id=tokenizer.eos_token_id)
    ans = tokenizer.batch_decode(res, skip_special_tokens=True)[0]

    return ans


server = flask.Flask(__name__)


@server.route('/api/aiAssist', methods=['get', 'post'])
def ai_assist():
    # 获取通过url请求传参的数据
    reqQue = request.values.get('q')

    if reqQue:
        r = generate_ans(reqQue)
        resu = {'code': 200, 'message': r}

    else:
        resu = {'code': 10001, 'message': '参数不能为空'}

    return json.dumps(resu, ensure_ascii=False)


server.run(debug=False, port=8888, host='0.0.0.0')

这段代码可以实现简单的一问一答。但存在一些问题：

首先，响应时间会很长，而且只能得到最终结果，可能会影响实际使用时的体验。因为它的相应用时完全取决于回复花了多久生成完毕，个人电脑配置再高，也得至少等个十几秒。

其次，不支持连续对话，无法接着上文内容继续进行提问。

想解决这些问题，我们需要进一步修改代码。

流式输出

想要实现一边生成一边输出，我们需要借助streamer和多线程实现。

首先，我们引入streamer和threading。

from transformers import TextIteratorStreamer
import threading

接着，我们定义一个streamer变量。

streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)

参数skip_prompt可以在输出时跳过system_prompt和用户的输入，直接从回复开始输出。

然后，我们使用多线程生成回复：

thread = threading.Thread(target=generator, args=(streamer, req, ))
thread.daemon = True
thread.start()

这段代码开启新的线程让generator工作。那么，我们接着实现generator并完善流式输出功能。

generator：

def generator(self, streamer, que):
    prompt = f"{system_prompt}{B_INST}{que.strip()}\n{E_INST}"
    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
    attention_mask = torch.ones_like(inputs).to("cuda")

    _ = self.model.generate(inputs, streamer=streamer, max_new_tokens=max_tokens,
                            attention_mask=attention_mask,
                            pad_token_id=self.tokenizer.eos_token_id)

    torch.cuda.empty_cache()

流式输出：

generated_text = ""

for new_text in streamer:
    generated_text += new_text

thread.join()

使用这种方式，generated_text会随着生成不断更新。

接下来，我们只需要将以上内容汇总进程序中即可。

一个简单的思路是：用户第一次带着问题访问时，返回一个key，之后用户每一次使用key访问时，返回当前生成的回复即可。下面进行实现和打包。

import gc
import random
import threading
import time
import flask
import json
import torch
from flask import request
from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer


class AIReply:
    def __init__(self):
        self.state = "Ready"
        self.content = ""
        self.last_use = time.time()


class AIBot:
    def __init__(self):
        self.model_addr = "./model"

        self.max_tokens = 300  # 最大单次输出token数
        self.system_prompt = "你是一个AI助理，你的功能是与用户对话交流。"

        self.reply: dict[str, AIReply] = {}

        self.model = None
        self.tokenizer = None

        self.r_cleaner_working = False

    def generate_key(self):
        r = random
        x = ''
        while True:
            for i in range(16):
                x += chr(r.randint(65, 90))
            if not self.reply.__contains__(x):
                return x

    def load_model(self):
        if self.model_addr != "":
            nf4_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16
            )

            self.tokenizer = AutoTokenizer.from_pretrained(self.model_addr,
                                                           trust_remote_code=True,
                                                           use_safetensors=True,
                                                           quantization_config=nf4_config)
            self.model = AutoModelForCausalLM.from_pretrained(self.model_addr,
                                                              trust_remote_code=True,
                                                              use_safetensors=True,
                                                              quantization_config=nf4_config)

            self.model = self.model.eval()
        gc.collect()
        torch.cuda.empty_cache()

    def generator(self, streamer, query):
        B_INST, E_INST = "[INST]", "[/INST]"

        prompt = self.system_prompt
        max_tokens = self.max_tokens

        prompt = f"{prompt}{B_INST}{query.strip()}\n{E_INST}"

        inputs = self.tokenizer.encode(prompt, return_tensors="pt").to("cuda")

        attention_mask = torch.ones_like(inputs).to("cuda")

        _ = self.model.generate(inputs, streamer=streamer, max_new_tokens=max_tokens,
                                attention_mask=attention_mask,
                                pad_token_id=self.tokenizer.eos_token_id)

        gc.collect()
        torch.cuda.empty_cache()

    def gen_reply(self, req, key):
        print("Request: " + req)

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True)
        thread = threading.Thread(target=self.generator, args=(streamer, req, ))
        thread.daemon = True
        thread.start()

        st = time.time()

        generated_text = ""
        count = 0

        for new_text in streamer:
            generated_text += new_text

            count += 1
            if count == 1:
                st = time.time()
                self.reply[key].state = "Generating"
            if count % 2 == 0:
                self.reply[key].content = generated_text
                self.reply[key].last_use = time.time()

        thread.join()
        tokens = count
        ed = time.time()

        if generated_text[-4:] == "</s>":
            generated_text = generated_text[0:-4]

        self.reply[key].state = "Finish"
        self.reply[key].content = generated_text
        self.reply[key].last_use = time.time()

        info_str = "Time:" + str(round(ed - st, 2)) + "\t" + "Tokens:" + str(
            tokens) + "\t" + "Speed:" + str(round(tokens / (ed - st), 4))

        print(info_str)
        torch.cuda.empty_cache()

    def view_reply(self, key):
        if self.reply.__contains__(key):
            r = self.reply[key]
            self.reply[key].last_use = time.time()
        else:
            r = AIReply()
        return r

    def new_reply(self, req):
        r_key = self.generate_key()

        self.reply[r_key] = AIReply()

        thread = threading.Thread(target=self.gen_reply, args=(req, r_key, ))
        thread.daemon = True
        thread.start()

        return r_key

    def main_func(self):
        self.load_model()


main_F = AIBot()
server = flask.Flask(__name__)


@server.route('/api/aiAssist', methods=['get', 'post'])
def ai_assist():
    reqQue = request.values.get('q')
    reqKey = request.values.get('k')
    
    resu = {'code': -2, 'message': '访问出错'}
    
    if reqQue:
        r_key = main_F.new_reply(reqQue)
        resu = {'code': 200, 'message': r_key}
    elif reqKey:
        if main_F.reply.__contains__(reqKey):
            if main_F.view_reply(reqKey).state == "Ready":
                resu = {'code': 10002, 'message': main_F.view_reply(reqKey).content}
            if main_F.view_reply(reqKey).state == "Generating":
                resu = {'code': 10003, 'message': main_F.view_reply(reqKey).content}
            if main_F.view_reply(reqKey).state == "Finish":
                resu = {'code': 200, 'message': main_F.view_reply(reqKey).content}
        else:
            resu = {'code': -1, 'message': '访问密钥不存在'}
    else:
        resu = {'code': 10001, 'message': '参数不能为空'}

    return json.dumps(resu, ensure_ascii=False)


if __name__ == "__main__":
    main_F.main_func()
    server.run(debug=False, port=8888, host='0.0.0.0')

上述代码可能相比基础版变化比较大，但实际上拆分来看，新增的内容并不算多。

首先创建了新的类AIReply和字典reply用于存储回复。有新的请求会创建一个AIReply、一串大写英文字母组成的Key，并将key返回，同时开始生成回复。之后凭借key访问即可查看回复的状态（准备开始、生成中和生成完毕）并获取到当前已经生成的内容。初步实现了获取当前回复状态的功能，还可以避免单次访问响应时间过长。

需要注意的是，这段代码虽然可以实现相应功能，但如果要实际大规模使用，还是需要自行改进，加入限制访问频率、从内存中清除过早的内容等等（根据AIReply的last_use属性，即上次访问的时间），建议借助数据库实现，在这里不多赘述了。

接下来，如果你还有连续对话的需求，可以继续在此基础上修改。

连续对话

连续对话只需要在上面的基础上稍加改动，为AIReply增加对话记录的记忆即可。当用户带着q访问，给出key，带着key访问给出最新回复，带着key和q的时候就结合上文对新的请求进行回复。

下面给出实现：

import gc
import random
import threading
import time
import flask
import json
import torch
from flask import request
from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer


class AIReply:
    def __init__(self):
        self.state = "Ready"
        self.content = ""
        self.last_use = time.time()
        self.history = []


class AIBot:
    def __init__(self):
        self.model_addr = "./model"

        self.max_tokens = 300  # 最大单次输出token数
        self.system_prompt = "你是一个AI助理，你的功能是与用户对话交流。"

        self.reply: dict[str, AIReply] = {}

        self.model = None
        self.tokenizer = None

        self.r_cleaner_working = False

    def generate_key(self):
        r = random
        x = ''
        while True:
            for i in range(16):
                x += chr(r.randint(65, 90))
            if not self.reply.__contains__(x):
                return x

    def load_model(self):
        if self.model_addr != "":
            nf4_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16
            )

            self.tokenizer = AutoTokenizer.from_pretrained(self.model_addr,
                                                           trust_remote_code=True,
                                                           use_safetensors=True,
                                                           quantization_config=nf4_config)
            self.model = AutoModelForCausalLM.from_pretrained(self.model_addr,
                                                              trust_remote_code=True,
                                                              use_safetensors=True,
                                                              quantization_config=nf4_config)

            self.model = self.model.eval()
        gc.collect()
        torch.cuda.empty_cache()

    def generator(self, streamer, query, hist):
        B_INST, E_INST = "[INST]", "[/INST]"

        prompt = self.system_prompt

        for q, r in hist:
            prompt += "<s>"
            prompt += B_INST
            prompt += q
            prompt += E_INST
            prompt += r
            prompt += "</s>"

        max_tokens = self.max_tokens

        prompt = f"{prompt}{B_INST}{query.strip()}\n{E_INST}"

        inputs = self.tokenizer.encode(prompt, return_tensors="pt").to("cuda")

        attention_mask = torch.ones_like(inputs).to("cuda")

        _ = self.model.generate(inputs, streamer=streamer, max_new_tokens=max_tokens,
                                attention_mask=attention_mask,
                                pad_token_id=self.tokenizer.eos_token_id)

        gc.collect()
        torch.cuda.empty_cache()

    def gen_reply(self, req, key):
        print("Request: " + req)

        hist = self.reply[key].history

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True)
        thread = threading.Thread(target=self.generator, args=(streamer, req, hist, ))
        thread.daemon = True
        thread.start()

        st = time.time()

        generated_text = ""
        count = 0

        for new_text in streamer:
            generated_text += new_text

            count += 1
            if count == 1:
                st = time.time()
                self.reply[key].state = "Generating"
            if count % 2 == 0:
                self.reply[key].content = generated_text
                self.reply[key].last_use = time.time()

        thread.join()
        tokens = count
        ed = time.time()

        if generated_text[-4:] == "</s>":
            generated_text = generated_text[0:-4]

        self.reply[key].state = "Finish"
        self.reply[key].content = generated_text

        self.reply[key].history.append((req, generated_text))

        self.reply[key].last_use = time.time()

        info_str = "Time:" + str(round(ed - st, 2)) + "\t" + "Tokens:" + str(
            tokens) + "\t" + "Speed:" + str(round(tokens / (ed - st), 4))

        print(info_str)
        torch.cuda.empty_cache()

    def view_reply(self, key):
        if self.reply.__contains__(key):
            r = self.reply[key]
            self.reply[key].last_use = time.time()
        else:
            r = AIReply()
        return r

    def new_reply(self, req):
        r_key = self.generate_key()

        self.reply[r_key] = AIReply()

        thread = threading.Thread(target=self.gen_reply, args=(req, r_key,))
        thread.daemon = True
        thread.start()

        return r_key

    def add_reply(self, req, key):
        self.reply[key].state = "Ready"
        self.reply[key].content = ""
        thread = threading.Thread(target=self.gen_reply, args=(req, key,))
        thread.daemon = True
        thread.start()
        return key

    def main_func(self):
        self.load_model()


main_F = AIBot()
server = flask.Flask(__name__)


@server.route('/api/aiAssist', methods=['get', 'post'])
def ai_assist():
    reqQue = request.values.get('q')
    reqKey = request.values.get('k')

    resu = {'code': -2, 'message': '访问出错'}

    if reqQue and reqKey:
        if main_F.reply.__contains__(reqKey):
            if main_F.view_reply(reqKey).state != "Finish":
                resu = {'code': 10004, 'message': "上次回复还没有结束"}
            else:
                resu = {'code': 200, 'message': main_F.add_reply(reqQue, reqKey)}
        else:
            resu = {'code': -1, 'message': '访问密钥不存在'}
    elif reqQue:
        r_key = main_F.new_reply(reqQue)
        resu = {'code': 200, 'message': r_key}
    elif reqKey:
        if main_F.reply.__contains__(reqKey):
            if main_F.view_reply(reqKey).state == "Ready":
                resu = {'code': 10002, 'message': main_F.view_reply(reqKey).content}
            if main_F.view_reply(reqKey).state == "Generating":
                resu = {'code': 10003, 'message': main_F.view_reply(reqKey).content}
            if main_F.view_reply(reqKey).state == "Finish":
                resu = {'code': 200, 'message': main_F.view_reply(reqKey).content}
        else:
            resu = {'code': -1, 'message': '访问密钥不存在'}
    else:
        resu = {'code': 10001, 'message': '参数不能为空'}

    return json.dumps(resu, ensure_ascii=False)


if __name__ == "__main__":
    main_F.main_func()
    server.run(debug=False, port=8888, host='0.0.0.0')

测试一下：

第一次请求：

http://127.0.0.1:8888/api/aiAssist?q=什么是AVL树

返回key：

{"code": 200, "message": "PIVUNBXNDCRIBELB"}

带着key请求（用你实际返回的key而不是我演示里的key）：

http://127.0.0.1:8888/api/aiAssist?k=PIVUNBXNDCRIBELB

返回输出结果：

{"code": 200, "message": "AVL树（Adelson-Velsky and Landis Tree）是一种自平衡二叉查找树，它能够保持平衡，使得树的高度最多为log2(n)，其中n是树中节点的数目。AVL树是在1962年由G.M. Adelson-Velsky和E.M. Landis提出的。\n\nAVL树的特点是：\n\n1. 每个节点上附加一个额外的balance factor，用来记录子树的高度差。\n2. 每个节点的balance factor的值是-1,0,1之一，表示左子树的高度与右子树的高度的差。\n3. 在插入或删除节点时，会进行旋转操作，以维持平衡。\n\nAVL树的插入、删除和查找操作的时间复杂度都是O(log n)，因此它是一种高效的数据结构。"}

code为200表示已经输出完毕，否则会是10002（加载中）或10003（已经开始输出但未完成）

接着，带着key和新问题请求：

http://127.0.0.1:8888/api/aiAssist?k=PIVUNBXNDCRIBELB&q=那splay呢

会将key返回（不会变，还是第一次那个key）：

{"code": 200, "message": "PIVUNBXNDCRIBELB"}

接着，再次带着key请求：

http://127.0.0.1:8888/api/aiAssist?k=PIVUNBXNDCRIBELB

得到新的回复：

{"code": 200, "message": "Splay树（堆压树）是一种自平衡二叉查找树，它的特点是通过旋转操作来将最近访问的节点移动到树的根部，以提高查找、插入和删除操作的效率。Splay树的平均时间复杂度为O(log* n)，其中log* n是logarithmic star，表示以任意基数log的n次幂来计算的对数。\n\nSplay树的特点是：\n\n1. 每个节点上附加一个额外的访问计数器，用来记录节点被访问的次数。\n2. 在查找、插入和删除节点时，会进行旋转操作，以将最近访问的节点移动到树的根部。\n3. 旋转操作包括左旋、右旋、双旋等多种形式，以维持平衡。\n\nSplay树的查找、插入和删除操作的时间复杂度都是O(log* n)，因此它是一种高效的数据结构。但是，由于需要维护访问计数器，因此Splay树的空间"}

上面的回复虽然显示输出完毕了，即code为200，但由于max_tokens的限制，句子并没有结束。有需要可以将该变量改大点生成更长的回复。

控制台截图：

值得注意的是，如果在回复生成期间尝试用当前key请求新问题，并不会开始新的回复，而是会返回：

{"code": 10004, "message": "上次回复还没有结束"}

以上代码的注意事项还是这些：如果要实际大规模使用，还是需要自行改进，加入限制访问频率、从内存中清除过早的内容等等（根据AIReply的last_use属性，即上次访问的时间），建议借助数据库实现数据持久化。

那么，以上就是这篇的全部内容了，后续可能会优化或者改进，也会发一些相关的文章。

再次感谢UCloud和compshare.cn平台让我能用得起GPU云服务器进行测试x

本人自己也在学习中，如有问题请各路大佬多加指点，感谢。

OrikasaYuki

关注

11
点赞
踩
13

收藏

觉得还不错? 一键收藏
打赏
0
评论
【0基础】使用Mistral-7B模型搭建简易的聊天API（进阶篇）

本文是在基础篇之上的改进，【0基础】使用Mistral-7B模型搭建简易的聊天API（基础篇）-CSDN博客MistralAI团队去年发布mistral 7b模型至今，mistral7b就一直是同参数量内表现最亮眼的模型，逻辑能力、代码能力都非常强大，官方还发布了为聊天对话而微调的版本，也没有为个人用户自行下载部署设置任何额外的门槛（甚至当时在推特上直接发出磁力链接的时候我都被惊到了）。
复制链接

扫一扫