使用UMAP降维可视化RAG嵌入

大型语言模型（LLMs）如 GPT-4 已经展示了出色的文本理解和生成能力。但它们在处理领域特定信息方面面临挑战，比如当查询超出训练数据范围时，它们会产生错误的答案。LLMs 的推理过程也缺乏透明度，使用户难以理解达成结论的方式。

检索增强生成（RAG）在 LLMS 的工作流程中添加了一个检索步骤，使其能够在响应查询时从其他来源（如私人文本文档）中查询相关数据。这些文档事先分成小段，然后使用embedding的 ML 模型生成嵌入。具有相似内容的段将具有相似的嵌入。当 RAG 应用程序收到一个问题时，它使用查询检索相关文档片段。然后 LLMS 使用这些文档片段作为上下文来回答问题。这种方法可以提供回答查询所需的信息，并通过展示使用的片段来增加回答的透明度。

对于RAG来说，可视化嵌入空间是一个非常重要的方法，因为RAG应用程序使用该空间来查找相关信息。查询的结果与文档片空间息息相关，所以可以使用像UMAP这样的可视化方法，将高维嵌入减少到更易于展示的2D进行可视化。虽然高维嵌入被简化为两个分量，但问题及其相关文档片段在嵌入空间中形成簇，仍然是可以被识别出来，尤其是这时肉眼可见的，所以这有助于深入了解数据的本质。

在本文中，我们将使用HTML格式的Wikipedia中的f1数据集，使用嵌入模型将它们转换为紧凑的矢量表示，并存储到ChromaDB中。使用LangChain构建RAG应用，并在2D中可视化嵌入，分析查询和文档片段之间的关系和接近度。

首先我们安装需要的库

 !pip install langchain langchain-openai chromadb renumics-spotlight

我们使用了Renumics-Spotlight python包，Renumics-Spotlight是一个交互式探索非结构化ML数据集的可视化工具。

而获得嵌入我们则直接使用了OpenAI的embedding-ada-002，所以这里要设置API的Key

 %env OPENAI_API_KEY=<your-api-key>

准备文件

对于数据集是使用wikipedia-api和BeautifulSoup创建的。

本数据集基于维基百科的文章，并在知识共享署名-相同方式共享许可下获得许可。原始文章和作者名单可以在各自的维基百科页面上找到。

将提取的html放到docs/子文件夹中，或者可以通过创建docs/子文件夹的方式将你自己的文件复制到其中来使用自己的Dataset。

创建嵌入

要创建嵌入，首先需要设置嵌入模型和vectorstore。这里我们使用OpenAIEmbeddings中的text- embeddings -ada-002和使用ChromaDB的vectorstore:

 from langchain_openai import OpenAIEmbeddings
 from langchain.vectorstores.chroma import Chroma
 
 embeddings_model = OpenAIEmbeddings(model="text-embedding-ada-002")
 docs_vectorstore = Chroma(
     collection_name="docs_store",
     embedding_function=embeddings_model,
     persist_directory="docs-db",
 )

vectorstore将持久化到docs-db 文件夹中。然后使用BSHTMLLoader加载html文档:

 from langchain_community.document_loaders import BSHTMLLoader, DirectoryLoader
 loader = DirectoryLoader(
     "docs",
     glob="*.html",
     loader_cls=BSHTMLLoader,
     loader_kwargs={"open_encoding": "utf-8"},
     recursive=True,
     show_progress=True,
 )
 docs = loader.load()

将文档分成小块

 from langchain.text_splitter import RecursiveCharacterTextSplitter
 text_splitter = RecursiveCharacterTextSplitter(
     chunk_size=1000, chunk_overlap=200, add_start_index=True
 )
 splits = text_splitter.split_documents(docs)

从元数据重建的id，这样可以在数据库中查找嵌入。

 import hashlib
 import json
 from langchain_core.documents import Document
 
 def stable_hash(doc: Document) -> str:
     """
     Stable hash document based on its metadata.
     """
     return hashlib.sha1(json.dumps(doc.metadata, sort_keys=True).encode()).hexdigest()
 
 split_ids = list(map(stable_hash, splits))
 docs_vectorstore.add_documents(splits, ids=split_ids)
 docs_vectorstore.persist()

LangChain

首先，你需要选择一个LLM模型。我们使用GPT-4

 from langchain_openai import ChatOpenAI
 
 llm = ChatOpenAI(model="gpt-4", temperature=0.0)
 retriever = docs_vectorstore.as_retriever(search_kwargs={"k": 20})

在初始化ChatOpenAI模型时，将temperature参数设置为0.0可确保确定性输出。

我们使用下面的提示：

 from langchain_core.prompts import ChatPromptTemplate
 
 template = """
 You are an assistant for question-answering tasks.
 Given the following extracted parts of a long document and a question, create a final answer with references ("SOURCES").
 If you don't know the answer, just say that you don't know. Don't try to make up an answer.
 ALWAYS return a "SOURCES" part in your answer.
 
 QUESTION: {question}
 =========
 {source_documents}
 =========
 FINAL ANSWER: """
 prompt = ChatPromptTemplate.from_template(template)

然后对检索到的文档进行格式化，使其包含页面内容和源文件路径，将这个格式化的输入输入到语言模型(LLM)，模型就可以根据合并的用户问题和文档上下文生成答案。

 from typing import List
 
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_core.output_parsers import StrOutputParser
 
 
 def format_docs(docs: List[Document]) -> str:
     return "\n\n".join(
         f"Content: {doc.page_content}\nSource: {doc.metadata['source']}" for doc in docs
     )
 
 
 rag_chain_from_docs = (
     RunnablePassthrough.assign(
         source_documents=(lambda x: format_docs(x["source_documents"]))
     )
     | prompt
     | llm
     | StrOutputParser()
 )
 rag_chain = RunnableParallel(
     {
         "source_documents": retriever,
         "question": RunnablePassthrough(),
     }
 ).assign(answer=rag_chain_from_docs)

RAG的为问题回答

RAG应用程序现在已经准备好回答问题了:

 question = "Who built the nuerburgring"
 response = rag_chain.invoke(question)
 response["answer"]

回答如下：

 'The Nürburgring was built in the 1920s, with the construction of the track beginning in September 1925. The track was designed by the Eichler Architekturbüro from Ravensburg, led by architect Gustav Eichler. The original Nürburgring was intended to be a showcase for German automotive engineering and racing talent (SOURCES: data/docs/Nürburgring.html).'

这个问题也将在下一步中用于进一步得到确认。

可视化

为了进行可视化，我们使用Pandas DataFrame来组织数据。从从矢量存储中提取文本片段及其嵌入。我们还要标出正确的答案:

 import pandas as pd
 
 response = docs_vectorstore.get(include=["metadatas", "documents", "embeddings"])
 df = pd.DataFrame(
     {
         "id": response["ids"],
         "source": [metadata.get("source") for metadata in response["metadatas"]],
         "page": [metadata.get("page", -1) for metadata in response["metadatas"]],
         "document": response["documents"],
         "embedding": response["embeddings"],
     }
 )
 df["contains_answer"] = df["document"].apply(lambda x: "Eichler" in x)
 df["contains_answer"].to_numpy().nonzero()

问题和相关的答案也被投影到嵌入空间中。它们的处理方式与文本片段相同:

 question_row = pd.DataFrame(
     {
         "id": "question",
         "question": question,
         "embedding": embeddings_model.embed_query(question),
     }
 )
 answer_row = pd.DataFrame(
     {
         "id": "answer",
         "answer": answer,
         "embedding": embeddings_model.embed_query(answer),
     }
 )
 df = pd.concat([question_row, answer_row, df])

首先打印问题与文档片段之间的距离:

 import numpy as np
 question_embedding = embeddings_model.embed_query(question)
 df["dist"] = df.apply(
     lambda row: np.linalg.norm(
         np.array(row["embedding"]) - question_embedding
     ),
     axis=1,
 )

这个距离还可以用于可视化:

 from renumics import spotlight
 spotlight.show(df)

上面的语句会打开一个新的浏览器窗口。左上角的表格部分显示数据集的所有字段。右上方的相似性图中突出显示它们。

可以看到，包含正确答案的单个文档片段和最相关的文件离问题和答案很近。

总结

使用降维技术可以使用户和开发人员访问嵌入空间。在可视化空间中，可以通过浏览相邻的数据点来进行检索增强的检查。降维可视化虽然有助于理解数据，但也可能存在信息损失，因为它将高维数据映射到一个较低维度的空间中。因此，在进行检查时，需要权衡信息丢失和效果提升之间的关系。

完整代码：

https://avoid.overfit.cn/post/31e45d66ef1547e397bbbef2ebcf38c8

作者:Markus Stoll