使用嵌入模型进行重排序

RAG的重排序的Python代码

目录

Reranking 是检索增强生成 (RAG) 系统 的第二步， 位于检索和生成之间。

数字空间中的电立方体

以上是 Flux-1 dev 对 数字空间中的电立方体 的想象。

带重排序的检索

如果我们从一开始就将文档以嵌入形式存储在向量数据库中，检索将直接给我们相似文档的列表。

独立重排序

但如果我们先从互联网下载文档，搜索系统的响应可能会受到搜索提供商的偏好/算法、赞助内容、SEO优化等因素的影响，因此我们需要进行搜索后的重排序。

我所做的事情：

  • 获取搜索查询的嵌入
  • 获取每个文档的嵌入。文档本身预计不会超过8k个token
  • 计算查询与每个文档嵌入之间的相似度
  • 按照相似度对文档进行排序。

这里没有向量数据库，我们继续。

示例代码

使用Langchain连接到Ollama和langchain的cosine_similarity函数。 你可以通过相似度进行过滤，但请注意，对于不同领域和嵌入式LLM，阈值会有所不同。

如果这段代码对你有任何帮助，我将非常高兴。 Copy/Paste/UseAnyWayYouWant许可证。 谢谢。

from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.utils.math import cosine_similarity
import numpy as np


def cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return 1.0 - cosine_similarity(a, b)

def compute_score(vectors: np.ndarray) -> float:
    score = cosine_distance(vectors[0].reshape(1, -1), vectors[1].reshape(1, -1)).item()
    return score

def list_to_array(lst):
    return np.array(lst, dtype=float)   

def compute_scorel(lists) -> float:
    v1 = list_to_array(lists[0])
    v2 = list_to_array(lists[1])
    return compute_score([v1, v2])

def filter_docs(emb_model_name, docs, query, num_docs):
    content_arr = [doc.page_content for doc in docs]

    ollama_emb = OllamaEmbeddings(
        model=emb_model_name
    )

    docs_embs = ollama_emb.embed_documents(content_arr)
    query_embs = ollama_emb.embed_query(query)
    sims = []
    for i, emb in enumerate(docs_embs):
        idx = docs[i].id
        s = compute_scorel([query_embs, docs_embs[i]])
        simstr = str(round(s, 4))
        docs[i].metadata["sim"] = simstr
        sim = {
            "idx": idx,
            "i": i,
            "sim": s,
        }
        sims.append(sim)

    sims.sort(key=sortFn)

    sorted_docs = [docs[x["i"]] for x in sims]
    filtered_docs = sorted_docs[:num_docs]
    return filtered_docs

最佳嵌入模型

对于我的任务，目前最佳的嵌入模型是 bge-large:335m-en-v1.5-fp16

第二名是 nomic-embed-text:137m-v1.5-fp16 和 jina/jina-embeddings-v2-base-en:latest。

但请根据你自己的领域和查询进行自己的测试。

有用的链接