0. 源代码文件

https://github.com/langchain-ai/langchain/blob/master/cookbook/RAPTOR.ipynb

1. 部分代码理解笔记

我们可以用一个简单的故事和易于理解的语言来解释每一部分。我们假设我们要处理一堆文本（比如一些文章），然后把它们分类成不同的类别，并总结每个类别的内容。

故事背景

假设我们有很多关于不同主题的文章，比如“动物”、“体育”、“科技”等。我们希望自动地把这些文章分成不同的类别，并总结每个类别的主要内容。

导入工具

首先，我们需要一些工具来帮助我们完成这项任务。就像我们要建一座房子需要锤子和钉子一样，我们在代码里导入了一些工具：

from typing import Dict, List, Optional, Tupleimport numpy as np
import pandas as pd
import umap
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from sklearn.mixture import GaussianMixture

这些工具会帮助我们处理数据、做数学计算、生成文本摘要等等。

固定种子（随机种子）

我们设置了一个固定的种子，这样每次运行代码时，结果都是一样的，便于重复和验证。

RANDOM_SEED = 224

全局降维函数

这个函数叫 global_cluster_embeddings，它的作用是把高维的嵌入向量（代表文本的数学表示）降到一个低维空间，方便后续处理。

def global_cluster_embeddings(embeddings: np.ndarray,dim: int,n_neighbors: Optional[int] = None,metric: str = "cosine",
) -> np.ndarray:"""使用UMAP进行全局降维"""if n_neighbors is None:n_neighbors = int((len(embeddings) - 1) ** 0.5)return umap.UMAP(n_neighbors=n_neighbors, n_components=dim, metric=metric).fit_transform(embeddings)

局部降维函数

这个函数叫 local_cluster_embeddings，它和上面的函数类似，但通常在全局降维之后使用，以进一步细化分类。

def local_cluster_embeddings(embeddings: np.ndarray, dim: int, num_neighbors: int = 10, metric: str = "cosine"
) -> np.ndarray:"""使用UMAP进行局部降维"""return umap.UMAP(n_neighbors=num_neighbors, n_components=dim, metric=metric).fit_transform(embeddings)

获取最佳聚类数函数

这个函数叫 get_optimal_clusters，它用来找出最适合的聚类数量（把数据分成几组）。

def get_optimal_clusters(embeddings: np.ndarray, max_clusters: int = 50, random_state: int = RANDOM_SEED
) -> int:"""使用高斯混合模型和贝叶斯信息准则来确定最佳聚类数"""max_clusters = min(max_clusters, len(embeddings))n_clusters = np.arange(1, max_clusters)bics = []for n in n_clusters:gm = GaussianMixture(n_components=n, random_state=random_state)gm.fit(embeddings)bics.append(gm.bic(embeddings))return n_clusters[np.argmin(bics)]

GMM聚类函数

这个函数叫 GMM_cluster，它用高斯混合模型（GMM）对数据进行聚类。

def GMM_cluster(embeddings: np.ndarray, threshold: float, random_state: int = 0):"""使用GMM进行聚类，并基于概率阈值进行分配"""n_clusters = get_optimal_clusters(embeddings)gm = GaussianMixture(n_components=n_clusters, random_state=random_state)gm.fit(embeddings)probs = gm.predict_proba(embeddings)labels = [np.where(prob > threshold)[0] for prob in probs]return labels, n_clusters

执行聚类函数

这个函数叫 perform_clustering，它结合了全局降维、GMM聚类和局部降维聚类的步骤，把所有的文本分成不同的类别。

def perform_clustering(embeddings: np.ndarray,dim: int,threshold: float,
) -> List[np.ndarray]:"""对嵌入向量进行全局降维、GMM聚类和局部降维聚类"""if len(embeddings) <= dim + 1:return [np.array([0]) for _ in range(len(embeddings))]reduced_embeddings_global = global_cluster_embeddings(embeddings, dim)global_clusters, n_global_clusters = GMM_cluster(reduced_embeddings_global, threshold)all_local_clusters = [np.array([]) for _ in range(len(embeddings))]total_clusters = 0for i in range(n_global_clusters):global_cluster_embeddings_ = embeddings[np.array([i in gc for gc in global_clusters])]if len(global_cluster_embeddings_) == 0:continueif len(global_cluster_embeddings_) <= dim + 1:local_clusters = [np.array([0]) for _ in global_cluster_embeddings_]n_local_clusters = 1else:reduced_embeddings_local = local_cluster_embeddings(global_cluster_embeddings_, dim)local_clusters, n_local_clusters = GMM_cluster(reduced_embeddings_local, threshold)for j in range(n_local_clusters):local_cluster_embeddings_ = global_cluster_embeddings_[np.array([j in lc for lc in local_clusters])]indices = np.where((embeddings == local_cluster_embeddings_[:, None]).all(-1))[1]for idx in indices:all_local_clusters[idx] = np.append(all_local_clusters[idx], j + total_clusters)total_clusters += n_local_clustersreturn all_local_clusters

嵌入函数

这个函数叫 embed，它把文本转成嵌入向量。

def embed(texts):"""为文本生成嵌入向量"""text_embeddings = embd.embed_documents(texts)text_embeddings_np = np.array(text_embeddings)return text_embeddings_np

嵌入并聚类文本函数

这个函数叫 embed_cluster_texts，它先嵌入文本，再进行聚类，并把结果存到一个数据框里。

def embed_cluster_texts(texts):"""嵌入文本并聚类，返回包含文本、嵌入向量和聚类标签的数据框"""text_embeddings_np = embed(texts)cluster_labels = perform_clustering(text_embeddings_np, 10, 0.1)df = pd.DataFrame()df["text"] = textsdf["embd"] = list(text_embeddings_np)df["cluster"] = cluster_labelsreturn df

格式化文本函数

这个函数叫 fmt_txt，它把数据框里的文本格式化成一个字符串。

def fmt_txt(df: pd.DataFrame) -> str:"""格式化文本"""unique_txt = df["text"].tolist()return "--- --- \n --- --- ".join(unique_txt)

嵌入、聚类并总结文本函数

这个函数叫 embed_cluster_summarize_texts，它嵌入、聚类并总结文本，返回两个数据框，一个包含文本及其聚类标签，另一个包含每个聚类的总结。

def embed_cluster_summarize_texts(texts: List[str], level: int
) -> Tuple[pd.DataFrame, pd.DataFrame]:"""嵌入、聚类并总结文本"""df_clusters = embed_cluster_texts(texts)expanded_list = []for index, row in df_clusters.iterrows():for cluster in row["cluster"]:expanded_list.append({"text": row["text"], "embd": row["embd"], "cluster": cluster})expanded_df = pd.DataFrame(expanded_list)all_clusters = expanded_df["cluster"].unique()template = """Here is a sub-set of LangChain Expression Language doc. LangChain Expression Language provides a way to compose chain in LangChain.Give a detailed summary of the documentation provided.Documentation:{context}"""prompt = ChatPromptTemplate.from_template(template)chain = prompt | model | StrOutputParser()summaries = []for i in all_clusters:df_cluster = expanded_df[expanded_df["cluster"] == i]formatted_txt = fmt_txt(df_cluster)summaries.append(chain.invoke({"context": formatted_txt}))df_summary = pd.DataFrame({"summaries": summaries,"level": [level] * len(summaries),"cluster": list(all_clusters),})return df_clusters, df_summary

递归嵌入、聚类并总结函数

这个函数叫 `recursive_embed_cluster_summar

ize`，它递归地进行嵌入、聚类并总结，直到达到指定的层级或聚类数变为1。

def recursive_embed_cluster_summarize(texts: List[str], level: int = 1, n_levels: int = 3
) -> Dict[int, Tuple[pd.DataFrame, pd.DataFrame]]:"""递归嵌入、聚类并总结文本"""results = {}df_clusters, df_summary = embed_cluster_summarize_texts(texts, level)results[level] = (df_clusters, df_summary)unique_clusters = df_summary["cluster"].nunique()if level < n_levels and unique_clusters > 1:new_texts = df_summary["summaries"].tolist()next_level_results = recursive_embed_cluster_summarize(new_texts, level + 1, n_levels)results.update(next_level_results)return results

总结

这段代码的总体思路是：

1. 把文本转成嵌入向量。
2. 对嵌入向量进行全局降维。
3. 使用GMM模型对降维后的嵌入向量进行聚类。
4. 对每个全局聚类的结果进行进一步的局部降维和聚类。
5. 对每个聚类的文本进行总结。
6. 如果还有更多层级，递归地对每个总结进行嵌入、聚类和总结，直到达到指定的层数或聚类数变为1。