NLP(一百零九)Embedding中的Late-Chunking(迟分)策略

本文将会介绍Embedding模型中的Late Chunking(迟分)策略,演示多个中文Late Chunking的例子,并搭建相关Gradio服务,最后再展示其在RAG框架中对于大模型回复质量的提升作用。

Late Chunking技术及原理

Late Chunking(迟分)技术是Jina AI公司(<https://jina.ai)于今年8月份介绍的在Embedding模型方面的新技术。

众所周知,在保留上下文信息的同时对长文本进行分块并保证召回的效果,是一项有难度的挑战。而Late Chunking技术利用长上下文嵌入模型生成上下文的分块嵌入,以实现更好的检索应用。

Late Chunking(迟分)是一种先通读全文再分块的新方法,包含两个主要步骤:

  • 编码全文:先编码整个文档,得到每个 token 的向量表示,保留完整的上下文信息。
  • 分块池化:根据分块边界,对同一个文本块的 token 向量进行平均池化,生成每个文本块的向量。由于每个 token 的向量是在全文的语境下生成的,因此迟分可以保留块之间的上下文信息。

其原理如下图所示:

迟分与朴素分块的原理对比

从上面的原理图中,我们可以看到,Late Chunking技术并没有改变Embedding模型的内部结构,而是在对文本进行嵌入时改变了嵌入方式。传统的嵌入方式(Naive Chunking)是先对文本进行切分,分别对每个chunk进行嵌入;而Late Chunking是先对文本进行token级别的嵌入,获取每个token的嵌入,再根据文本块的token向量进行平均池化,生成每个文本块的嵌入向量,这也是它被称为“迟分”的原因。当然,Late Chunking并不是对所有Embedding模型都会生效,目前只有Jina AI的Embedding模型能做到。

Jina AI官网给出了一个英语方面的生动例子,使用的Embedding模型为jina-embeddings-v2-base-en,对于输入的文本,按照句子进行切分,共生成3个chunk,输入的query为Berlin,朴素嵌入(即我们现在在用的常见的嵌入方式)和Late Chunking的相似度分数计算如下:

Query Chunk Sim. on naive chunking Sim. on late chunking
Berlin Berlin is the capital and largest city of Germany, both by area and by population. 0.849 0.850
Berlin Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits. 0.708 0.825
Berlin The city is also one of the states of Germany, and is the third smallest state in the country in terms of area. 0.753 0.850

从上面的英语例子中,我们可以看到Late Chunking很好地保留了上下文之间的信息,每个chunk与query之间的相似度都比较高,而朴素嵌入时其余两个文本中不含Berlin,因此与query的相似度较低。这个例子很好地展示了Late Chunking技术有不错的上下文信息保留能力。

网络上能搜索到的关于Late Chunking技术大概就这么多。笔者想要在这基础上再深入一步,探索Late Chunking的更多应用。

中文Late Chunking的例子

官网给出了Late Chunking在英语方面的例子,这里我们将其扩充至中文。笔者选用的中文Embedding模型为jinaai/jina-embeddings-v2-base-zh,使用的示例文本为(来源于王安石的百度词条):

王安石(1021年12月19日-1086年5月21日),字介甫,号半山。抚州临川县(今属江西省抚州市)人。中国北宋时期政治家、文学家、思想家、改革家。庆历二年(1042年),王安石中进士,历任扬州签判、鄞县知县、舒州通判等职,政绩显著。宋仁宗末年,曾作《上仁宗皇帝言事书》,要求对宋初以来的法度进行全盘改革,但未被采纳。

  1. 加载模型:
1
2
3
4
5
6
from transformers import AutoModel
from transformers import AutoTokenizer

# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)
  1. 按句子进行切分
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def chunk_by_sentences(input_text: str, tokenizer: callable):
"""
Split the input text into sentences using the tokenizer
:param input_text: The text snippet to split into sentences
:param tokenizer: The tokenizer to use
:return: A tuple containing the list of text chunks and their corresponding token spans
"""
inputs = tokenizer(input_text, return_tensors='pt', return_offsets_mapping=True)
punctuation_mark_id = tokenizer.convert_tokens_to_ids('。')
sep_id = tokenizer.eos_token_id
token_offsets = inputs['offset_mapping'][0]
token_ids = inputs['input_ids'][0]
chunk_positions = [
(i, int(start + 1))
for i, (token_id, (start, end)) in enumerate(zip(token_ids, token_offsets))
if token_id == punctuation_mark_id
and (
token_offsets[i + 1][0] - token_offsets[i][1] >= 0
or token_ids[i + 1] == sep_id
)
]
chunks = [
input_text[x[1] : y[1]]
for x, y in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
]
span_annotations = [
(x[0], y[0]) for (x, y) in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
]
return chunks, span_annotations
  1. 对示例文本进行切分
1
2
3
4
5
input_text = "王安石(1021年12月19日-1086年5月21日),字介甫,号半山。抚州临川县(今属江西省抚州市)人。中国北宋时期政治家、文学家、思想家、改革家。庆历二年(1042年),王安石中进士,历任扬州签判、鄞县知县、舒州通判等职,政绩显著。宋仁宗末年,曾作《上仁宗皇帝言事书》,要求对宋初以来的法度进行全盘改革,但未被采纳。"

# determine chunks
chunks, span_annotations = chunk_by_sentences(input_text, tokenizer)
print('Chunks:\n- "' + '"\n- "'.join(chunks) + '"')

输出结果如下:

Chunks:
- "王安石(1021年12月19日-1086年5月21日),字介甫,号半山。"
- "抚州临川县(今属江西省抚州市)人。"
- "中国北宋时期政治家、文学家、思想家、改革家。"
- "庆历二年(1042年),王安石中进士,历任扬州签判、鄞县知县、舒州通判等职,政绩显著。"
- "宋仁宗末年,曾作《上仁宗皇帝言事书》,要求对宋初以来的法度进行全盘改革,但未被采纳。"
  1. 定义late_chunking函数
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def late_chunking(
model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
token_embeddings = model_output[0]
outputs = []
for embeddings, annotations in zip(token_embeddings, span_annotation):
if (
max_length is not None
): # remove annotations which go bejond the max-length of the model
annotations = [
(start, min(end, max_length - 1))
for (start, end) in annotations
if start < (max_length - 1)
]
pooled_embeddings = [
embeddings[start:end].sum(dim=0) / (end - start)
for start, end in annotations
if (end - start) >= 1
]
pooled_embeddings = [
embedding.detach().cpu().numpy() for embedding in pooled_embeddings
]
outputs.append(pooled_embeddings)

return outputs
  1. 对比朴素嵌入与Late Chunking的结果
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# chunk before
embeddings_traditional_chunking = model.encode(chunks)

# chunk afterwards (context-sensitive chunked pooling)
inputs = tokenizer(input_text, return_tensors='pt')
model_output = model(**inputs)
embeddings = late_chunking(model_output, [span_annotations])[0]

import numpy as np

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))

query = "王安石是哪个朝代的"
# query = "王安石是哪里人"
query_embedding = model.encode(query)

for chunk, new_embedding, trad_embeddings in zip(chunks, embeddings, embeddings_traditional_chunking):
print(f'similarity_new("{query}", "{chunk}"):', cos_sim(query_embedding, new_embedding))
print(f'similarity_trad("{query}", "{chunk}"):', cos_sim(query_embedding, trad_embeddings))

输出结果如下:

similarity_new("王安石是哪个朝代的", "王安石(1021年12月19日-1086年5月21日),字介甫,号半山。"): 0.6774667
similarity_trad("王安石是哪个朝代的", "王安石(1021年12月19日-1086年5月21日),字介甫,号半山。"): 0.7342801
similarity_new("王安石是哪个朝代的", "抚州临川县(今属江西省抚州市)人。"): 0.61272216
similarity_trad("王安石是哪个朝代的", "抚州临川县(今属江西省抚州市)人。"): 0.27474773
similarity_new("王安石是哪个朝代的", "中国北宋时期政治家、文学家、思想家、改革家。"): 0.63981277
similarity_trad("王安石是哪个朝代的", "中国北宋时期政治家、文学家、思想家、改革家。"): 0.49549717
similarity_new("王安石是哪个朝代的", "庆历二年(1042年),王安石中进士,历任扬州签判、鄞县知县、舒州通判等职,政绩显著。"): 0.61709845
similarity_trad("王安石是哪个朝代的", "庆历二年(1042年),王安石中进士,历任扬州签判、鄞县知县、舒州通判等职,政绩显著。"): 0.57014936
similarity_new("王安石是哪个朝代的", "宋仁宗末年,曾作《上仁宗皇帝言事书》,要求对宋初以来的法度进行全盘改革,但未被采纳。"): 0.5486519
similarity_trad("王安石是哪个朝代的", "宋仁宗末年,曾作《上仁宗皇帝言事书》,要求对宋初以来的法度进行全盘改革,但未被采纳。"): 0.36279958

根据上面的对比结果,我们输入的query为王安石是哪个朝代的,在朴素嵌入结果中,正确答案对应文本排在第三位,相似度分数为0.4955,而在Late Chunking的结果中,正确答案对应文本排在第二位,相似度分数为0.6398。

由此可见,Late Chunking比朴素嵌入更能保留上下文信息,尤其是上下文之间存在指代关系的文本,Late Chunking的表现更为出色。

使用Gradio实现中文Late Chunking服务

上面的例子仅仅是中文Late Chunking的一个简单例子,让我们来使用Gradio工具,实现中文Late Chunking服务,将query的召回结果按照文本相似度排序,获得更为直观的展示。

搭建Gradio服务的Python代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import gradio as gr
import numpy as np
from transformers import AutoModel, AutoTokenizer

# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)


def chunk_by_sentences(input_text: str, tokenizer: callable, separator: str):
inputs = tokenizer(input_text, return_tensors='pt', return_offsets_mapping=True)
punctuation_mark_id = tokenizer.convert_tokens_to_ids(separator)
print(f"separator: {separator}, punctuation_mark_id: {punctuation_mark_id}")
sep_id = tokenizer.eos_token_id
token_offsets = inputs['offset_mapping'][0]
token_ids = inputs['input_ids'][0]
chunk_positions = [
(i, int(start + 1))
for i, (token_id, (start, end)) in enumerate(zip(token_ids, token_offsets))
if token_id == punctuation_mark_id
and (
token_offsets[i + 1][0] - token_offsets[i][1] >= 0
or token_ids[i + 1] == sep_id
)
]
chunks = [
input_text[x[1]: y[1]]
for x, y in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
]
span_annotations = [
(x[0], y[0]) for (x, y) in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
]
return chunks, span_annotations


def late_chunking(model_output, span_annotation, max_length=None):
token_embeddings = model_output[0]
outputs = []
for embeddings, annotations in zip(token_embeddings, span_annotation):
if max_length is not None:
annotations = [
(start, min(end, max_length - 1))
for (start, end) in annotations
if start < (max_length - 1)
]
pooled_embeddings = [
embeddings[start:end].sum(dim=0) / (end - start)
for start, end in annotations
if (end - start) >= 1
]
pooled_embeddings = [
embedding.detach().cpu().numpy() for embedding in pooled_embeddings
]
outputs.append(pooled_embeddings)

return outputs


def embedding_retriever(query_input, text_input, separator):
chunks, span_annotations = chunk_by_sentences(text_input, tokenizer, separator)
print(f"chunks: ", chunks)
inputs = tokenizer(text_input, return_tensors='pt', max_length=4096, truncation=True)
model_output = model(**inputs)
late_chunking_embeddings = late_chunking(model_output, [span_annotations])[0]

query_inputs = tokenizer(query_input, return_tensors='pt')
query_embedding = model(**query_inputs)[0].detach().cpu().numpy().mean(axis=1)

traditional_chunking_embeddings = model.encode(chunks)

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))

naive_embedding_score_dict = {}
late_chunking_embedding_score_dict = {}
for chunk, trad_embed, new_embed in zip(chunks, traditional_chunking_embeddings, late_chunking_embeddings):
# 计算query和每个chunk的embedding的cosine similarity,相似度分数转化为float类型
naive_embedding_score_dict[chunk] = round(cos_sim(query_embedding, trad_embed).tolist()[0], 4)
late_chunking_embedding_score_dict[chunk] = round(cos_sim(query_embedding, new_embed).tolist()[0], 4)

naive_embedding_order = sorted(
naive_embedding_score_dict.items(), key=lambda x: x[1], reverse=True
)
late_chunking_order = sorted(
late_chunking_embedding_score_dict.items(), key=lambda x: x[1], reverse=True
)

df_data = []
for i in range(len(naive_embedding_order)):
df_data.append([i+1, naive_embedding_order[i][0], naive_embedding_order[i][1],
late_chunking_order[i][0], late_chunking_order[i][1]])
return df_data


if __name__ == '__main__':
with gr.Blocks() as demo:
query = gr.TextArea(lines=1, placeholder="your query", label="Query")
text = gr.TextArea(lines=3, placeholder="your text", label="Text")
sep = gr.TextArea(lines=1, placeholder="your separator", label="Separator")
submit = gr.Button("Submit")
result = gr.DataFrame(headers=["order", "naive_embedding_text", "naive_embedding_score",
"late_chunking_text", "late_chunking_score"],
label="Retrieve Result",
wrap=True)

submit.click(fn=embedding_retriever,
inputs=[query, text, sep],
outputs=[result])
demo.launch()

下面笔者将借助这个Gradio服务,来展示几个Late Chunking与朴素嵌入的对比结果的例子。其中示例文本分别来自王安石清明上河图密码百度词条。

late chunking例子1
late chunking例子2
late chunking例子3

上面的几个例子中,Late Chunking的召回结果都比朴素嵌入的要好,这是因为这些文本块之间存在着明显的指代关系,而Late Chunking此时能很好地保留这些文本块之间的上下文信息。

上述的Gradio服务,笔者后续将会放在HuggingFace Spaces中进行部署,有兴趣的读者到时可以试用。

Late Chunking在RAG框架中的作用

下面笔者将会来介绍Late Chunking在RAG框架中的作用,看看Late Chunking是如何在RAG过程中提升召回效果,保证回复质量的。

我们的示例文本是关于蔚来ET9的,其文章标题为蔚来ET9正式上市 售78.8万元起,网址为 https://news.qq.com/rain/a/20241221A07RW900

我们对上面的Late Chunking中的切分方式做个小小的改动,之前是按照分隔符进行句子级别的切分,这里我们保留原文中的段落切分的方式。对于Embedding的召回结果,我们取top 4合并成参考文本,并使用大模型进行问题回复。

笔者实现了对比Late Chunking与朴素嵌入在RAG过程的回复结果的Python代码,如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# -*- coding: utf-8 -*-
# @place: Pudong, Shanghai
# @file: late_chunking_exp.py
# @time: 2024/12/22 22:48
from transformers import AutoModel
from transformers import AutoTokenizer
import os
import numpy as np
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)

chunks = [
"蔚来ET9正式上市 售78.8万元起",
"易车讯 12月21日,蔚来ET9正式上市,售价区间78.8-81.8万元。蔚来ET9定位蔚来品牌科技行政旗舰轿车,新车搭载众多顶尖的黑科技,是中国首款搭载线控转向技术的量产车型,并搭载先进数字架构。",
"蔚来ET9保留了蔚来家族式设计,标志性的一体式X-Bar和Double Dash日间行车灯,让新车看起来富有力量感。“Design for AD”的设计理念得以延续,前瞭望塔式传感器布局,将3颗激光雷达、摄像头等感应硬件巧妙融入外观造型设计中。",
"车头大灯组采用了行业首发MicroLED智能高像素大灯,结合Aqulia2.0超感系统可以实现“广、亮、准、远”的精细化照明。新车整体造型非常流畅,车顶流线从车头一直延伸到车尾,像一张巨大的弓箭,在保持了经典轿车造型商务感的同时,又带来强大的气场和未来气息。",
"超感系统天鹰座Aquila 2.0新增双侧广角激光雷达,通过两侧金属翼子板集成,即提升了安全性,又提升了辅助驾驶能力。超远距激光雷达,搭载蔚来自研杨戬主控芯片,成像效果更佳清晰。新车首次搭载4D毫米波成像雷达,大大增加前向感知能力。",
"车身尺寸方面,蔚来ET9长宽高分别为5325*2017*1621mm,轴距达到了3250mm。此外,新车还配备了23寸的铝合金锻造轮毂,且搭配同级最大的790mm轮胎直径,极具视觉冲击力。来到车尾,新车延续了家族式设计,贯穿式的尾灯组极具辨识度。值得一提的是,新车搭配了同级唯一的鹅颈式全主动尾翼,运动感十足。蔚来ET9首发感应式电动前备箱,支持脚踢感应和车外语音开启,前备箱容积达到105L。",
"内饰方面,蔚来ET9首次采用了矩形方向盘,同时,新车还首发搭载蓝宝石全焦段 AR HUD,能够实现远焦面15米处等效120寸AR-HUD效果。",
"作为行政旗舰轿车,蔚来ET9采用四座布局,创造性的采用了“天空岛”和“行政桥”的设计,配合拱式车身设计,后排的乘坐体验堪比商务舱。在'行政桥'内部,蔚来为二排乘客精心设计了飞机头等舱座椅,拥有582mm超宽坐垫,拥有前排22向,后排20向电动调节。此外,二排座椅还拥有135°超大躺角,可一键尊享11项功能联动。后排还配备了一张360°无级调节的行政桌案,能在任意角度随心调节。“行政桥”下方集成智能冰箱,最大容积达到10L,温度调节范围在-2°C到55°C,此外还首发了常温模式,总计拥有6种预设模式。",
"此外,全车配备七扇电动遮阳帘,支持一键开启。专为后排商务场景开发的全景互联行政屏,应用14.5英寸OLED高清显示屏,屏幕角度可随座椅位置调节,任意姿态下都能拥有舒适的视角。",
"蔚来ET9还首发九霄天琴蔚来8.2.4.8旗舰沉浸声系统。配备了35个扬声器,采用8.2.4.8声学布局,功率可达2800W。在ET9后排的行政桥内,还设置了中置环绕单元,配备了2个高音扬声器+1个中音扬声器。",
"蔚来ET9还首发搭载cedar 雪松全新智能系统,包含全新一代感知硬件、全新一代中央计算器、SkyOS 天枢整车操作系统等。ET9搭载了蔚来首款5nm车规级智能驾驶芯片——神玑NX9031,与全球首个以车为中心的智能电动汽车整车全域操作系统SkyOS·天枢相结合,实现算力与自身算法的紧密结合,智驾、座舱跨域计算资源的共享,带来极致安全和极致效率。",
"蔚来ET9搭载先进数字架构,定义了一层解耦的计算与通信框架,能够支持智能硬件、操作系统、算法和应用等各层次独立迭代。具体来看,蔚来ET9的先进数字架构由大脑——中央计算平台、小脑与脊髓——高效区域控制器、神经网络——高速冗余的通信网络、血液循环——双冗余低压电源、感知器官——超感系统、灵魂和思想——整车全域操作系统六大部分组成,具备强大的算力、超高带宽与极低时延、极致可靠、精准到点的能源管理等特点。在先进数字架构的支撑下,蔚来ET9实现了多项全球首发与同级领先的技术。",
"SkyOS是蔚来整车底层操作系统,包含了整车系统、智驾系统、智能座舱系统、联通服务补能和移动互联,解决整车各个系统不同域之间的安全性、实时性和应用的复杂性问题,以及将软件定义汽车有效落实到造车的各个环节,建立全方位的、立体的技术体系,使得各种设备能够有机地融合在一起,实现高效的协同工作。",
"蔚来ET9搭载国内首个“全域900V高压架构”,包含电池、电机、线束、空调、DC-DC、车载充电机等核心电子电器元件,拥有最高电压925V、充电峰值功率600kW、充电峰值电流765A的三项全球第一。",
"具体来看,蔚来ET9搭载了前180千瓦感应异步电机,后340千瓦永磁同步电机,综合功率520千瓦,综合扭矩达700牛·米,百公里加速4.3秒。电池方面,蔚来ET9搭载自研46105大圆柱电芯。补能方面,新车的闪电充峰值功率高达600kW,充电峰值电流765A,900V支持充电5分钟补能255公里。",
"蔚来ET9搭载“SkyRide·天行智能底盘系统”,首次将线控转向、后轮转向和全主动悬架三大核心硬件系统集成在一起,是目前全球唯一的全线控智能底盘。全球首创智能化、高集成度的主动悬架系统,每个减振器高度集成独立电动液压泵,无刷直流电机响应迅速,可以在1毫秒内完成信息处理、计算和响应。同时,悬架支持大幅度高度调节,每秒可进行1000次扭矩调整,且四轮独立控制,满足多场景驾驶需求。",
"蔚来ET9首次应用的航空工业级“线控转向”技术,方向盘与转向电机之间采用电讯号传输,不仅结构重量轻,传递效率也能提升40%,并支持OTA迭代升级。在低速泊车、掉头等场景中,“线控转向”技术提供灵敏便捷的转向,无需交叉手打方向盘,配合标配最大后轮转角8.3°的后轮转向系统,实现最小10.9米的转弯直径。",
"天行全主动悬架的每个减振器高度集成独立电动液压泵,无刷直流电机响应迅速,可以在1毫秒内完成信息处理、计算和响应。同时,悬架支持大幅度高度调节,每秒可进行1000次扭矩调整,且四轮独立控制,满足多场景驾驶需求。",
"车身强度方面,新车采用高强度钢铝镁合金车身与空间力学设计,扭转刚度达52600Nm/Deg。车身强度达2000MPa,全面提升乘员舱保护。侧气帘长2.3m,高0.67m,可100%覆盖前后排乘客保护区域。同时,新车搭载了行业首创“V腔”设计的二排专属侧气囊。"
]

input_text = ''.join(chunks)

chunk_inputs = tokenizer(chunks[0], return_tensors='pt')
first_length = chunk_inputs['input_ids'].shape[1]
span_annotations = [(1, first_length)]

for i in range(1, len(chunks)):
chunk_inputs = tokenizer(chunks[i], return_tensors='pt')
length = chunk_inputs['input_ids'].shape[1]
start = span_annotations[i-1][1]
end = start + length
span_annotations.append((start, end))

print(span_annotations)

def late_chunking(
model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
token_embeddings = model_output[0]
outputs = []
for embeddings, annotations in zip(token_embeddings, span_annotation):
if (
max_length is not None
): # remove annotations which go bejond the max-length of the model
annotations = [
(start, min(end, max_length - 1))
for (start, end) in annotations
if start < (max_length - 1)
]
pooled_embeddings = [
embeddings[start:end].sum(dim=0) / (end - start)
for start, end in annotations
if (end - start) >= 1
]
pooled_embeddings = [
embedding.detach().cpu().numpy() for embedding in pooled_embeddings
]
outputs.append(pooled_embeddings)

return outputs

# chunk before
embeddings_traditional_chunking = model.encode(chunks)

# chunk after wards (context-sensitive chunked pooling)
inputs = tokenizer(input_text, return_tensors='pt', max_length=4096, truncation=True)
model_output = model(**inputs)
embeddings = late_chunking(model_output, [span_annotations])[0]

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))

query = "蔚来ET9的车身强度是多少?"
query_embedding = model.encode(query)

naive_embedding_score_dict = {}
late_chunking_embedding_score_dict = {}
for chunk, trad_embed, new_embed in zip(chunks, embeddings_traditional_chunking, embeddings):
# 计算query和每个chunk的embedding的cosine similarity,相似度分数转化为float类型
naive_embedding_score_dict[chunk] = cos_sim(query_embedding, trad_embed)
late_chunking_embedding_score_dict[chunk] = cos_sim(query_embedding, new_embed)

naive_embedding_order = sorted(
naive_embedding_score_dict.items(), key=lambda x: x[1], reverse=True
)
late_chunking_order = sorted(
late_chunking_embedding_score_dict.items(), key=lambda x: x[1], reverse=True
)


def get_answer(query, retrieve_result):
top_k = 4
text = ''.join([_[0] for _ in retrieve_result[:top_k]])
prompt = f"给定下面的文本,请问答用户的问题。\n\n{text}\n\n问题:{query}"

client = OpenAI(
api_key=os.environ.get("OPENAI_API_KEY"), # This is the default and can be omitted
)

chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": prompt,
}
],
model="gpt-4o-mini",
)
return chat_completion.choices[0].message.content


naive_embedding_answer = get_answer(query=query, retrieve_result=naive_embedding_order)
print(f"query: {query}, 朴素嵌入时RAG过程中LLM的回复:{naive_embedding_answer}")
late_chunking_answer = get_answer(query=query, retrieve_result=late_chunking_order)
print(f"query: {query}, 迟分嵌入时RAG过程中LLM的回复:{late_chunking_answer}")

根据示例文本,笔者测试了三个简单的问题,对比结果如下:

序号 问题 朴素嵌入时LLM回复 Late Chunking时LLM回复
1 蔚来ET9的车身强度是多少? 文本中并没有提供关于蔚来ET9车身强度的具体信息,所以无法回答这个问题。 蔚来ET9的车身强度达到2000MPa。
2 蔚来ET9有多少个电动遮阳帘? 文本中并未提到蔚来ET9的电动遮阳帘数量。因此,无法回答这个问题。 蔚来ET9配备了七扇电动遮阳帘。
3 蔚来ET9中的冰箱的最大容积是多少? 文本中并没有提到蔚来ET9中包含冰箱或相关信息。因此,无法提供关于冰箱最大容积的数据。提到的是前备箱的容积为105L。 如果您有其他问题,请告知! 蔚来ET9中的冰箱的最大容积达到10升。

当然,这只是演示了几个Late Chunking的召回效果比朴素嵌入效果好的例子,并不是说Late Chunking的召回效果就一定会比朴素嵌入效果好。

那么,沃我们在实际场景中该如何选择分块策略呢?

  • 对于朴素分块,适合场景为:主题多样,需要检索特定信息;需要展示局部文本片段。
  • 对于Late Chunking,适合场景为:主题连贯,需要上下文信息;需要平衡局部细节和全局语义。

读者可以仔细观察上面关于蔚来ET9的文章,测试的三个例子都来自这样的段落:段落中未提及蔚来ET9,而是用指代,但人类不难用上下文得到这些信息。说到来,这是一篇蔚来ET9的文章,主体连贯,因此很适合用Late Chunking,其保留上下文的能力在这种场景下会比朴素分块好。

总结

本文主要介绍了Embedding模型中的Late Chunking(迟分)策略,演示多个中文Late Chunking的例子,并搭建相关Gradio服务,最后再展示其在RAG框架中对于大模型回复质量的提升作用。

上面给出的Python代码均已开源至Github,网址为: https://github.com/percent4/embedding_rerank_retrieval

参考文献

  1. Jina AI官网: https://jina.ai/
  2. 在 Notebook的例子: https://colab.research.google.com/drive/15vNZb6AsU7byjYoaEtXuNu567JWNzXOz?usp=sharing#scrollTo=abe3d93b9e6609b9
  3. Jina AI官网关于Late Chunking的介绍: https://jina.ai/news/late-chunking-in-long-context-embedding-models/
  4. Late Chunking论文: https://arxiv.org/pdf/2409.04701
  5. 卷起来了!长文本向量模型分块策略大比拼: https://mp.weixin.qq.com/s/tWToc7Lu18nb6TwuZ_bz1g

NLP(一百零九)Embedding中的Late-Chunking(迟分)策略
https://percent4.github.io/NLP(一百零九)Embedding中的Late-Chunking(迟分)策略/
作者
Jclian91
发布于
2025年1月8日
许可协议