wehub-resource-sync
268 lines
10 KiB
Python
268 lines
10 KiB
Python
import tqdm
|
|
from torch.utils.data import Dataset
|
|
from markuplmft.data.tag_utils import tags_dict
|
|
import pickle
|
|
import os
|
|
import constants
|
|
|
|
|
|
class SwdeFeature(object):
|
|
def __init__(self,
|
|
html_path,
|
|
input_ids,
|
|
token_type_ids,
|
|
attention_mask,
|
|
xpath_tags_seq,
|
|
xpath_subs_seq,
|
|
labels,
|
|
involved_first_tokens_pos,
|
|
involved_first_tokens_xpaths,
|
|
involved_first_tokens_types,
|
|
involved_first_tokens_text,
|
|
):
|
|
"""
|
|
html_path: indicate which page the feature belongs to
|
|
input_ids: RT
|
|
token_type_ids: RT
|
|
attention_mask: RT
|
|
xpath_tags_seq: RT
|
|
xpath_subs_seq: RT
|
|
labels: RT
|
|
involved_first_tokens_pos: a list, indicate the positions of the first-tokens in this feature
|
|
involved_first_tokens_xpaths: the xpaths of the first-tokens, used to build dict
|
|
involved_first_tokens_types: the types of the first-tokens
|
|
involved_first_tokens_text: the text of the first tokens
|
|
|
|
Note that `involved_xxx` are not fixed-length array, so they shouldn't be sent into our model
|
|
They are just used for evaluation
|
|
"""
|
|
self.html_path = html_path
|
|
self.input_ids = input_ids
|
|
self.token_type_ids = token_type_ids
|
|
self.attention_mask = attention_mask
|
|
self.xpath_tags_seq = xpath_tags_seq
|
|
self.xpath_subs_seq = xpath_subs_seq
|
|
self.labels = labels
|
|
self.involved_first_tokens_pos = involved_first_tokens_pos
|
|
self.involved_first_tokens_xpaths = involved_first_tokens_xpaths
|
|
self.involved_first_tokens_types = involved_first_tokens_types
|
|
self.involved_first_tokens_text = involved_first_tokens_text
|
|
|
|
|
|
class SwdeDataset(Dataset):
|
|
def __init__(self,
|
|
all_input_ids,
|
|
all_attention_mask,
|
|
all_token_type_ids,
|
|
all_xpath_tags_seq,
|
|
all_xpath_subs_seq,
|
|
all_labels=None,
|
|
):
|
|
'''
|
|
print(type(all_input_ids))
|
|
print(type(all_attention_mask))
|
|
print(type(all_token_type_ids))
|
|
print(type(all_xpath_tags_seq))
|
|
print(type(all_xpath_subs_seq))
|
|
print(type(all_labels))
|
|
raise ValueError
|
|
'''
|
|
self.tensors = [all_input_ids, all_attention_mask, all_token_type_ids,
|
|
all_xpath_tags_seq, all_xpath_subs_seq]
|
|
|
|
if not all_labels is None:
|
|
self.tensors.append(all_labels)
|
|
|
|
def __len__(self):
|
|
return len(self.tensors[0])
|
|
|
|
def __getitem__(self, index):
|
|
return tuple(tensor[index] for tensor in self.tensors)
|
|
|
|
|
|
def process_xpath(xpath: str):
|
|
if xpath.endswith("/tail"):
|
|
xpath = xpath[:-5]
|
|
|
|
xpath_tags_seq, xpath_subs_seq = [], []
|
|
units = xpath.split("/")
|
|
|
|
for unit in units:
|
|
if not unit:
|
|
continue
|
|
if '[' not in unit:
|
|
xpath_tags_seq.append(tags_dict.get(unit, 215))
|
|
xpath_subs_seq.append(0)
|
|
else:
|
|
xx = unit.split('[')
|
|
name = xx[0]
|
|
id = int(xx[1][:-1])
|
|
xpath_tags_seq.append(tags_dict.get(name, 215))
|
|
xpath_subs_seq.append(min(id, 1000))
|
|
|
|
assert len(xpath_subs_seq) == len(xpath_tags_seq)
|
|
|
|
if len(xpath_tags_seq) > 50:
|
|
xpath_tags_seq = xpath_tags_seq[-50:]
|
|
xpath_subs_seq = xpath_subs_seq[-50:]
|
|
|
|
xpath_tags_seq = xpath_tags_seq + [216] * (50 - len(xpath_tags_seq))
|
|
xpath_subs_seq = xpath_subs_seq + [1001] * (50 - len(xpath_subs_seq))
|
|
|
|
return xpath_tags_seq, xpath_subs_seq
|
|
|
|
|
|
def get_swde_features(root_dir, vertical, website, tokenizer,
|
|
doc_stride, max_length, prev_nodes, n_pages):
|
|
real_max_token_num = max_length - 2 # for cls and sep
|
|
padded_xpath_tags_seq = [216] * 50
|
|
padded_xpath_subs_seq = [1001] * 50
|
|
|
|
filename = os.path.join(root_dir, f"{vertical}-{website}-{n_pages}.pickle")
|
|
with open(filename, "rb") as f:
|
|
raw_data = pickle.load(f)
|
|
|
|
features = []
|
|
|
|
for index in tqdm.tqdm(raw_data, desc=f"Processing {vertical}-{website}-{n_pages} features ..."):
|
|
html_path = f"{vertical}-{website}-{index}.htm"
|
|
needed_docstrings_id_set = set()
|
|
for i in range(len(raw_data[index])):
|
|
doc_string_type = raw_data[index][i][2]
|
|
if doc_string_type == "fixed-node":
|
|
continue
|
|
# we take i-3, i-2, i-1 into account
|
|
|
|
needed_docstrings_id_set.add(i)
|
|
|
|
used_prev = 0
|
|
prev_id = i - 1
|
|
while prev_id >= 0 and used_prev < prev_nodes:
|
|
if raw_data[index][prev_id][0].strip():
|
|
needed_docstrings_id_set.add(prev_id)
|
|
used_prev += 1
|
|
prev_id -= 1
|
|
|
|
needed_docstrings_id_list = sorted(list(needed_docstrings_id_set))
|
|
|
|
all_token_ids_seq = []
|
|
all_xpath_tags_seq = []
|
|
all_xpath_subs_seq = []
|
|
token_to_ori_map_seq = []
|
|
all_labels_seq = []
|
|
|
|
first_token_pos = []
|
|
first_token_xpaths = []
|
|
first_token_type = []
|
|
first_token_text = []
|
|
|
|
for i, needed_id in enumerate(needed_docstrings_id_list):
|
|
text = raw_data[index][needed_id][0]
|
|
xpath = raw_data[index][needed_id][1]
|
|
type = raw_data[index][needed_id][2]
|
|
token_ids = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(text))
|
|
|
|
xpath_tags_seq, xpath_subs_seq = process_xpath(xpath)
|
|
|
|
all_token_ids_seq += token_ids
|
|
all_xpath_tags_seq += [xpath_tags_seq] * len(token_ids)
|
|
all_xpath_subs_seq += [xpath_subs_seq] * len(token_ids)
|
|
|
|
token_to_ori_map_seq += [i] * len(token_ids)
|
|
|
|
if type == "fixed-node":
|
|
all_labels_seq += [-100] * len(token_ids)
|
|
else:
|
|
# we always use the first token to predict
|
|
first_token_pos.append(len(all_labels_seq))
|
|
first_token_type.append(type)
|
|
first_token_xpaths.append(xpath)
|
|
first_token_text.append(text)
|
|
|
|
all_labels_seq += [constants.ATTRIBUTES_PLUS_NONE[vertical].index(type)] * len(token_ids)
|
|
|
|
assert len(all_token_ids_seq) == len(all_xpath_tags_seq)
|
|
assert len(all_token_ids_seq) == len(all_xpath_subs_seq)
|
|
assert len(all_token_ids_seq) == len(all_labels_seq)
|
|
|
|
# we have all the pos of variable nodes in all_token_ids_seq
|
|
# now we need to assign them into each feature
|
|
|
|
start_pos = 0
|
|
flag = False
|
|
|
|
curr_first_token_index = 0
|
|
|
|
while True:
|
|
|
|
# invloved is [ start_pos , end_pos )
|
|
|
|
token_type_ids = [0] * max_length # that is always this
|
|
|
|
end_pos = start_pos + real_max_token_num
|
|
# add start_pos ~ end_pos as a feature
|
|
splited_token_ids_seq = [tokenizer.cls_token_id] + all_token_ids_seq[start_pos:end_pos] + [
|
|
tokenizer.sep_token_id]
|
|
splited_xpath_tags_seq = [padded_xpath_tags_seq] + all_xpath_tags_seq[start_pos:end_pos] + [
|
|
padded_xpath_tags_seq]
|
|
splited_xpath_subs_seq = [padded_xpath_subs_seq] + all_xpath_subs_seq[start_pos:end_pos] + [
|
|
padded_xpath_subs_seq]
|
|
splited_labels_seq = [-100] + all_labels_seq[start_pos:end_pos] + [-100]
|
|
|
|
# locate first-tokens in them
|
|
involved_first_tokens_pos = []
|
|
involved_first_tokens_xpaths = []
|
|
involved_first_tokens_types = []
|
|
involved_first_tokens_text = []
|
|
|
|
while curr_first_token_index < len(first_token_pos) \
|
|
and end_pos > first_token_pos[curr_first_token_index] >= start_pos:
|
|
involved_first_tokens_pos.append(
|
|
first_token_pos[curr_first_token_index] - start_pos + 1) # +1 for [cls]
|
|
involved_first_tokens_xpaths.append(first_token_xpaths[curr_first_token_index])
|
|
involved_first_tokens_types.append(first_token_type[curr_first_token_index])
|
|
involved_first_tokens_text.append(first_token_text[curr_first_token_index])
|
|
curr_first_token_index += 1
|
|
|
|
# we abort this feature if no useful node in it
|
|
if len(involved_first_tokens_pos) == 0:
|
|
break
|
|
|
|
if end_pos >= len(all_token_ids_seq):
|
|
flag = True
|
|
# which means we need to pad in this feature
|
|
current_len = len(splited_token_ids_seq)
|
|
splited_token_ids_seq += [tokenizer.pad_token_id] * (max_length - current_len)
|
|
splited_xpath_tags_seq += [padded_xpath_tags_seq] * (max_length - current_len)
|
|
splited_xpath_subs_seq += [padded_xpath_subs_seq] * (max_length - current_len)
|
|
splited_labels_seq += [-100] * (max_length - current_len)
|
|
attention_mask = [1] * current_len + [0] * (max_length - current_len)
|
|
|
|
else:
|
|
# no need to pad, the splited seq is exactly with the length `max_length`
|
|
assert len(splited_token_ids_seq) == max_length
|
|
attention_mask = [1] * max_length
|
|
|
|
features.append(
|
|
SwdeFeature(
|
|
html_path=html_path,
|
|
input_ids=splited_token_ids_seq,
|
|
token_type_ids=token_type_ids,
|
|
attention_mask=attention_mask,
|
|
xpath_tags_seq=splited_xpath_tags_seq,
|
|
xpath_subs_seq=splited_xpath_subs_seq,
|
|
labels=splited_labels_seq,
|
|
involved_first_tokens_pos=involved_first_tokens_pos,
|
|
involved_first_tokens_xpaths=involved_first_tokens_xpaths,
|
|
involved_first_tokens_types=involved_first_tokens_types,
|
|
involved_first_tokens_text=involved_first_tokens_text,
|
|
)
|
|
)
|
|
|
|
start_pos = end_pos - doc_stride
|
|
|
|
if flag:
|
|
break
|
|
|
|
return features
|