Skip to content

Commit b91d032

Browse files
jereliucopybara-github
jereliu
authored and
copybara-github
committed
Add retrieval utils.
PiperOrigin-RevId: 480620564
1 parent 7ae4fc4 commit b91d032

1 file changed

Lines changed: 379 additions & 0 deletions

File tree

‎baselines/t5/data/deepbank/retrieval_utils.py‎

Lines changed: 379 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,379 @@
1+
# coding=utf-8
2+
# Copyright 2022 The Uncertainty Baselines Authors.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
"""Retrieval utils."""
17+
import re
18+
from typing import Text, Dict, Any, List, Optional
19+
20+
import t5.data
21+
from data.deepbank import graph_utils # local file import from baselines.t5
22+
from data.deepbank import meta_graph_utils # local file import from baselines.t5
23+
from data.deepbank import penman_utils # local file import from baselines.t5
24+
25+
DEFAULT_VOCAB = t5.data.get_default_vocabulary()
26+
27+
28+
def _check_pred_input(pred: Dict[Text, Any], beam_id: int = 0):
29+
"""Checks if prediction input is well-formed."""
30+
if not (f'prediction_{beam_id}_ids' in pred and
31+
f'prediction_{beam_id}' in pred):
32+
raise ValueError(f'Some of the following items are missing in {pred}:'
33+
f'`prediction_{beam_id}_ids`, `prediction_{beam_id}`.')
34+
35+
36+
def _get_pred_graph_info(sentence: Text,
37+
pred: Dict[Text, Any],
38+
tgt_penman: penman_utils.PENMANStr,
39+
beam_id: int = 0,
40+
data_version: str = 'v0',
41+
prefix: str = 'x'):
42+
"""Gets prediction graph infos from prediction Dict."""
43+
return meta_graph_utils.GraphInfo(
44+
token_ids=pred[f'prediction_{beam_id}_ids'],
45+
beam_scores=pred['beam_scores'][beam_id],
46+
sentence=sentence,
47+
prediction=pred[f'prediction_{beam_id}'],
48+
target=tgt_penman.retokened_variable_free_penman,
49+
data_version=data_version,
50+
prefix=prefix
51+
)
52+
53+
54+
def _random_update_src(src: Text, tgt: Text,
55+
total_num_examplar: int = 0, max_num_examplar: int = 1,
56+
depth: int = 1, use_alignment: bool = True,
57+
use_custom_token: bool = True,
58+
max_seq_length: int = 512, data_version: str = 'v0'):
59+
"""Updates input sentence with new randomly retrieved examplars."""
60+
src_length = len(DEFAULT_VOCAB.encode(src))
61+
while True:
62+
if use_alignment:
63+
subgraph_penman_str, align_sent = graph_utils.get_random_linear_subgraph(
64+
tgt, src, level=depth)
65+
examplar_str = ' @@ %s' % align_sent
66+
else:
67+
subgraph_penman_str = graph_utils.get_random_linear_subgraph(
68+
tgt, src, level=depth, return_align_sent=False)
69+
examplar_str = ''
70+
subgraph_penman = penman_utils.PENMANStr(
71+
subgraph_penman_str,
72+
variable_free=False,
73+
data_version=data_version)
74+
subgraph_output = (
75+
subgraph_penman.retokened_variable_free_penman if use_custom_token
76+
else subgraph_penman.variable_free_penman)
77+
examplar_str += ' ## %s' % subgraph_output
78+
examplar_length = len(DEFAULT_VOCAB.encode(examplar_str))
79+
if src_length + examplar_length < max_seq_length and (total_num_examplar <
80+
max_num_examplar):
81+
total_num_examplar += 1
82+
src += examplar_str
83+
src_length += examplar_length
84+
else:
85+
break
86+
return src
87+
88+
89+
def _oracle_update_src(src: Text, tgt: Text,
90+
pred_graph_info: meta_graph_utils.GraphInfo,
91+
tgt_dag: graph_utils.DAG,
92+
tgt_prefix: Text = 'x', pred_prefix: Text = 'y',
93+
total_num_examplar: int = 0,
94+
max_num_examplar: int = 1,
95+
depth: int = 1, use_alignment: bool = True,
96+
use_custom_token: bool = True,
97+
max_seq_length: int = 512, data_version: str = 'v0',
98+
with_uncertain: bool = False):
99+
"""Updates input sentence with new retrieved examplars based on oracle."""
100+
src_length = len(DEFAULT_VOCAB.encode(src))
101+
tgt_instances, tgt_attributes, tgt_relations = tgt_dag.get_triples()
102+
pred_graph = meta_graph_utils.MetaGraph(pred_graph_info, pred_prefix)
103+
# Sets `compate_attribute` to False to avoid compate alignment information,
104+
# which the prediction does not have.
105+
mapping, _ = graph_utils.get_best_match(tgt_instances, tgt_attributes,
106+
tgt_relations,
107+
pred_graph_info.instances,
108+
pred_graph_info.attributes,
109+
pred_graph_info.relations, tgt_prefix,
110+
pred_prefix,
111+
compare_attribute=False)
112+
oracle_node_idxs = list(
113+
graph_utils.find_mismatched_node_idxs(mapping, tgt_prefix, pred_prefix,
114+
tgt_instances, tgt_attributes,
115+
tgt_relations,
116+
pred_graph_info.instances,
117+
pred_graph_info.attributes,
118+
pred_graph_info.relations))
119+
if with_uncertain:
120+
# If `with_uncertain`, re-order `oracle_node_idxs` based on probabilties
121+
# (ascending order).
122+
mapping_dict = graph_utils.get_mapping_dict(
123+
tgt_prefix, pred_prefix, mapping, reverse=True)
124+
uncertain_node_idxs = graph_utils.find_uncertain_node_idxs(
125+
pred_graph.instance_prob_dict,
126+
pred_graph.attribute_prob_dict,
127+
pred_graph.relation_prob_dict,
128+
mapping_dict)
129+
oracle_node_idxs = [
130+
idx for idx in uncertain_node_idxs if idx in oracle_node_idxs]
131+
excluded_node_idxs = []
132+
while oracle_node_idxs:
133+
oracle_node_idx = oracle_node_idxs.pop()
134+
excluded_node_idxs.append(oracle_node_idx)
135+
if use_alignment:
136+
(subgraph_penman_str, align_sent,
137+
oracle_node_idxs) = graph_utils.get_oracle_linear_subgraph(
138+
tgt_instances, tgt_attributes, tgt_relations,
139+
src, oracle_node_idx, oracle_node_idxs, level=depth)
140+
examplar_str = ' @@ %s' % align_sent
141+
else:
142+
(subgraph_penman_str,
143+
oracle_node_idxs) = graph_utils.get_oracle_linear_subgraph(
144+
tgt_instances, tgt_attributes, tgt_relations,
145+
src, oracle_node_idx, oracle_node_idxs, level=depth,
146+
return_align_sent=False)
147+
examplar_str = ''
148+
subgraph_penman = penman_utils.PENMANStr(
149+
subgraph_penman_str,
150+
variable_free=False,
151+
data_version=data_version)
152+
subgraph_output = (
153+
subgraph_penman.retokened_variable_free_penman if use_custom_token
154+
else subgraph_penman.variable_free_penman)
155+
examplar_str += ' ## %s' % subgraph_output
156+
examplar_length = len(DEFAULT_VOCAB.encode(examplar_str))
157+
if src_length + examplar_length < max_seq_length and (total_num_examplar <
158+
max_num_examplar):
159+
total_num_examplar += 1
160+
src += examplar_str
161+
src_length += examplar_length
162+
else:
163+
break
164+
if src_length < max_seq_length and total_num_examplar < max_num_examplar:
165+
# If the number of oracle retrival examplars has not reached the budget,
166+
# and the input sequence length has not reached the max sequence length,
167+
# we left the rest for uncertain retrieval or random retrieval based
168+
# on gold.
169+
if with_uncertain:
170+
src = _uncertain_update_src(src, tgt, pred_graph_info,
171+
tgt_dag, tgt_prefix,
172+
pred_prefix, total_num_examplar,
173+
max_num_examplar, depth, use_alignment,
174+
use_custom_token, max_seq_length,
175+
data_version, excluded_node_idxs)
176+
else:
177+
src = _random_update_src(src, tgt, total_num_examplar, max_num_examplar,
178+
depth, use_alignment, use_custom_token,
179+
max_seq_length, data_version)
180+
return src
181+
182+
183+
def _uncertain_update_src(src: Text, tgt: Text,
184+
pred_graph_info: meta_graph_utils.GraphInfo,
185+
tgt_dag: graph_utils.DAG,
186+
tgt_prefix: Text = 'x', pred_prefix: Text = 'y',
187+
total_num_examplar: int = 0,
188+
max_num_examplar: int = 1,
189+
depth: int = 1, use_alignment: bool = True,
190+
use_custom_token: bool = True,
191+
max_seq_length: int = 512, data_version: str = 'v0',
192+
excluded_node_idxs: Optional[List[Text]] = None):
193+
"""Updates input sentence with new retrieved examplars based on uncertainty."""
194+
src_length = len(DEFAULT_VOCAB.encode(src))
195+
tgt_instances, tgt_attributes, tgt_relations = tgt_dag.get_triples()
196+
pred_graph = meta_graph_utils.MetaGraph(pred_graph_info, pred_prefix)
197+
# Sets `compate_attribute` to False to avoid compate alignment information,
198+
# which the prediction does not have.
199+
mapping, _ = graph_utils.get_best_match(pred_graph.instances,
200+
pred_graph.attributes,
201+
pred_graph.relations,
202+
tgt_instances, tgt_attributes,
203+
tgt_relations,
204+
pred_prefix,
205+
tgt_prefix,
206+
compare_attribute=False)
207+
mapping_dict = graph_utils.get_mapping_dict(pred_prefix, tgt_prefix, mapping)
208+
uncertain_node_idxs = graph_utils.find_uncertain_node_idxs(
209+
pred_graph.instance_prob_dict, pred_graph.attribute_prob_dict,
210+
pred_graph.relation_prob_dict, mapping_dict)
211+
# Excludes node indexes in predefined node index list.
212+
if excluded_node_idxs:
213+
uncertain_node_idxs = [
214+
idx for idx in uncertain_node_idxs if idx not in excluded_node_idxs]
215+
while uncertain_node_idxs:
216+
uncertain_node_idx = uncertain_node_idxs.pop()
217+
if use_alignment:
218+
(subgraph_penman_str, align_sent,
219+
uncertain_node_idxs) = graph_utils.get_uncertain_linear_subgraph(
220+
tgt_instances, tgt_attributes, tgt_relations,
221+
src, uncertain_node_idx, uncertain_node_idxs, level=depth)
222+
examplar_str = ' @@ %s' % align_sent
223+
else:
224+
(subgraph_penman_str,
225+
uncertain_node_idxs) = graph_utils.get_uncertain_linear_subgraph(
226+
tgt_instances, tgt_attributes, tgt_relations,
227+
src, uncertain_node_idx, uncertain_node_idxs, level=depth,
228+
return_align_sent=False)
229+
examplar_str = ''
230+
subgraph_penman = penman_utils.PENMANStr(
231+
subgraph_penman_str,
232+
variable_free=False,
233+
data_version=data_version)
234+
subgraph_output = (
235+
subgraph_penman.retokened_variable_free_penman if use_custom_token
236+
else subgraph_penman.variable_free_penman)
237+
examplar_str += ' ## %s' % subgraph_output
238+
examplar_length = len(DEFAULT_VOCAB.encode(examplar_str))
239+
if src_length + examplar_length < max_seq_length and (total_num_examplar <
240+
max_num_examplar):
241+
total_num_examplar += 1
242+
src += examplar_str
243+
src_length += examplar_length
244+
else:
245+
break
246+
if src_length < max_seq_length and total_num_examplar < max_num_examplar:
247+
# If the number of uncertain retrival examplars has not reached the budget,
248+
# and the input sequence length has not reached the max sequence length,
249+
# we left the rest for random retrieval based on gold.
250+
src = _random_update_src(src, tgt, total_num_examplar, max_num_examplar,
251+
depth, use_alignment, use_custom_token,
252+
max_seq_length, data_version)
253+
return src
254+
255+
256+
def random_retrieval_on_gold(src: Text,
257+
tgt: Text,
258+
max_num_examplar: int = 1,
259+
depth: int = 1,
260+
use_alignment: bool = True,
261+
use_custom_token: bool = True,
262+
max_seq_length: int = 512,
263+
data_version: str = 'v0'):
264+
"""Retrieves random subgraphs based on gold graphs."""
265+
if not max_num_examplar:
266+
max_num_examplar = max_seq_length
267+
tgt_no_alignment = re.sub(r' :lnk "<[0-9]+:[0-9]+>"', '', tgt)
268+
tgt_penman = penman_utils.PENMANStr(
269+
tgt_no_alignment, variable_free=False, data_version=data_version)
270+
src = _random_update_src(src, tgt, 0, max_num_examplar,
271+
depth, use_alignment, use_custom_token,
272+
max_seq_length, data_version)
273+
if use_custom_token:
274+
return src, tgt_penman.retokened_variable_free_penman
275+
else:
276+
return src, tgt_penman.variable_free_penman
277+
278+
279+
def oracle_retrieval_on_gold(src: Text,
280+
tgt: Text,
281+
pred: Dict[Text, Any],
282+
max_num_examplar: int = 1,
283+
depth: int = 1,
284+
use_alignment: bool = True,
285+
use_custom_token: bool = True,
286+
max_seq_length: int = 512,
287+
data_version: str = 'v0',
288+
beam_id: int = 0,
289+
with_uncertain: bool = False):
290+
"""Retrieves oracle subgraphs based on gold graphs."""
291+
if not max_num_examplar:
292+
max_num_examplar = max_seq_length
293+
tgt_prefix = 'x'
294+
pred_prefix = 'y'
295+
tgt_no_alignment = re.sub(r' :lnk "<[0-9]+:[0-9]+>"', '', tgt)
296+
tgt_penman = penman_utils.PENMANStr(
297+
tgt_no_alignment, variable_free=False, data_version=data_version)
298+
tgt_dag = graph_utils.parse_string_to_dag(tgt)
299+
tgt_dag.change_node_prefix(tgt_prefix)
300+
_check_pred_input(pred, beam_id)
301+
pred_graph_info = _get_pred_graph_info(src, pred, tgt_penman, beam_id,
302+
data_version, pred_prefix)
303+
if not pred_graph_info.pred_parsed:
304+
# The prediction is an ill-formed graph, if `with_uncertain` is True,
305+
# use uncertain retrieval on gold instead, otherwise use random retrieval
306+
# on gold instead.
307+
if with_uncertain:
308+
return uncertain_retrieval_on_gold(src, tgt, pred, max_num_examplar,
309+
depth, use_alignment,
310+
use_custom_token, max_seq_length,
311+
data_version)
312+
else:
313+
return random_retrieval_on_gold(src, tgt, max_num_examplar, depth,
314+
use_alignment, use_custom_token,
315+
max_seq_length, data_version)
316+
src = _oracle_update_src(src, tgt, pred_graph_info, tgt_dag, tgt_prefix,
317+
pred_prefix, 0, max_num_examplar,
318+
depth, use_alignment, use_custom_token,
319+
max_seq_length, data_version, with_uncertain)
320+
if use_custom_token:
321+
return src, tgt_penman.retokened_variable_free_penman
322+
else:
323+
return src, tgt_penman.variable_free_penman
324+
325+
326+
def uncertain_retrieval_on_gold(src: Text,
327+
tgt: Text,
328+
pred: Dict[Text, Any],
329+
max_num_examplar: int = 1,
330+
depth: int = 1,
331+
use_alignment: bool = True,
332+
use_custom_token: bool = True,
333+
max_seq_length: int = 512,
334+
data_version: str = 'v0',
335+
beam_id: int = 0):
336+
"""Retrieves uncertain subgraphs based on gold graphs."""
337+
if not max_num_examplar:
338+
max_num_examplar = max_seq_length
339+
tgt_prefix = 'x'
340+
pred_prefix = 'y'
341+
tgt_no_alignment = re.sub(r' :lnk "<[0-9]+:[0-9]+>"', '', tgt)
342+
tgt_penman = penman_utils.PENMANStr(
343+
tgt_no_alignment, variable_free=False, data_version=data_version)
344+
tgt_dag = graph_utils.parse_string_to_dag(tgt)
345+
tgt_dag.change_node_prefix(tgt_prefix)
346+
_check_pred_input(pred, beam_id)
347+
pred_graph_info = _get_pred_graph_info(src, pred, tgt_penman, beam_id,
348+
data_version, pred_prefix)
349+
if not pred_graph_info.pred_parsed:
350+
# The prediction is an ill-formed graph, use random retrieval
351+
# on gold instead.
352+
return random_retrieval_on_gold(src, tgt, max_num_examplar, depth,
353+
use_alignment, use_custom_token,
354+
max_seq_length, data_version)
355+
src = _uncertain_update_src(src, tgt, pred_graph_info, tgt_dag, tgt_prefix,
356+
pred_prefix, 0, max_num_examplar,
357+
depth, use_alignment, use_custom_token,
358+
max_seq_length, data_version)
359+
if use_custom_token:
360+
return src, tgt_penman.retokened_variable_free_penman
361+
else:
362+
return src, tgt_penman.variable_free_penman
363+
364+
365+
def oracle_uncertain_retrieval_on_gold(src: Text,
366+
tgt: Text,
367+
pred: Dict[Text, Any],
368+
max_num_examplar: int = 1,
369+
depth: int = 1,
370+
use_alignment: bool = True,
371+
use_custom_token: bool = True,
372+
max_seq_length: int = 512,
373+
data_version: str = 'v0',
374+
beam_id: int = 0):
375+
"""Retrieves oracle subgraphs based on uncertainty-ordered gold graphs."""
376+
return oracle_retrieval_on_gold(src, tgt, pred, max_num_examplar, depth,
377+
use_alignment, use_custom_token,
378+
max_seq_length, data_version, beam_id,
379+
with_uncertain=True)

0 commit comments

Comments
 (0)