基于 ChatGLM2 和 OpenVINO™ 打造中文聊天助手

点击蓝字

关注我们,让开发变得更有趣

作者 | 英特尔 AI 软件工程师 杨亦诚

排版 | 李擎

基于ChatGLM2和OpenVINO™打造中文聊天助手

ChatGLM 是由清华大学团队开发的是一个开源的、支持中英双语的类 ChatGPT 大语言模型，它能生成相当符合人类偏好的回答,  ChatGLM2 是开源中英双语对话模型 ChatGLM 的第二代版本，在保留了初代模型对话流畅、部署门槛较低等众多优秀特性的基础之上，通过全面升级的基座模型，带来了更强大的性能，更长的上下文，并且该模型对学术研究完全开放，登记后亦允许免费商业使用。接下来我们分享一下如何基于 ChatGLM2-6B 和 OpenVINO™ 工具套件来打造一款聊天机器人。

项目仓库地址：https://github.com/OpenVINO-dev-contest/chatglm2.openvino

注1：由于 ChatGLM2-6B 对在模型转换和运行过程中对内存的占用较高，推荐使用支持 128Gb 以上内存的的服务器终端作为测试平台。

注2：本文仅分享部署 ChatGLM2-6B 原始预训练模型的方法，如需获得自定义知识的能力，需要对原始模型进行 Fine-tune；如需获得更好的推理性能，可以使用量化后的模型版本。

OpenVINO™

模型导出

第一步，我们需要下载 ChatGLM2-6B 模型，并将其导出为 OpenVINO™ 所支持的IR格式模型进行部署，由于 ChatGLM 团队已经将 6B 版本的预训练模型发布在 Hugging Face 平台上，支持通过 Transformer 库进行推理，但不支持基于Optimum 的部署方式（可以参考 Llama2 的文章），因此这里我们需要提取 Transformer 中的 ChatGLM2 的 PyTorch 模型对象，并实现模型文件的序列化。主要步骤可以分为：

1.获取 PyTorch 模型对象

通过Transformer库获取PyTorch对象，由于目前Transformer中原生的ModelForCausalLM类并不支持ChatGLM2模型架构，因此需要添加trust_remote_code=True参数，从远程模型仓库中获取模型结构信息，并下载权重。

model = AutoModel.from_pretrained(args.model_id,
                                  trust_remote_code=True).float()

2.模拟并获取模型的输入输出参数

在调用 torch.onnx.export 接口将模型对象导出为 ONNX 文件之前，我们首先需要获取模型的输入和输出信息。由于 ChatGLM2 存在 KV cache 机制，因此这个步骤中会模拟第一次文本生成时不带 cache 的输入，并将其输出作为第二次迭代时