知识库

提升RAG应用性能：使用智谱AI的GLM-4和Embedding-3模型优化文档检索

2024 年 10 月 8 日

上文提速 RAG 应用：用 DeepSeek API 替换本地 Ollama 模型，LlamaIndex 实战解析我们介绍了如何通过 DeepSeek 的 API 调用 DeepSeek v2.5 模型替换通过 Ollama 调用本地下载好的 Qwen2.5 模型。

这样做的目的是想通过 API 调用远程部署好的 LLM 给我们的 RAG 应用提提速。不然由于本地个人电脑计算资源的不足（我的电脑没有 GPU）会导致 RAG 应用运行缓慢。

在我们的 RAG 应用中分别使用了两个模型，一个是 embedding 模型，它的作用有这么几点：

文档嵌入（Document Embedding）

表示文档：将文档转换为高维向量（embeddings），这些向量能够捕捉文档的语义信息。
相似度计算：通过计算查询和文档嵌入之间的相似度，找到与查询最相关的文档。
查询嵌入（Query Embedding）

表示查询：将用户的查询转换为高维向量，这些向量能够捕捉查询的语义信息。
检索相关文档：通过计算查询嵌入和文档嵌入之间的相似度，找到与查询最相关的文档

文档检索（Document Retrieval）

高效检索：通过向量数据库（如 Faiss、Annoy 等），快速找到与查询最相关的文档。
相关性排序：根据相似度得分对检索到的文档进行排序，选择最相关的文档作为生成回答的依据。

生成回答（Answer Generation）

融合信息：将检索到的相关文档与查询结合，生成高质量的回答。
上下文感知：利用检索到的文档作为上下文，生成更加准确和丰富的回答。

其中第 4 点，要结合 LLM 来完成。所以这也是我们在 RAG 应用中使用第二个模型--大语言模型（LLM）的意义。

我们再通过回顾 2 张图片来比较直观地了解下 embedding 和 LLM 在 RAG 中的作用：

embedding 过程
RAG

问题

上文遗留的问题很明显，因为我们需要使用的 2 个模型通过 DeekSeek 的 API 只替换了其中的LLM，而 embedding 模型仍然用的是本地的。没有替换是因为 DeepSeek 的 API 不支持：

虽然我们下载的 embedding 模型 BAAI/bge-base-zh-v1.5 比较小巧，在本地运行的速度也还行，但我还是想试一下调用远程部署好的更大更优秀的 embedding 模型后会怎样？

于是我将目光转向了另一个很知名，同样很优秀的国产 AI 公司 智谱 AI

智谱 AI

这两年 AI 的发展如火如荼，以 ChatGPT 为代表的一众 AIGC 应用深入人心，这些应用的背后都少不了大语言模型的支持。然而对于国内用户使用这些产品仍然有门槛。大家不禁想找到一个能打的国产 AI 产品。

去年秋天我还在迷信 ChatGPT 的能力是“宇宙无敌”，直到我体验了智谱 AI 旗下的 智谱清言 我才觉得国产 AI 产品在中文语料下的能力并不比别人差。

智谱是由清华大学计算机系技术成果转化而来的公司。它的发展很快。目前可供用户使用的各类模型 20 余个。其中包括：

大规模语言模型 GLM-4
视频生成模型 CogVideoX
代码模型 CodeGeeX-4
图片生成模型 CogView-3
嵌入式模型 Embedding-3
......

智谱在开源领域也做出了极大贡献，上面列举的这些模型都能在 HuggingFace 或 GitHub 上找到开源的版本。

智谱 AI 最让我们熟悉的产品是其 C 端 AIGC 产品 智谱清言

在中文语料下，它的问答质量不比 GPT-4 差！

LlamaIndex 集成 Zhipu embedding

通过查看智谱 AI 大模型开放平台的文档得知它有两款 embedding 模型可以通过 API 调用

于是决定将 Embedding-3 试着集成到 LlamaIndex 中。

当然，调用 API 首先你要有 API Key 以及可用的 tokens，这个我们在之前的文章我介绍过，一般是需要付费的，智谱 AI 会给新老用户赠送一些 tokens，之前赠送给了我 1000w tokens ，所以下面的示例我就用这些免费的 tokens。

简单 demo

我们先根据文档写一个最简单的模型调用 demo

from zhipuai import ZhipuAI
import os

client = ZhipuAI(api_key=os.getenv("GLM_4_PLUS_API_KEY"))
response = client.embeddings.create(
    model="embedding-3",
    input=[
        "美食非常美味，服务员也很友好。",
        "这部电影既刺激又令人兴奋。",
        "阅读书籍是扩展知识的好方法。",
    ],
)
print(response)

它响应的输出是这样的：

这输出的一片数字是啥？

这里简单解释一下：嵌入是将文字、图像或其他类型的数据转换成一系列数字（向量）的过程。这个向量在高维空间中代表了原始数据的语义信息。你看到的那一长串数字（如 -0.019210815, -0.0023460388, 0.010299683 等）就是嵌入向量的具体值。每个数字代表向量在某个维度上的值，这些数字虽然看起来没有明显意义，但它们在高维空间中编码了输入文本的语义信息。相似的文本会产生相似的向量，这使得我们可以进行语义相似度比较。这种表示方法使得机器能够更好地"理解"和处理文本数据。

能够正常输出，代表模型调用成功。

和 LlamaIndex 集成

在之前的文章中我们已经通过 Custom LLM 的方式将 LlamaIndex 和 GLM-4 集成在一起了，也就是在 RAG 应用中使用的框架是 LlamaIndex ，调用的 LLM 是 GLM-4。

同理，现在我们要把 embedding 模型也同 LlamaIndex 集成起来，这样我们自己写的这个 RAG 应用的技术组合就是 RAG = LlamaIndex +GLM-4 + Embedding-3

和 LLM 一样，在 LlamaIndex 文档的 embedding 模型兼容列表中并没有 Zhipu 的 Embedding-3 ，仍然需要通过自定义的方式来实现。

这是文档中给的自定义 embedding 的例子：

from typing import Any, List
from InstructorEmbedding import INSTRUCTOR
from llama_index.core.embeddings import BaseEmbedding

class InstructorEmbeddings(BaseEmbedding):
    def __init__(
        self,
        instructor_model_name: str = "hkunlp/instructor-large",
        instruction: str = "Represent the Computer Science documentation or question:",
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._model = INSTRUCTOR(instructor_model_name)
        self._instruction = instruction

        def _get_query_embedding(self, query: str) -> List[float]:
            embeddings = self._model.encode([[self._instruction, query]])
            return embeddings[0]

        def _get_text_embedding(self, text: str) -> List[float]:
            embeddings = self._model.encode([[self._instruction, text]])
            return embeddings[0]

        def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
            embeddings = self._model.encode(
                [[self._instruction, text] for text in texts]
            )
            return embeddings

        async def _get_query_embedding(self, query: str) -> List[float]:
            return self._get_query_embedding(query)

        async def _get_text_embedding(self, text: str) -> List[float]:
            return self._get_text_embedding(text)

仔细看的话，实际上只需要实现 2 个方法即可，下面的方法都会调用这两个方法：

        def _get_query_embedding(self, query: str) -> List[float]:
            embeddings = self._model.encode([[self._instruction, query]])
            return embeddings[0]

        def _get_text_embedding(self, text: str) -> List[float]:
            embeddings = self._model.encode([[self._instruction, text]])
            return embeddings[0]

这里我们可以新建一个自定义的 embedding 类：

class ZhipuEmbeddings(BaseEmbedding):
    client: ZhipuAI = Field(default_factory=lambda: ZhipuAI(api_key=API_KEY))

    def __init__(
        self,
        model_name: str = "embedding-3",
        **kwargs: Any,
    ) -> None:
        super().__init__(model_name=model_name, **kwargs)
        self._model = model_name

    def invoke_embedding(self, query: str) -> List[float]:
        response = self.client.embeddings.create(model=self._model, input=[query])

        # 检查响应是否成功
        if response.data and len(response.data) > 0:
            return response.data[0].embedding
        else:
            raise ValueError("Failed to get embedding from ZhipuAI API")

    def _get_query_embedding(self, query: str) -> List[float]:
        return self.invoke_embedding(query)

    def _get_text_embedding(self, text: str) -> List[float]:
        return self.invoke_embedding(text)

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return [self._get_text_embedding(text) for text in texts]

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)

    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return self._get_text_embeddings(texts)

在利用 LlamaIndex 调用时，将 embed_model 设置为自定义类就可以了：

  # 设置 LLM 和嵌入模型
    Settings.llm = GLM4LLM()
    Settings.embed_model = ZhipuEmbeddings()

这样我们的 RAG 应用就把智谱 AI 的 GLM-4 和 Embedding-3 一起使用上了。

以下是完整代码：

import os
import sys
import logging
from zhipuai import ZhipuAI
from typing import Any, List
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.llms import (
    CustomLLM,
    CompletionResponse,
    CompletionResponseGen,
    LLMMetadata,
)
from llama_index.core.llms.callbacks import llm_completion_callback
from dotenv import load_dotenv
from functools import cached_property
from pydantic import Field

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 从环境变量获取 API 密钥
load_dotenv()

API_KEY = os.getenv("GLM_4_PLUS_API_KEY")
if not API_KEY:
    raise ValueError("GLM_4_PLUS_API_KEY environment variable is not set")

class GLM4LLM(CustomLLM):
    @cached_property
    def client(self):
        return ZhipuAI(api_key=API_KEY)

    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata()

    def chat_with_glm4(self, system_message, user_message):
        response = self.client.chat.completions.create(
            model="glm-4-plus",
            messages=[
                {
                    "role": "system",
                    "content": system_message,
                },
                {
                    "role": "user",
                    "content": user_message,
                },
            ],
            stream=True,
        )
        return response

    @llm_completion_callback()
    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        response = self.chat_with_glm4("你是一个聪明的 AI 助手", prompt)
        full_response = "".join(
            chunk.choices[0].delta.content
            for chunk in response
            if chunk.choices[0].delta.content
        )
        return CompletionResponse(text=full_response)

    @llm_completion_callback()
    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
        response = self.chat_with_glm4("你是一个聪明的 AI 助手", prompt)

        def response_generator():
            response_content = ""
            for chunk in response:
                if chunk.choices[0].delta.content:
                    response_content += chunk.choices[0].delta.content
                    yield CompletionResponse(
                        text=response_content, delta=chunk.choices[0].delta.content
                    )

        return response_generator()

class ZhipuEmbeddings(BaseEmbedding):
    client: ZhipuAI = Field(default_factory=lambda: ZhipuAI(api_key=API_KEY))

    def __init__(
        self,
        model_name: str = "embedding-3",
        **kwargs: Any,
    ) -> None:
        super().__init__(model_name=model_name, **kwargs)
        self._model = model_name

    def invoke_embedding(self, query: str) -> List[float]:
        response = self.client.embeddings.create(model=self._model, input=[query])

        # 检查响应是否成功
        if response.data and len(response.data) > 0:
            return response.data[0].embedding
        else:
            raise ValueError("Failed to get embedding from ZhipuAI API")

    def _get_query_embedding(self, query: str) -> List[float]:
        return self.invoke_embedding(query)

    def _get_text_embedding(self, text: str) -> List[float]:
        return self.invoke_embedding(text)

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return [self._get_text_embedding(text) for text in texts]

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)

    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return self._get_text_embeddings(texts)

# 设置环境变量，禁用 tokenizers 的并行处理
os.environ["TOKENIZERS_PARALLELISM"] = "false"

def run_glm4_query_with_embeddings(query: str):
    # 从指定目录加载文档数据
    documents = SimpleDirectoryReader("data").load_data()

    # 设置 LLM 和嵌入模型
    Settings.llm = GLM4LLM()
    Settings.embed_model = ZhipuEmbeddings()

    # 创建索引和查询引擎
    index = VectorStoreIndex.from_documents(documents)
    query_engine = index.as_query_engine(streaming=True)

    # 执行查询
    print("GLM-4 查询结果：")
    response = query_engine.query(query)

    # 处理并输出响应
    if hasattr(response, "response_gen"):
        # 流式输出
        for text in response.response_gen:
            print(text, end="", flush=True)
            sys.stdout.flush()  # 确保立即输出
    else:
        # 非流式输出
        print(response.response, end="", flush=True)

    print("\n 查询完成")

效果

从最终的使用效果上看，速度上不如之前使用本地 embedding 模型 BAAI/bge-base-zh-v1.5 快。因为执行了多次 Http 远程调用：

所以我又查了一下文档看看有没有办法提提速：

虽然有一个 dimensions 参数，虽然我感觉设置的越小维度越小数据也越少，那么速度可能更快，但实际测试下来速度并没有明显变化。其主要原因还是：它是同步调用的

看来从 API 上没办法提速了，只能在编程模型上想办法了，这里就不多说了。这里我认为，最好的方式还是在一个资源充足的服务器中部署一个开源的 embedding 模型，这样方便模型的微调及不限量的调用。速度也会快许多

立即咨询 JOTO

JOTO 提供覆盖企业智能体规划与搭建、AI 平台私有化部署、RAG 知识工程、AI 安全治理、FDE 驻场共创及持续运营优化的全周期 AI 落地服务，帮助企业把验证中的 AI 能力转化为安全、可控、可持续迭代的生产力。联系 JOTO 获取 AI 落地咨询

想把这些做法用到你的业务里?

留下你的场景和痛点,我们帮你判断从哪一步开始。

联系我们