Elasticsearch 混合检索优化大模型 RAG 任务_基于倒数的融合排序

Elastic 社区在自然语言处理上面做的很不错官方博客更新速度也很快，现阶段大模型的应用场景主要在 Rag 和 Agent 上，国内 Rag（Retrieval-Augmented Generation 检索增强生成） 的尤其多，而搜索对于 Elasticsearch 来说是强项特别是 8.9 之后的版本提供了 ESRE 模块（集成了高级相关性排序如 BM25f、强大的矢量数据库、自然语言处理技术、与第三方模型如 GPT-3 和 GPT-4 的集成，并支持开发者自定义模型与应用），经过我的各种尝试在 Elasticsearch 上做 NLP 是一个很不错的选择，要做大规模的 RAG 任务甚至是针对图像、声音、多模态、关键词等大数据量的向量召回且搭配生成式模型这种复杂的业务场景 Elasticsearch 是天生支持的。此篇文章主要记录混合检索（BM25 +HNSW）倒数融合排序（RRF）完整测试。

官博有几篇不错的文章可以看看：

先说一下 RAG 任务的流程，以民法典为例 LLM 可以在现有资料上分析出确切的回答：
文档分割 -> 文本向量化 -> 问句向量化 -> 向量相似 top k个 -> 拼接 prompt 上下文  -> 提交给 LLM 生成回答。

1.混合检索

全文检索 + ANN 检索。因为全文检索能查找更加准确的文档，直观都会感觉比单一的相似度检索更强。一个混合检索的查询语句例如：

{
  "query": {
    "bool": {
      "must": [
        { "match": {"content": {"query": "结婚领证登记需要双发到场吗？","boost": 1}}}
      ]
    }
  },
  "knn": {
    "field": "content_embed",
    "k": 5,
    "num_candidates": 100,
    "query_vector": []   // 向量、省略
  },
  "size": 5
}

2.倒数融合排序

倒数排序融合 - Reciprocal rank fusion：
由于全文搜索及向量搜索是使用不同的算法进行打分的，这就造成把两个不同搜索结果综合起来统一排名的困难。向量搜索的分数处于 0-1.0 之间，而全文搜索的结果排名分数可能是高于10或者更大的值。我们需要一种方法把两种搜索方法的结果进行综合处理，并得出一个唯一的排名。
倒数排序融合（RRF）是一种将具有不同相关性指标的多个结果集组合成单个结果集的方法。 
RRF 无需调优，不同的相关性指标也不必相互关联即可获得高质量的结果。该方法的优势在于不利用相关分数，而仅靠排名计算。相关分数存在的问题在于不同模型的分数范围差。
针对不同的 RAG 任务有不同的处理方式比如 法律、历史、人文类型的任务还可以加入命名实体识别 。或者使用其他语义转换模型将长文本总结为短文本。将拆分的长文本先调用 embed 转为向量后存储到 index 上。然后执行混合检索。

3.Embedding

第一步是文本向量化，这一步可以放在客户端做也可以放在 Elasticsearch 服务端做，不过模型推理是 Elasticsearch 新版中的重大功能，下面演示如何做。

在抱脸上直接搜索  sentence-similarity 模型，最靠前的就是 bge 由智源开源，基本上从去年开始一直是榜一，输入 zh 筛选中文：

使用 langchain 测试推理，模型输出是  dim=1024：

ElasticSearch支持最大 2048，目前 Es 还不支持非固定长度的向量，Elasticsearch 提供了 Eland 工具用于 pytorch 模型的推理和上传，源码安装该工具：

git clone https://github.com/elastic/eland
pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/
python setup.py install

然后执行上传脚本：
eland_import_hub_model --url http://192.168.197.128:9200 --hub-model-id .\Langchain-Chatchat-0.2.10\model\bge-large-zh-v1.5 --task-type text_embedding --start --clear-previous

上传过程不太顺利发现源码有一些问题需要修改，大致两处：
eland_import_hub_model.py  =>    上传前会把模型和一些文件放到临时目录，因为我的 windwos user name 是中文会找不到路径。直接将 tmp 写死即可。
            with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_dir = 'C:\\tmp' 
transformers.py    =>        函数里面将 token 这个参数去掉
            # model = AutoModel.from_pretrained(model_id, token=token, torchscript=True)
            model = AutoModel.from_pretrained(model_id, torchscript=True)

等待执行完成：

上传成功后在 kibana 模型管理位置点击 Synchronize your jobs and trained models.，同步一下刚刚上传的模型看到，调用推理接口，复制模型id，可以看到模型输出和前面 embed_demo.py 中测试的一样：
POST _ml/trained_models/m_workspace__langchain-chatchat-0.2.10__model__bge-large-zh-v1.5/_infer
{
  "docs": [
    {"text_field": "你好，请问你在干什么？"}
  ]
}

4.文本分割

向量 dim=1024 是无法将一个超长文本完整的语义全部嵌入的，且大模型 token 的限制需要将文档进行分割，最简单的做法是指定 chunk_size（单个文档token数） 和 chunk_overlap（向量文档重叠token数）对文档进行分割，也有按句分割的做法，更加准确的是使用现成的语义分割模型，可以看看 github 上 Langchain-Chatchat 这个项目，提供了多种分割方式：

5.部署 LLM 做增强生成

对于 RAG 任务，更大参数量的 LLM 对效果并没有显著提升， 即使是最小参数量的大模型也涵盖了基本的理解能力，这里部署清华 ChatGLM-6b  int4 量化模型 6G显存就够，这样可以将 token 开到很大。

git clone https://github.com/THUDM/ChatGLM-6B
pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/
git clone https://huggingface.co/THUDM/chatglm-6b-int4

模型 README.md 中有测试代码，替换一下模型路径就可以了：

然后写一个 ELasticsearch Query 例子，根据搜索文档拼接 Prompt 做问答，Java 完整代码：

package tool.elk;
 
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.nio.entity.NStringEntity;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.client.*;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.xcontent.XContentType;
import java.io.BufferedReader;
import java.io.FileReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
 
/**
 *   @desc : elatcisearch rag 测试
 *   @auth : tyf
 *   @date : 2024-04-16 10:06:24
*/
public class RAGDemo {
 
    public static String es_host = "192.168.197.128";
    public static Integer es_port = 9200;
 
    public static String llm_host = "http://0.0.0.0:8000";
 
    public static RestHighLevelClient highLevelClient;
    public static RestClient lowLevelClient;
    static {
        String[] ipArr = es_host.split(",");
        HttpHost[] httpHosts = new HttpHost[ipArr.length];
        for (int i = 0; i < ipArr.length; i++) {
            httpHosts[i] = new HttpHost(ipArr[i], es_port, "http");
        }
        RestClientBuilder builder = RestClient.builder(httpHosts);
        highLevelClient = new RestHighLevelClient(builder);
        lowLevelClient = highLevelClient.getLowLevelClient();
        System.out.println("初始化成功");
    }
 
    
    // 索引名称
    public static String indexName = "doc_split";
    // 索引 mapping
    public static String indexMapping =
            "{\n" +
                    "  \"settings\": {\n" +
                    "    \"number_of_shards\": 1,\n" +
                    "    \"number_of_replicas\": 0\n" +
                    "  },\n" +
                    "  \"mappings\": {\n" +
                    "    \"properties\": {\n" +
                    "      \"content\": {\n" +
                    "        \"type\": \"text\"\n" +
                    "      },\n" +
                    "      \"timestamp\": {\n" +
                    "        \"type\": \"long\"\n" +
                    "      },\n" +
                    "      \"content_embed\": {\n" +
                    "        \"type\": \"dense_vector\",\n" +
                    "        \"dims\": 1024,\n" +
                    "        \"index\": true,\n" +
                    "        \"similarity\": \"cosine\"\n" +
                    "      }\n" +
                    "    }\n" +
                    "  }\n" +
                    "}";
 
    // embed 模型编号
    public static String modelId = "m_workspace__langchain-chatchat-0.2.10__model__bge-large-zh-v1.5";
 
    // 文档召回 _score 阈值
    public static double scoreThreshold = 3d;
 
    // 本地文档路径
    public static String docPath = "C:\\Users\\唐于凡\\Desktop\\中华人民共和国民法典.txt";
 
 
 
    // 创建索引
    public static void createIndex() throws Exception{
//        System.out.println(indexMapping);
        // 索引不存在则创建
        GetIndexRequest request1 = new GetIndexRequest(indexName);
        boolean response1 = highLevelClient.indices().exists(request1, RequestOptions.DEFAULT);
        if(!response1){
            CreateIndexRequest request2 = new CreateIndexRequest(indexName);
            request2.source(indexMapping, XContentType.JSON);
            highLevelClient.indices().create(request2, RequestOptions.DEFAULT);
        }
    }
 
    // 读取并拆分文档
    public static List<String> parseDoc(int chunkSize,int chunkOverlap) throws Exception{
        List<String> splitTexts = new ArrayList<>();
        try (BufferedReader br = new BufferedReader(new FileReader(docPath))) {
            StringBuilder sb = new StringBuilder();
            String line;
            while ((line = br.readLine()) != null) {
                // 去掉没用的空格
                line = line.trim();
                if (!line.isEmpty()) {
                    sb.append(line).append(" "); // 可以根据需要调整分隔符
                }
            }
            String fullText = sb.toString().trim();
            // 拆分文本
            for (int i = 0; i < fullText.length(); i += chunkSize - chunkOverlap) {
                if (i + chunkSize < fullText.length()) {
                    splitTexts.add(fullText.substring(i, i + chunkSize));
                } else {
                    splitTexts.add(fullText.substring(i));
                }
            }
        }
        System.out.println("文档总数："+splitTexts.size());
        return splitTexts;
    }
 
 
    // 调用 embed 模型转为向量
    public static Object embedDoc(String text){
 
        Object rt = null;
        // POST
        try {
            String entity = "{ \"docs\": [{\"text_field\": \""+text+"\"}]}";
            Request req = new Request("POST","_ml/trained_models/"+modelId+"/_infer");
            HttpEntity params = new NStringEntity(entity, ContentType.APPLICATION_JSON);
            req.setEntity(params);
            Response rsp = lowLevelClient.performRequest(req);
            HttpEntity en = rsp.getEntity();
            String body = EntityUtils.toString(en);
            JSONObject data = JSON.parseObject(body);
            rt = data.getJSONArray("inference_results").getJSONObject(0).getJSONArray("predicted_value");
        }
        catch (Exception e){
            e.printStackTrace();
        }
        return rt;
    }
 
    // 提交 Elasticsearch
    public static void uploadDoc(List<String> docSplits) throws Exception{
 
        // 遍历每个文档
        for (int i = 0; i < docSplits.size(); i++) {
            // 原始文本
            String content = docSplits.get(i);
            // 转为向量
            Object content_embed = embedDoc(content);
            // 时间
            Long timestamp = System.currentTimeMillis();
 
            // 上传
            JSONObject data = new JSONObject();
            data.put("content",content);
            data.put("content_embed",content_embed);
            data.put("timestamp",timestamp);
 
            Request req = new Request("POST","/"+indexName+"/_doc");
            HttpEntity params = new NStringEntity(data.toJSONString(), ContentType.APPLICATION_JSON);
            req.setEntity(params);
            Response res = lowLevelClient.performRequest(req);
            System.out.println("上传第"+i+"条："+res);
        }
 
    }
 
    // 执行混合检索
    public static List<String> search(String q) throws Exception{
 
        // 转为向量
        Object vector = embedDoc(q);
 
        // 查询语句
        String query =
                "{\n" +
                        "  \"query\": {\n" +
                        "    \"bool\": {\n" +
                        "      \"must\": [\n" +
                        "        {\n" +
                        "          \"match\": {\n" +
                        "            \"content\": {\n" +
                        "              \"query\": \""+q+"\",\n" +
                        "              \"boost\": 1\n" +
                        "            }\n" +
                        "          }\n" +
                        "        }\n" +
                        "      ]\n" +
                        "    }\n" +
                        "  },\n" +
                        "  \"knn\": {\n" +
                        "    \"field\": \"content_embed\",\n" +
                        "    \"k\": 5,\n" +
                        "    \"num_candidates\": 100,\n" +
                        "    \"query_vector\": "+vector+"\n" +
                        "  },\n" +
                        "  \"size\": 5\n" +
                        "}\n";
 
//        System.out.println("查询语句:");
//        System.out.println(query);
 
        // 调用查询
        Request req = new Request("POST","/"+indexName+"/_search?pretty");
        HttpEntity params = new NStringEntity(query, ContentType.APPLICATION_JSON);
        req.setEntity(params);
        Response res = lowLevelClient.performRequest(req);
 
        // 解析
        String body = EntityUtils.toString(res.getEntity());
        JSONArray data = JSON.parseObject(body).getJSONObject("hits").getJSONArray("hits");
 
        // 遍历每个文档、将高的分的文档保存
        List<String> contents = new ArrayList<>();
        data.stream().map(n->JSONObject.parseObject(n.toString())).forEach(n->{
            // 得分高的才作为资料避免 llm 幻觉
            Double _score = n.getDouble("_score");
            if(_score >= scoreThreshold){
                // 文本
                String content = n.getJSONObject("_source").getString("content");
                contents.add(content);
                System.out.println("召回文档数据:"+n);
            }
        });
 
        System.out.println();
        return contents;
    }
 
    // 拼接 prompt
    public static String prompt(List<String> content,String q){
 
        StringBuilder question = new StringBuilder();
 
        question.append("你好，下面是我搜索得到的资料：\n");
        if(content.size()==0){
            question.append("无。\n");
        }
        for (int i = 0; i < content.size() ; i++) {
            question.append("（"+(i+1)+"）").append(content.get(i)).append("\n");
        }
        question.append("\n");
        question.append("请帮我根据上面的资料分析下面的问题，并帮我根据资料列出相关依据：\n");
        question.append(q).append("\n");
        question.append("\n");
        question.append("如果根据资料无法分析请回复不知道!");
 
        return question.toString();
    }
 
    // 调用 LLM 生成回答
    public static String llmAnswer(String question) throws Exception{
 
        JSONObject data = new JSONObject();
        data.put("prompt",question);
        data.put("history",null);
 
        HttpPost httpPost = new HttpPost(llm_host);
        httpPost.addHeader("Content-Type", "application/json;charset=utf-8");
        httpPost.setEntity(new StringEntity(data.toString(), StandardCharsets.UTF_8));
        CloseableHttpResponse response = HttpClients.createDefault().execute(httpPost);
        HttpEntity resEntity = response.getEntity();
        String resp = EntityUtils.toString(resEntity,"utf-8");
        return JSONObject.parseObject(resp).getString("response");
    }
 
    public static void main(String[] args) throws Exception{
 
        // 创建索引
//        createIndex();
 
        // 读取并拆分文档、提交 Elasticsearch
//        uploadDoc(parseDoc(500,100));
 
        // 执行混合检索
        String question = "结婚领证登记需要双发到场吗？";
        List<String> contents = search(question);
 
        // 执行混合检索并拼接 prompt
        String prompt = prompt(contents,question);
 
        // 调用 LLM 生成回答
        String answer = llmAnswer(prompt);
 
        System.out.println("-----------");
        System.out.println("Question:");
        System.out.println(question);
        System.out.println("-----------");
        System.out.println("Prompt:");
        System.out.println(prompt);
        System.out.println("-----------");
        System.out.println("Answer:");
        System.out.println(answer);
    }
 
}