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

回顾

上文提速 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 不支持：

引自 deepseek 文档

虽然我们下载的 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

 1from zhipuai import ZhipuAI
 2import os
 3
 4client = ZhipuAI(api_key=os.getenv("GLM_4_PLUS_API_KEY"))
 5response = client.embeddings.create(
 6    model="embedding-3",
 7    input=[
 8        "美食非常美味，服务员也很友好。",
 9        "这部电影既刺激又令人兴奋。",
10        "阅读书籍是扩展知识的好方法。",
11    ],
12)
13print(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 的例子：

 1from typing import Any, List
 2from InstructorEmbedding import INSTRUCTOR
 3from llama_index.core.embeddings import BaseEmbedding
 4
 5class InstructorEmbeddings(BaseEmbedding):
 6    def __init__(
 7        self,
 8        instructor_model_name: str = "hkunlp/instructor-large",
 9        instruction: str = "Represent the Computer Science documentation or question:",
10        **kwargs: Any,
11    ) -> None:
12        super().__init__(**kwargs)
13        self._model = INSTRUCTOR(instructor_model_name)
14        self._instruction = instruction
15
16        def _get_query_embedding(self, query: str) -> List[float]:
17            embeddings = self._model.encode([[self._instruction, query]])
18            return embeddings[0]
19
20        def _get_text_embedding(self, text: str) -> List[float]:
21            embeddings = self._model.encode([[self._instruction, text]])
22            return embeddings[0]
23
24        def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
25            embeddings = self._model.encode(
26                [[self._instruction, text] for text in texts]
27            )
28            return embeddings
29
30        async def _get_query_embedding(self, query: str) -> List[float]:
31            return self._get_query_embedding(query)
32
33        async def _get_text_embedding(self, text: str) -> List[float]:
34            return self._get_text_embedding(text)

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

1        def _get_query_embedding(self, query: str) -> List[float]:
2            embeddings = self._model.encode([[self._instruction, query]])
3            return embeddings[0]
4
5        def _get_text_embedding(self, text: str) -> List[float]:
6            embeddings = self._model.encode([[self._instruction, text]])
7            return embeddings[0]

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

 1class ZhipuEmbeddings(BaseEmbedding):
 2    client: ZhipuAI = Field(default_factory=lambda: ZhipuAI(api_key=API_KEY))
 3
 4    def __init__(
 5        self,
 6        model_name: str = "embedding-3",
 7        **kwargs: Any,
 8    ) -> None:
 9        super().__init__(model_name=model_name, **kwargs)
10        self._model = model_name
11
12    def invoke_embedding(self, query: str) -> List[float]:
13        response = self.client.embeddings.create(model=self._model, input=[query])
14
15        # 检查响应是否成功
16        if response.data and len(response.data) > 0:
17            return response.data[0].embedding
18        else:
19            raise ValueError("Failed to get embedding from ZhipuAI API")
20
21    def _get_query_embedding(self, query: str) -> List[float]:
22        return self.invoke_embedding(query)
23
24    def _get_text_embedding(self, text: str) -> List[float]:
25        return self.invoke_embedding(text)
26
27    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
28        return [self._get_text_embedding(text) for text in texts]
29
30    async def _aget_query_embedding(self, query: str) -> List[float]:
31        return self._get_query_embedding(query)
32
33    async def _aget_text_embedding(self, text: str) -> List[float]:
34        return self._get_text_embedding(text)
35
36    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
37        return self._get_text_embeddings(texts)

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

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

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

以下是完整代码：

  1import os
  2import sys
  3import logging
  4from zhipuai import ZhipuAI
  5from typing import Any, List
  6from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
  7from llama_index.core.embeddings import BaseEmbedding
  8from llama_index.core.llms import (
  9    CustomLLM,
 10    CompletionResponse,
 11    CompletionResponseGen,
 12    LLMMetadata,
 13)
 14from llama_index.core.llms.callbacks import llm_completion_callback
 15from dotenv import load_dotenv
 16from functools import cached_property
 17from pydantic import Field
 18
 19# 配置日志
 20logging.basicConfig(level=logging.INFO)
 21logger = logging.getLogger(__name__)
 22
 23# 从环境变量获取 API 密钥
 24load_dotenv()
 25
 26API_KEY = os.getenv("GLM_4_PLUS_API_KEY")
 27if not API_KEY:
 28    raise ValueError("GLM_4_PLUS_API_KEY environment variable is not set")
 29
 30class GLM4LLM(CustomLLM):
 31    @cached_property
 32    def client(self):
 33        return ZhipuAI(api_key=API_KEY)
 34
 35    @property
 36    def metadata(self) -> LLMMetadata:
 37        return LLMMetadata()
 38
 39    def chat_with_glm4(self, system_message, user_message):
 40        response = self.client.chat.completions.create(
 41            model="glm-4-plus",
 42            messages=[
 43                {
 44                    "role": "system",
 45                    "content": system_message,
 46                },
 47                {
 48                    "role": "user",
 49                    "content": user_message,
 50                },
 51            ],
 52            stream=True,
 53        )
 54        return response
 55
 56    @llm_completion_callback()
 57    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
 58        response = self.chat_with_glm4("你是一个聪明的 AI 助手", prompt)
 59        full_response = "".join(
 60            chunk.choices[0].delta.content
 61            for chunk in response
 62            if chunk.choices[0].delta.content
 63        )
 64        return CompletionResponse(text=full_response)
 65
 66    @llm_completion_callback()
 67    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
 68        response = self.chat_with_glm4("你是一个聪明的 AI 助手", prompt)
 69
 70        def response_generator():
 71            response_content = ""
 72            for chunk in response:
 73                if chunk.choices[0].delta.content:
 74                    response_content += chunk.choices[0].delta.content
 75                    yield CompletionResponse(
 76                        text=response_content, delta=chunk.choices[0].delta.content
 77                    )
 78
 79        return response_generator()
 80
 81class ZhipuEmbeddings(BaseEmbedding):
 82    client: ZhipuAI = Field(default_factory=lambda: ZhipuAI(api_key=API_KEY))
 83
 84    def __init__(
 85        self,
 86        model_name: str = "embedding-3",
 87        **kwargs: Any,
 88    ) -> None:
 89        super().__init__(model_name=model_name, **kwargs)
 90        self._model = model_name
 91
 92    def invoke_embedding(self, query: str) -> List[float]:
 93        response = self.client.embeddings.create(model=self._model, input=[query])
 94
 95        # 检查响应是否成功
 96        if response.data and len(response.data) > 0:
 97            return response.data[0].embedding
 98        else:
 99            raise ValueError("Failed to get embedding from ZhipuAI API")
100
101    def _get_query_embedding(self, query: str) -> List[float]:
102        return self.invoke_embedding(query)
103
104    def _get_text_embedding(self, text: str) -> List[float]:
105        return self.invoke_embedding(text)
106
107    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
108        return [self._get_text_embedding(text) for text in texts]
109
110    async def _aget_query_embedding(self, query: str) -> List[float]:
111        return self._get_query_embedding(query)
112
113    async def _aget_text_embedding(self, text: str) -> List[float]:
114        return self._get_text_embedding(text)
115
116    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
117        return self._get_text_embeddings(texts)
118
119# 设置环境变量，禁用 tokenizers 的并行处理
120os.environ["TOKENIZERS_PARALLELISM"] = "false"
121
122def run_glm4_query_with_embeddings(query: str):
123    # 从指定目录加载文档数据
124    documents = SimpleDirectoryReader("data").load_data()
125
126    # 设置 LLM 和嵌入模型
127    Settings.llm = GLM4LLM()
128    Settings.embed_model = ZhipuEmbeddings()
129
130    # 创建索引和查询引擎
131    index = VectorStoreIndex.from_documents(documents)
132    query_engine = index.as_query_engine(streaming=True)
133
134    # 执行查询
135    print("GLM-4 查询结果：")
136    response = query_engine.query(query)
137
138    # 处理并输出响应
139    if hasattr(response, "response_gen"):
140        # 流式输出
141        for text in response.response_gen:
142            print(text, end="", flush=True)
143            sys.stdout.flush()  # 确保立即输出
144    else:
145        # 非流式输出
146        print(response.response, end="", flush=True)
147
148    print("\n 查询完成")

效果

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

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

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

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

最后

我已将文章中涉及到的相关代码上传至：https://github.com/xiaobox/llamaindex_test

这个仓库中包含了最新几篇文章中的所有 demo 代码，大家可以自行查看。