基于 LlamaIndex 的嵌入和重排查询

Marshal
1,756 阅读4分钟

本文是基于MarshalW/jupyterlab-demos/blob/main/proto-embedding-reranker.ipynb 笔记的说明。

笔记的大部分内容，都可以执行在 4GB 显存的 nvidia 设备上（有关配置见 在 4GB 显存下运行 LLM 基础开发环境 ）。

本文的目的是，通过 JupyterLab 笔记，最简明的介绍基于 LlamaIndex 使用嵌入（Embedding）和重排（Rerank），对大文档中的知识点进行检索和查询的方法。

另外，给出了初步结论：目前至少需要 Qwen2:7B 以上的模型，才能达到基本可用的效果。

低显存设备可用于学习，但因为模型偏小，准确率很低。

基本思路

  • 大模型在不断提升上下文长度，可以将更大的文档一次性加入提示词上下文
    • kimi 可以最多加载20万字上下文（甚至100万字）
    • 本地模型，Qwen2:7b 也达到了32K
  • 但是 RAG 依然有存在的价值，原因是：
    • 云端模型在对话过程中将消费大量 token，费用昂贵
    • 本地模型上下文长度受硬件限制很难大幅提高
    • 很多场景的数据量还是会超过最大上下文长度限制
  • 最基本的 RAG 一般涉及4个步骤
    • 加载数据，比如通过文件系统加载本地的 txt 文件
      • 加载的文本文档，切分为一定长度的片段
      • 切分方式，基于字符长度，基于自然语言断句，或者基于语义等
    • 索引阶段
      • 通过嵌入模型，将文档分片用向量（嵌入）表示
      • 不同的模型，生成的嵌入是不同的，好的嵌入模型，召回率更高
      • 嵌入保存在向量数据库中，笔记中使用 Qdrant
      • 创建索引很耗时，会将相关数据存储，供后续阶段读取
    • 检索阶段
      • 将用户的查询文字，和索引中的向量做运算，文本内容一般使用余弦相似度算法
      • 运算结果取 top_k，就是取前多少个相邻向量
      • 运算结果通常与语义相似性紧密相关
      • 但是，基于词向量的运算，可能无法完全捕捉到所有语义细节
    • 查询阶段
      • 将检索到的向量的原文，加入到提示词上下文
      • 提示词将包括：检索到的文本+用户输入的提示词
      • 将合成的提示词传给 LLM，让 LLM 组织文字形成最终的回答
  • rerank 是对 RAG 的检索阶段的改进
    • 余弦相似度算法的 top_k，可能存在更准确的文本片段排在后面
    • rerank 模型基于查询提示词的语义，对 top_k 结果重新排序，并取其中的 top_n
    • rerank 适用的场景
      • 必须取大量的向量才能涵盖到需要的文本片段，即 top_k 很大，比如 1000
      • 如果直接将这些片段合成为提示词，会造成
        • 性能下降
        • 干扰 LLM 正确的回答
        • 超出提示词上下文最大值
      • 在上述情况下，rerank 能提升大模型回答的准确率，并显著提高响应速度

对显存的要求

笔记代码中使用的模型，最多占用 2.5GB 左右显存：

  • Qwen2:1.5b: 1508MB
  • quentinz/bge-large-zh-v1.5: 1080MB

reranker，可以不使用 GPU 加载，

  • BAAI/bge-reranker-base: 需要大约 1.5GB 系统内存

配置使用上述模型，准确率很低，但是可以跑通代码。

实现效果

完整代码笔记见 MarshalW/jupyterlab-demos/blob/main/proto-embedding-reranker.ipynb

演示视频：

Kapture 2024-07-15 at 17.20.52.gif

实现步骤说明

基本步骤按照笔记 MarshalW/jupyterlab-demos/blob/main/proto-embedding-reranker.ipynb 即可。

有一些需要配置的，下面说明一下。

Ollama 加载多个模型

Ollama 默认只加载一个模型，如果再加载其他模型，则会替换掉当前模型。这样做的目的是节省显存。

为了能同时加载 LLM（比如 Qwen2:1.5b） 模型和嵌入模型（比如 bge-large-zh-v1.5）, 需要编辑 llm-proto .env 文件：

OLLAMA_FLASH_ATTENTION=1
OLLAMA_NUM_PARALLEL=4
OLLAMA_MAX_LOADED_MODELS=2 # 可同时加载2个模型

# jubyterlab
JUPYTER_TOKEN="password"

Ollama 下载嵌入模型

# 进入 ollama 容器
docker exec -it llms bash

# 拉取模型
ollama pull quentinz/bge-large-zh-v1.5

reranker 模型的下载和加载

Ollama 没有可支持的 reranker 模型，只能自行加载加载。

bge-reranker-base 下载地址。

在 JupyterLab 的命令行中运行：

# 安装git lfs支持
apt-get update && apt-get install git-lfs -y

mkdir /models
cd /models

# 下载 - 可能需要代理
git clone https://huggingface.co/BAAI/bge-reranker-base

按照绝对路径加载模型：

rerank = SentenceTransformerRerank(
    model="/models/bge-reranker-base", top_n=5, 
    device='cpu',
)

reranker 模型在 GPU 加载

如果显存足够大，可以考虑 reranker 模型也通过 GPU 加载，在笔记中为了适配低显存，用的 CPU 加载, 注视掉 device='cpu',，默认情况下是 GPU 加载。：

rerank = SentenceTransformerRerank(
    model="/models/bge-reranker-base", top_n=5, 
   # device='cpu',
)

总结

  • 通过笔记可以掌握 RAG 的基本做法，嵌入检索和查询
  • 执行效果很依赖模型质量，一般越大的模型准确率越高
  • reranker 可以通过 LLM 执行，LlamaIndex 分装了 LLMRerank
  • 笔记最后的部分 - 基于更大的本地模型
    • 展示了 Qwen2:7b 等更大尺寸模型的效果，准确率较高
    • 这部分加载模型的显存至少要 7GB
avatar
文章
阅读
粉丝