AI知识库

53AI知识库

学习大模型的前沿技术与行业应用场景


提升RAG应用性能:使用智谱AI的GLM-4和Embedding-3模型优化文档检索
发布日期:2024-10-08 16:03:16 浏览次数: 1878 来源:小盒子的技术分享


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

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

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

  1. 文档嵌入(Document Embedding)

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

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

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

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

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

    我们再通过回顾 2 张图片来比较直观地了解下 embeddingLLM 在 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 模型 ,这样方便模型的微调及不限量的调用。速度也会快许多



    53AI,企业落地应用大模型首选服务商

    产品:大模型应用平台+智能体定制开发+落地咨询服务

    承诺:先做场景POC验证,看到效果再签署服务协议。零风险落地应用大模型,已交付160+中大型企业

    联系我们

    售前咨询
    186 6662 7370
    预约演示
    185 8882 0121

    微信扫码

    与创始人交个朋友

    回到顶部

     
    扫码咨询