Spaces:
Runtime error
Runtime error
File size: 4,582 Bytes
fc8c192 |
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 |
from collections import defaultdict
class VQASerTokenChunk(object):
def __init__(self, max_seq_len=512, infer_mode=False, **kwargs):
self.max_seq_len = max_seq_len
self.infer_mode = infer_mode
def __call__(self, data):
encoded_inputs_all = []
seq_len = len(data["input_ids"])
for index in range(0, seq_len, self.max_seq_len):
chunk_beg = index
chunk_end = min(index + self.max_seq_len, seq_len)
encoded_inputs_example = {}
for key in data:
if key in [
"label",
"input_ids",
"labels",
"token_type_ids",
"bbox",
"attention_mask",
]:
if self.infer_mode and key == "labels":
encoded_inputs_example[key] = data[key]
else:
encoded_inputs_example[key] = data[key][chunk_beg:chunk_end]
else:
encoded_inputs_example[key] = data[key]
encoded_inputs_all.append(encoded_inputs_example)
if len(encoded_inputs_all) == 0:
return None
return encoded_inputs_all[0]
class VQAReTokenChunk(object):
def __init__(
self, max_seq_len=512, entities_labels=None, infer_mode=False, **kwargs
):
self.max_seq_len = max_seq_len
self.entities_labels = (
{"HEADER": 0, "QUESTION": 1, "ANSWER": 2}
if entities_labels is None
else entities_labels
)
self.infer_mode = infer_mode
def __call__(self, data):
# prepare data
entities = data.pop("entities")
relations = data.pop("relations")
encoded_inputs_all = []
for index in range(0, len(data["input_ids"]), self.max_seq_len):
item = {}
for key in data:
if key in [
"label",
"input_ids",
"labels",
"token_type_ids",
"bbox",
"attention_mask",
]:
if self.infer_mode and key == "labels":
item[key] = data[key]
else:
item[key] = data[key][index : index + self.max_seq_len]
else:
item[key] = data[key]
# select entity in current chunk
entities_in_this_span = []
global_to_local_map = {} #
for entity_id, entity in enumerate(entities):
if (
index <= entity["start"] < index + self.max_seq_len
and index <= entity["end"] < index + self.max_seq_len
):
entity["start"] = entity["start"] - index
entity["end"] = entity["end"] - index
global_to_local_map[entity_id] = len(entities_in_this_span)
entities_in_this_span.append(entity)
# select relations in current chunk
relations_in_this_span = []
for relation in relations:
if (
index <= relation["start_index"] < index + self.max_seq_len
and index <= relation["end_index"] < index + self.max_seq_len
):
relations_in_this_span.append(
{
"head": global_to_local_map[relation["head"]],
"tail": global_to_local_map[relation["tail"]],
"start_index": relation["start_index"] - index,
"end_index": relation["end_index"] - index,
}
)
item.update(
{
"entities": self.reformat(entities_in_this_span),
"relations": self.reformat(relations_in_this_span),
}
)
if len(item["entities"]) > 0:
item["entities"]["label"] = [
self.entities_labels[x] for x in item["entities"]["label"]
]
encoded_inputs_all.append(item)
if len(encoded_inputs_all) == 0:
return None
return encoded_inputs_all[0]
def reformat(self, data):
new_data = defaultdict(list)
for item in data:
for k, v in item.items():
new_data[k].append(v)
return new_data
|