类功能
功能描述
通过聚类模型将嵌入后的文本进行聚类,将其划分为多个语义簇。随后,计算context的切片与question(总结长文本的指令)的余弦相似度。根据设定的压缩率,在每个簇内删除相似度较低的切片,从而保留与指令最相关的信息,实现长文本的压缩式总结。
函数原型
from mx_rag.compress.cluster_compressor import ClusterCompressor class ClusterCompressor(cluster_func, embed, splitter, dev_id):
输入参数说明
参数名 |
数据类型 |
可选/必选 |
说明 |
|---|---|---|---|
cluster_func |
Callable[[List[List[float]]], List[int]] |
必选 |
聚类函数,将嵌入后的文本进行聚类,将其划分为多个语义簇。 |
embed |
Embeddings |
必选 |
嵌入对象,把文本切片转换为向量,只能为继承自langchain_core.embeddings的Embeddings的子类。 |
splitter |
TextSplitter |
可选 |
文档切分对象,只能为继承自langchain_text_splitters.base的TextSplitter的子类。默认为langchain.text_splitter的RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0, separators=["。", "!", "?", "\n", ",", ";", " ", ""]) |
dev_id |
int |
可选 |
NPU id,通过npu-smi info查询可用ID,取值范围[0, 63],默认为卡0。 |
调用示例
from langchain.text_splitter import RecursiveCharacterTextSplitter
from sklearn.cluster import HDBSCAN
from mx_rag.compress.cluster_compressor import ClusterCompressor
from mx_rag.embedding.local import TextEmbedding
from mx_rag.embedding.service import TEIEmbedding
from mx_rag.utils import ClientParam
context="""需要压缩的prompt文本"""
question="请给上述内容起一个标题"
tei_emb=False
if tei_emb:
emb = TEIEmbedding.create(url="https://ip:port/embed", client_param=ClientParam(ca_file="/path/to/ca.crt"))
else:
emb = TextEmbedding(model_path="embedding_path", dev_id=0)
def _get_community(sentences_embedding):
# 社区划分
node_num=len(sentences_embedding)
min_cluster_size=2
hdbscan = HDBSCAN(min_cluster_size=min(min_cluster_size, node_num))
labels = hdbscan.fit_predict(sentences_embedding)
return labels
splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0, separators=["。", "!", "?", "\n", ",", ";", " ", ""], )
compressor=ClusterCompressor(cluster_func=_get_community, embed=emb, splitter=splitter, dev_id=0)
res=compressor.compress_texts(context, question, 0.6)
print(res)
父主题: ClusterCompressor类