UMAP可以对生成的 embedding 降维，而 matplotlib 可以对图表进行绘制，因此可以结合着两个库去对 RAG 的效果进行可视化。在这里，为了方便，我们使用 chromadb 作为向量数据库去存储向量，并展示。

UMAP 简介

UMAP (Uniform Manifold Approximation and Projection) 是一个用于降维的非线性技术，特别适用于可视化高维数据。UMAP 通过保持数据的全局和局部结构来尝试保持数据的流形结构，并可以将这些结构以二维或三维形式可视化。UMAP 在数据科学和机器学习中非常流行，特别是在需要探索性数据分析和数据可视化时。

UMAP 的 Python 实现提供了以下几个关键特性：

降维：将高维数据集映射到较低维度（例如 2D 或 3D）。

可视化：帮助开发者以图形方式理解高维数据。

近邻搜索：UMAP 在计算过程中使用了近邻搜索算法来保留数据的局部结构。

效率：UMAP 设计为在大型数据集上高效运行。

可解释性：与一些其他降维技术相比，UMAP 更易于理解，并且它的内部机制相对透明。

Embedding 模型定义

我们还是使用 BGE 的向量模型进行实验，可以从 HF 将模型下载到本地，然后调用就十分方便了。

class BaaiEmbedding(Embeddings):

    def __init__(self, max_length=512, batch_size=256):
        self.model = FlagModel(BGE_MODEL_PATH, query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：")
        self.max_length = max_length
        self.batch_size = batch_size

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode_corpus(texts, self.batch_size, self.max_length).tolist()

    def embed_query(self, text: str) -> List[float]:
        return self.model.encode_queries(text, self.batch_size, self.max_length).tolist()

使用 UMAP 压缩 embedding

可以使用 umap 将向量压缩成二维的，方便可视化。

umap_transform = umap.UMAP(random_state=0, transform_seed=0).fit(embeddings)
def project_embeddings(embeddings, umap_transform):
    umap_embeddings = np.empty((len(embeddings),2))
    for i, embedding in enumerate(tqdm(embeddings)):
        umap_embeddings[i] = umap_transform.transform([embedding])
    return umap_embeddings

总代码

import umap
import numpy as np
import matplotlib.pyplot as plt
from typing import List
from tqdm import tqdm
from FlagEmbedding import FlagModel
from langchain.schema.embeddings import Embeddings
from algo_common.text2vec.feature_dict import BGE_MODEL_PATH
from chroma import ChromaDB

class BaaiEmbedding(Embeddings):

    def __init__(self, max_length=512, batch_size=256):
        self.model = FlagModel(BGE_MODEL_PATH, query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：")
        self.max_length = max_length
        self.batch_size = batch_size

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode_corpus(texts, self.batch_size, self.max_length).tolist()

    def embed_query(self, text: str) -> List[float]:
        return self.model.encode_queries(text, self.batch_size, self.max_length).tolist()

db = ChromaDB("./zsxq_index")
print(db.collection.count())
embed_model = BaaiEmbedding()

def project_embeddings(embeddings, umap_transform):
    umap_embeddings = np.empty((len(embeddings),2))
    for i, embedding in enumerate(tqdm(embeddings)):
        umap_embeddings[i] = umap_transform.transform([embedding])
    return umap_embeddings

embeddings = db.collection.get(include=['embeddings'])['embeddings']
print(len(embeddings))
umap_transform = umap.UMAP(random_state=0, transform_seed=0).fit(embeddings)
# 可视化 query and doc
query = "测试用例 1？"
query_embedding = embed_model.embed_query(query)
results = db.collection.query(query_embeddings=query_embedding, n_results=5, include=['documents', 'embeddings'])


retrieved_embeddings = results['embeddings'][0]
projected_dataset_embeddings = project_embeddings(embeddings, umap_transform)
projected_query_embedding = project_embeddings([query_embedding], umap_transform)
projected_retrieved_embeddings = project_embeddings(retrieved_embeddings, umap_transform)

# Plot the projected query and retrieved documents in the embedding space
plt.figure()
plt.scatter(projected_dataset_embeddings[:, 0], projected_dataset_embeddings[:, 1], s=10, color='gray')
plt.scatter(projected_query_embedding[:, 0], projected_query_embedding[:, 1], s=150, marker='X', color='r')
plt.scatter(projected_retrieved_embeddings[:, 0], projected_retrieved_embeddings[:, 1], s=100, facecolors='none',
                edgecolors='g')

plt.gca().set_aspect('equal', 'datalim')
plt.title('projection of query and retrieved documents')
plt.axis('off')
plt.savefig("./query.png")
plt.show()

数据可视化示例

下面是用示例代码运行的一个实际例子，其中灰色的点表示所有的文档向量，带有绿色框的表示召回的文档，红色 X 表示的是用户输入 query 的向量。

从下图的结果来说，有一些在用户 query 附近的点，实际没有被召回，因此，我们在实际使用的时候，可以自己去进行一定的优化和调整。

好了，我写完了，感谢阅读，有帮助欢迎点赞评论，一键三连。

利用 UMAP 和 matplotlib 可视化 RAG 检索的效果

UMAP 简介

Embedding 模型定义

使用 UMAP 压缩 embedding

总代码

数据可视化示例