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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179 | class AutoMergingRetriever(BaseRetriever):
"""这个检索器将尝试将上下文合并到父上下文中。
检索器首先从向量存储中检索块。
然后，它将尝试将这些块合并成一个单一的上下文。"""
def __init__(
self,
vector_retriever: VectorIndexRetriever,
storage_context: StorageContext,
simple_ratio_thresh: float = 0.5,
verbose: bool = False,
callback_manager: Optional[CallbackManager] = None,
object_map: Optional[dict] = None,
objects: Optional[List[IndexNode]] = None,
) -> None:
"""初始化参数。"""
self._vector_retriever = vector_retriever
self._storage_context = storage_context
self._simple_ratio_thresh = simple_ratio_thresh
super().__init__(
callback_manager=callback_manager,
object_map=object_map,
objects=objects,
verbose=verbose,
)
def _get_parents_and_merge(
self, nodes: List[NodeWithScore]
) -> Tuple[List[NodeWithScore], bool]:
"""获取父节点并合并节点。"""
# retrieve all parent nodes
parent_nodes: Dict[str, BaseNode] = {}
parent_cur_children_dict: Dict[str, List[NodeWithScore]] = defaultdict(list)
for node in nodes:
if node.node.parent_node is None:
continue
parent_node_info = node.node.parent_node
# Fetch actual parent node if doesn't exist in `parent_nodes` cache yet
parent_node_id = parent_node_info.node_id
if parent_node_id not in parent_nodes:
parent_node = self._storage_context.docstore.get_document(
parent_node_id
)
parent_nodes[parent_node_id] = cast(BaseNode, parent_node)
# add reference to child from parent
parent_cur_children_dict[parent_node_id].append(node)
# compute ratios and "merge" nodes
# merging: delete some children nodes, add some parent nodes
node_ids_to_delete = set()
nodes_to_add: Dict[str, BaseNode] = {}
for parent_node_id, parent_node in parent_nodes.items():
parent_child_nodes = parent_node.child_nodes
parent_num_children = len(parent_child_nodes) if parent_child_nodes else 1
parent_cur_children = parent_cur_children_dict[parent_node_id]
ratio = len(parent_cur_children) / parent_num_children
# if ratio is high enough, merge
if ratio > self._simple_ratio_thresh:
node_ids_to_delete.update(
set({n.node.node_id for n in parent_cur_children})
)
parent_node_text = truncate_text(parent_node.text, 100)
info_str = (
f"> Merging {len(parent_cur_children)} nodes into parent node.\n"
f"> Parent node id: {parent_node_id}.\n"
f"> Parent node text: {parent_node_text}\n"
)
logger.info(info_str)
if self._verbose:
print(info_str)
# add parent node
# can try averaging score across embeddings for now
avg_score = sum(
[n.get_score() or 0.0 for n in parent_cur_children]
) / len(parent_cur_children)
parent_node_with_score = NodeWithScore(
node=parent_node, score=avg_score
)
nodes_to_add[parent_node_id] = parent_node_with_score
# delete old child nodes, add new parent nodes
new_nodes = [n for n in nodes if n.node.node_id not in node_ids_to_delete]
# add parent nodes
new_nodes.extend(list(nodes_to_add.values()))
is_changed = len(node_ids_to_delete) > 0
return new_nodes, is_changed
def _fill_in_nodes(
self, nodes: List[NodeWithScore]
) -> Tuple[List[NodeWithScore], bool]:
"""填充节点。"""
new_nodes = []
is_changed = False
for idx, node in enumerate(nodes):
new_nodes.append(node)
if idx >= len(nodes) - 1:
continue
cur_node = cast(BaseNode, node.node)
# if there's a node in the middle, add that to the queue
if (
cur_node.next_node is not None
and cur_node.next_node == nodes[idx + 1].node.prev_node
):
is_changed = True
next_node = self._storage_context.docstore.get_document(
cur_node.next_node.node_id
)
next_node = cast(BaseNode, next_node)
next_node_text = truncate_text(next_node.get_text(), 100)
info_str = (
f"> Filling in node. Node id: {cur_node.next_node.node_id}"
f"> Node text: {next_node_text}\n"
)
logger.info(info_str)
if self._verbose:
print(info_str)
# set score to be average of current node and next node
avg_score = (node.get_score() + nodes[idx + 1].get_score()) / 2
new_nodes.append(NodeWithScore(node=next_node, score=avg_score))
return new_nodes, is_changed
def _try_merging(
self, nodes: List[NodeWithScore]
) -> Tuple[List[NodeWithScore], bool]:
"""尝试不同的方法来合并节点。"""
# first try filling in nodes
nodes, is_changed_0 = self._fill_in_nodes(nodes)
# then try merging nodes
nodes, is_changed_1 = self._get_parents_and_merge(nodes)
return nodes, is_changed_0 or is_changed_1
def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
"""获取给定查询的节点。
由用户实现。
"""
initial_nodes = self._vector_retriever.retrieve(query_bundle)
cur_nodes, is_changed = self._try_merging(initial_nodes)
# cur_nodes, is_changed = self._get_parents_and_merge(initial_nodes)
while is_changed:
cur_nodes, is_changed = self._try_merging(cur_nodes)
# cur_nodes, is_changed = self._get_parents_and_merge(cur_nodes)
# sort by similarity
cur_nodes.sort(key=lambda x: x.get_score(), reverse=True)
return cur_nodes
|