chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,608 @@
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
import glob
|
||||
import re
|
||||
|
||||
import torch
|
||||
import tqdm
|
||||
import torch.utils.data
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Seq2seqDatasetForBert(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self, features, max_source_len, max_target_len,
|
||||
vocab_size, cls_id, sep_id, pad_id, mask_id,
|
||||
random_prob, keep_prob, offset, num_training_instances,
|
||||
span_len=1, span_prob=1.0):
|
||||
self.features = features
|
||||
self.max_source_len = max_source_len
|
||||
self.max_target_len = max_target_len
|
||||
self.offset = offset
|
||||
if offset > 0:
|
||||
logger.info(" **** Set offset %d in Seq2seqDatasetForBert **** ", offset)
|
||||
self.cls_id = cls_id
|
||||
self.sep_id = sep_id
|
||||
self.pad_id = pad_id
|
||||
self.random_prob = random_prob
|
||||
self.keep_prob = keep_prob
|
||||
self.mask_id = mask_id
|
||||
self.vocab_size = vocab_size
|
||||
self.num_training_instances = num_training_instances
|
||||
self.span_len = span_len
|
||||
self.span_prob = span_prob
|
||||
|
||||
def __len__(self):
|
||||
return int(self.num_training_instances)
|
||||
|
||||
def __trunk(self, ids, max_len):
|
||||
if len(ids) > max_len - 1:
|
||||
ids = ids[:max_len - 1]
|
||||
ids = ids + [self.sep_id]
|
||||
return ids
|
||||
|
||||
def __pad(self, ids, max_len):
|
||||
if len(ids) < max_len:
|
||||
return ids + [self.pad_id] * (max_len - len(ids))
|
||||
else:
|
||||
assert len(ids) == max_len
|
||||
return ids
|
||||
|
||||
def __getitem__(self, idx):
|
||||
idx = (self.offset + idx) % len(self.features)
|
||||
feature = self.features[idx]
|
||||
source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len)
|
||||
target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
|
||||
pseudo_ids = []
|
||||
for tk_id in target_ids:
|
||||
p = random.random()
|
||||
if p < self.keep_prob:
|
||||
pseudo_ids.append(tk_id)
|
||||
elif p < self.keep_prob + self.random_prob:
|
||||
pseudo_ids.append(random.randint(0, self.vocab_size - 1))
|
||||
else:
|
||||
pseudo_ids.append(self.mask_id)
|
||||
|
||||
num_source_tokens = len(source_ids)
|
||||
num_target_tokens = len(target_ids)
|
||||
|
||||
source_ids = self.__pad(source_ids, self.max_source_len)
|
||||
target_ids = self.__pad(target_ids, self.max_target_len)
|
||||
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)
|
||||
|
||||
if self.span_len > 1:
|
||||
span_ids = []
|
||||
span_id = 1
|
||||
while len(span_ids) < num_target_tokens:
|
||||
p = random.random()
|
||||
if p < self.span_prob:
|
||||
span_len = random.randint(2, self.span_len)
|
||||
span_len = min(span_len, num_target_tokens - len(span_ids))
|
||||
else:
|
||||
span_len = 1
|
||||
span_ids.extend([span_id] * span_len)
|
||||
span_id += 1
|
||||
span_ids = self.__pad(span_ids, self.max_target_len)
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids
|
||||
else:
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens
|
||||
|
||||
|
||||
# DONE: finish this!!! the 2D input id settings.
|
||||
class Seq2seqDatasetForLayoutlm(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self, features, max_source_len, max_target_len,
|
||||
vocab_size, cls_id, sep_id, pad_id, mask_id,
|
||||
random_prob, keep_prob, offset, num_training_instances, layout_flag=True,
|
||||
span_len=1, span_prob=1.0):
|
||||
|
||||
self.layout_flag = layout_flag
|
||||
|
||||
self.features = features
|
||||
self.max_source_len = max_source_len
|
||||
self.max_target_len = max_target_len
|
||||
self.offset = offset
|
||||
if offset > 0:
|
||||
logger.info(" **** Set offset %d in Seq2seqDatasetForBert **** ", offset)
|
||||
self.cls_id = cls_id
|
||||
self.sep_id = sep_id
|
||||
self.pad_id = pad_id
|
||||
self.random_prob = random_prob
|
||||
self.keep_prob = keep_prob
|
||||
self.mask_id = mask_id
|
||||
self.vocab_size = vocab_size
|
||||
self.num_training_instances = num_training_instances
|
||||
self.span_len = span_len
|
||||
self.span_prob = span_prob
|
||||
|
||||
self.index_sp_id = 0
|
||||
|
||||
def __len__(self):
|
||||
return int(self.num_training_instances)
|
||||
|
||||
def __clip_index(self, ids):
|
||||
replace_value = 0
|
||||
for i in range(len(ids)):
|
||||
if ids[i] > self.max_source_len - 1:
|
||||
ids[i] = replace_value
|
||||
return ids
|
||||
|
||||
def __trunk(self, ids, max_len, simple=False, value=None):
|
||||
trunk_value = value if value is not None else self.sep_id
|
||||
if len(ids) > max_len - 1:
|
||||
ids = ids[:max_len - 1]
|
||||
if simple:
|
||||
ids = ids + [trunk_value]
|
||||
else:
|
||||
ids = ids + [[trunk_value, 1000, 1000, 1000, 1000]]
|
||||
return ids
|
||||
|
||||
def __pad(self, ids, max_len, simple=False, value=None):
|
||||
pad_value = value if value is not None else self.pad_id
|
||||
if len(ids) < max_len:
|
||||
if simple:
|
||||
return ids + [pad_value] * (max_len - len(ids))
|
||||
else:
|
||||
return ids + [[pad_value, 0, 0, 0, 0]] * (max_len - len(ids))
|
||||
else:
|
||||
assert len(ids) == max_len
|
||||
return ids
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if self.layout_flag:
|
||||
return self.__getitem_layout__(idx)
|
||||
else:
|
||||
return self.__getitem_bert__(idx)
|
||||
|
||||
def __getitem_bert__(self, idx):
|
||||
idx = (self.offset + idx) % len(self.features)
|
||||
feature = self.features[idx]
|
||||
source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len, simple=True)
|
||||
target_ids = self.__trunk(feature["target_ids"], self.max_target_len, simple=True)
|
||||
target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)
|
||||
|
||||
pseudo_ids = []
|
||||
for tk_id in target_ids:
|
||||
p = random.random()
|
||||
if p < self.keep_prob:
|
||||
pseudo_ids.append(tk_id)
|
||||
elif p < self.keep_prob + self.random_prob:
|
||||
pseudo_ids.append(random.randint(0, self.vocab_size - 1))
|
||||
else:
|
||||
pseudo_ids.append(self.mask_id)
|
||||
|
||||
num_source_tokens = len(source_ids)
|
||||
num_target_tokens = len(target_ids)
|
||||
|
||||
source_ids = self.__pad(source_ids, self.max_source_len, simple=True)
|
||||
target_ids = self.__pad(target_ids, self.max_target_len, simple=True)
|
||||
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len, simple=True)
|
||||
target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
|
||||
target_index = self.__clip_index(target_index)
|
||||
|
||||
if self.span_len > 1:
|
||||
span_ids = []
|
||||
span_id = 1
|
||||
while len(span_ids) < num_target_tokens:
|
||||
p = random.random()
|
||||
if p < self.span_prob:
|
||||
span_len = random.randint(2, self.span_len)
|
||||
span_len = min(span_len, num_target_tokens - len(span_ids))
|
||||
else:
|
||||
span_len = 1
|
||||
span_ids.extend([span_id] * span_len)
|
||||
span_id += 1
|
||||
span_ids = self.__pad(span_ids, self.max_target_len)
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
|
||||
else:
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index
|
||||
|
||||
def __getitem_layout__(self, idx):
|
||||
# TODO: how to initialize the random and masked tokens' pos emb
|
||||
# Simple Solution: only mask the text
|
||||
idx = (self.offset + idx) % len(self.features)
|
||||
feature = self.features[idx]
|
||||
source_ids = self.__trunk([[self.cls_id, 0, 0, 0, 0]] + feature["source_ids"], self.max_source_len)
|
||||
target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
|
||||
target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)
|
||||
|
||||
pseudo_ids = []
|
||||
for tk_id in target_ids:
|
||||
p = random.random()
|
||||
if p < self.keep_prob:
|
||||
pseudo_ids.append(tk_id)
|
||||
elif p < self.keep_prob + self.random_prob:
|
||||
pseudo_ids.append([random.randint(0, self.vocab_size - 1)] + [0, 0, 0, 0]) # tk_id[1:])
|
||||
else:
|
||||
pseudo_ids.append([self.mask_id] + [0, 0, 0, 0]) # tk_id[1:])
|
||||
|
||||
num_source_tokens = len(source_ids)
|
||||
num_target_tokens = len(target_ids)
|
||||
|
||||
source_ids = self.__pad(source_ids, self.max_source_len)
|
||||
target_ids = self.__pad(target_ids, self.max_target_len)
|
||||
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)
|
||||
target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
|
||||
target_index = self.__clip_index(target_index)
|
||||
|
||||
if self.span_len > 1:
|
||||
span_ids = []
|
||||
span_id = 1
|
||||
while len(span_ids) < num_target_tokens:
|
||||
p = random.random()
|
||||
if p < self.span_prob:
|
||||
span_len = random.randint(2, self.span_len)
|
||||
span_len = min(span_len, num_target_tokens - len(span_ids))
|
||||
else:
|
||||
span_len = 1
|
||||
span_ids.extend([span_id] * span_len)
|
||||
span_id += 1
|
||||
span_ids = self.__pad(span_ids, self.max_target_len)
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
|
||||
else:
|
||||
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index
|
||||
|
||||
|
||||
def batch_list_to_batch_tensors(batch):
|
||||
batch_tensors = []
|
||||
for x in zip(*batch):
|
||||
if isinstance(x[0], torch.Tensor):
|
||||
batch_tensors.append(torch.stack(x))
|
||||
else:
|
||||
batch_tensors.append(torch.tensor(x, dtype=torch.long))
|
||||
return batch_tensors
|
||||
|
||||
|
||||
def get_max_epoch_model(output_dir):
|
||||
fn_model_list = glob.glob(os.path.join(output_dir, "model.*.bin"))
|
||||
fn_optim_list = glob.glob(os.path.join(output_dir, "optim.*.bin"))
|
||||
if (not fn_model_list) or (not fn_optim_list):
|
||||
return None
|
||||
os.path.basename(output_dir)
|
||||
both_set = set([int(os.path.basename(fn).split('.')[1]) for fn in fn_model_list]
|
||||
) & set([int(os.path.basename(fn).split('.')[1]) for fn in fn_optim_list])
|
||||
if both_set:
|
||||
return max(both_set)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def load_and_cache_examples(
|
||||
example_file, tokenizer, local_rank, cached_features_file, shuffle=True):
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank not in [-1, 0]:
|
||||
torch.distributed.barrier()
|
||||
|
||||
if cached_features_file is not None and os.path.exists(cached_features_file):
|
||||
logger.info("Loading features from cached file %s", cached_features_file)
|
||||
features = torch.load(cached_features_file)
|
||||
else:
|
||||
logger.info("Creating features from dataset file at %s", example_file)
|
||||
|
||||
examples = []
|
||||
with open(example_file, mode="r", encoding="utf-8") as reader:
|
||||
for i, line in enumerate(reader):
|
||||
if i == 100:
|
||||
break
|
||||
examples.append(json.loads(line))
|
||||
features = []
|
||||
|
||||
for example in tqdm.tqdm(examples):
|
||||
if isinstance(example["src"], list):
|
||||
source_tokens = example["src"]
|
||||
target_tokens = example["tgt"]
|
||||
else:
|
||||
source_tokens = tokenizer.tokenize(example["src"])
|
||||
target_tokens = tokenizer.tokenize(example["tgt"])
|
||||
features.append({
|
||||
"source_ids": tokenizer.convert_tokens_to_ids(source_tokens),
|
||||
"target_ids": tokenizer.convert_tokens_to_ids(target_tokens),
|
||||
})
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(features)
|
||||
|
||||
if local_rank in [-1, 0] and cached_features_file is not None:
|
||||
logger.info("Saving features into cached file %s", cached_features_file)
|
||||
torch.save(features, cached_features_file)
|
||||
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank == 0:
|
||||
torch.distributed.barrier()
|
||||
|
||||
return features
|
||||
|
||||
|
||||
def load_and_cache_line_order_examples(
|
||||
example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
|
||||
layout_flag=True, shuffle=True,
|
||||
src_shuffle_rate=0,
|
||||
file_info_flag=False,
|
||||
):
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank not in [-1, 0]:
|
||||
torch.distributed.barrier()
|
||||
|
||||
if cached_features_file is not None and os.path.exists(cached_features_file) and False:
|
||||
logger.info("Loading features from cached file %s", cached_features_file)
|
||||
features = torch.load(cached_features_file)
|
||||
else:
|
||||
logger.info("Creating features from dataset at %s", example_path)
|
||||
|
||||
examples = []
|
||||
|
||||
with open(example_path, 'r') as layout_reader:
|
||||
logger.info(f'Start loading {example_path}')
|
||||
for i, line in enumerate(layout_reader):
|
||||
examples.append(json.loads(line))
|
||||
|
||||
features = []
|
||||
|
||||
for layout in tqdm.tqdm(examples):
|
||||
bleu = layout['bleu']
|
||||
|
||||
if random.random() < src_shuffle_rate:
|
||||
# print('Random!!!')
|
||||
# DONE: the random src! here has bug! index also need shuffle
|
||||
src_layout = layout['src']
|
||||
tgt_index = layout['tgt_index']
|
||||
|
||||
source_length = len(src_layout)
|
||||
shuffle_index = list(range(source_length))
|
||||
random.shuffle(shuffle_index)
|
||||
|
||||
shuffle_layout = ['' for _ in range(source_length)]
|
||||
for i, j in enumerate(shuffle_index):
|
||||
# NOTE: map i-th token to j-th token
|
||||
shuffle_layout[j] = src_layout[i]
|
||||
|
||||
shuffle_target_index = [shuffle_index[i] for i in tgt_index]
|
||||
|
||||
layout['tgt_index'] = shuffle_target_index
|
||||
layout['src'] = shuffle_layout
|
||||
|
||||
mask = tokenizer.mask_token_id
|
||||
src_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + src_layout for tmp_i, src_layout in enumerate(layout['src'])]
|
||||
tgt_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + tgt_layout for tmp_i, tgt_layout in enumerate(layout['tgt'])]
|
||||
tgt_index = layout['tgt_index']
|
||||
|
||||
feature = {
|
||||
"source_ids": src_ids,
|
||||
"target_ids": tgt_ids,
|
||||
"target_index": tgt_index,
|
||||
'bleu': bleu
|
||||
}
|
||||
|
||||
if file_info_flag:
|
||||
file_info = {'original_filename': layout['filename'], 'filename': layout['filename'],
|
||||
'page_idx': 0}
|
||||
feature['file_info'] = file_info
|
||||
|
||||
features.append(feature)
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(features)
|
||||
|
||||
if local_rank in [-1, 0] and cached_features_file is not None:
|
||||
logger.info("Saving features into cached file %s", cached_features_file)
|
||||
torch.save(features, cached_features_file)
|
||||
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank == 0:
|
||||
torch.distributed.barrier()
|
||||
|
||||
return features
|
||||
|
||||
|
||||
def load_and_cache_layoutlm_examples(
|
||||
example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
|
||||
layout_flag=True, shuffle=True,
|
||||
src_shuffle_rate=0,
|
||||
file_info_flag=False
|
||||
):
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank not in [-1, 0]:
|
||||
torch.distributed.barrier()
|
||||
|
||||
if cached_features_file is not None and os.path.exists(cached_features_file):
|
||||
logger.info("Loading features from cached file %s", cached_features_file)
|
||||
features = torch.load(cached_features_file)
|
||||
else:
|
||||
logger.info("Creating features from dataset at %s", example_path)
|
||||
|
||||
examples = []
|
||||
|
||||
if os.path.isdir(example_path):
|
||||
text_files = glob.glob(f'{example_path}/*text*.json')
|
||||
layout_files = [re.sub('text|txt', 'layout', x, 1) for x in text_files]
|
||||
else:
|
||||
text_files = [example_path]
|
||||
layout_files = [re.sub('text|txt', 'layout', example_path, 1)]
|
||||
for text_file, layout_file in zip(text_files, layout_files):
|
||||
with open(text_file, mode='r', encoding='utf-8') as text_reader, \
|
||||
open(layout_file, mode='r', encoding='utf-8') as layout_reader:
|
||||
logger.info(f'Start loading {text_file}')
|
||||
for i, (text_line, layout_line) in enumerate(zip(text_reader, layout_reader)):
|
||||
if (i + 1) % 10000 == 0:
|
||||
logger.info(f'{i + 1} lines ...')
|
||||
examples.append((json.loads(text_line), json.loads(layout_line)))
|
||||
|
||||
features = []
|
||||
|
||||
def tokenize_text_and_layout_src(_text, _layout, _layout_flag):
|
||||
ret = []
|
||||
index_split = {}
|
||||
words = _text.split()
|
||||
# note: (OLD) the index should start from 1: 0-the cls token in src
|
||||
# note: (NEW) we need to remove the src embedding's CLS SEP token so we can still start from 0
|
||||
# note: (NEWER) we need to at least one blank pos for ignore index in loss function (we use sep's index)
|
||||
# NOTE: (NEWER-ER) 1 for all padding tgt index
|
||||
new_token_index = 1 # first ordinary index
|
||||
for i, (word, box) in enumerate(zip(words, _layout)):
|
||||
|
||||
if (not box[2] >= box[0]) or (not box[3] >= box[1]):
|
||||
continue
|
||||
|
||||
tokens = tokenizer.tokenize(word)
|
||||
tokens = tokenizer.convert_tokens_to_ids(tokens)
|
||||
new_token_ids = []
|
||||
for token in tokens:
|
||||
if _layout_flag:
|
||||
ret.append([token] + box)
|
||||
else:
|
||||
ret.append(token)
|
||||
new_token_ids.append(new_token_index)
|
||||
new_token_index += 1
|
||||
index_split[i] = new_token_ids
|
||||
|
||||
return ret, index_split
|
||||
|
||||
def tokenize_text_and_layout_tgt(_text, _layout, _index, _index_split, _layout_flag):
|
||||
ret = []
|
||||
ret_index = []
|
||||
words = _text.split()
|
||||
for word, box, i in zip(words, _layout, _index):
|
||||
|
||||
if (not box[2] >= box[0]) or (not box[3] >= box[1]):
|
||||
continue
|
||||
|
||||
tokens = tokenizer.tokenize(word)
|
||||
tokens = tokenizer.convert_tokens_to_ids(tokens)
|
||||
for token, ii in zip(tokens, _index_split[i]):
|
||||
if _layout_flag:
|
||||
ret.append([token] + box)
|
||||
else:
|
||||
ret.append(token)
|
||||
ii = min(ii, max_src_length - 1)
|
||||
ret_index.append(ii)
|
||||
return ret, ret_index
|
||||
|
||||
for text, layout in tqdm.tqdm(examples):
|
||||
if 'bleu' in text:
|
||||
bleu = text['bleu']
|
||||
else:
|
||||
bleu = 0
|
||||
|
||||
if random.random() < src_shuffle_rate:
|
||||
# print('Random!!!')
|
||||
# DONE: the random src! here has bug! index also need shuffle
|
||||
src_text = text['src']
|
||||
src_layout = layout['src']
|
||||
tgt_index = text['tgt_index']
|
||||
|
||||
src_text = src_text.split()
|
||||
source_length = len(src_text)
|
||||
shuffle_index = list(range(source_length))
|
||||
random.shuffle(shuffle_index)
|
||||
|
||||
shuffle_text = ['' for _ in range(source_length)]
|
||||
shuffle_layout = ['' for _ in range(source_length)]
|
||||
for i, j in enumerate(shuffle_index):
|
||||
# NOTE: map i-th token to j-th token
|
||||
shuffle_text[j] = src_text[i]
|
||||
shuffle_layout[j] = src_layout[i]
|
||||
|
||||
shuffle_target_index = [shuffle_index[i] for i in tgt_index]
|
||||
|
||||
text['src'] = ' '.join(shuffle_text)
|
||||
text['tgt_index'] = shuffle_target_index
|
||||
layout['src'] = shuffle_layout
|
||||
|
||||
src_ids, src_index_split = tokenize_text_and_layout_src(text['src'], layout['src'],
|
||||
_layout_flag=layout_flag)
|
||||
tgt_ids, tgt_index = tokenize_text_and_layout_tgt(text['tgt'], layout['tgt'], text['tgt_index'],
|
||||
src_index_split, _layout_flag=layout_flag)
|
||||
|
||||
feature = {
|
||||
"source_ids": src_ids,
|
||||
"target_ids": tgt_ids,
|
||||
"target_index": tgt_index,
|
||||
'bleu': bleu
|
||||
}
|
||||
|
||||
if file_info_flag:
|
||||
file_info = {'original_filename': text['original_filename'], 'filename': text['filename'], 'page_idx': text['page_idx']}
|
||||
feature['file_info'] = file_info
|
||||
|
||||
features.append(feature)
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(features)
|
||||
|
||||
if local_rank in [-1, 0] and cached_features_file is not None:
|
||||
if not os.path.exists(os.path.dirname(cached_features_file)):
|
||||
os.makedirs(os.path.dirname(cached_features_file))
|
||||
logger.info("Saving features into cached file %s", cached_features_file)
|
||||
torch.save(features, cached_features_file)
|
||||
|
||||
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
||||
if local_rank == 0:
|
||||
torch.distributed.barrier()
|
||||
|
||||
return features
|
||||
|
||||
|
||||
def convert_src_layout_inputs_to_tokens(inputs, converter, max_src_length, layout_flag=True):
|
||||
ret = []
|
||||
if not layout_flag:
|
||||
for line in inputs:
|
||||
ret.append(converter(line["source_ids"])[: max_src_length])
|
||||
else:
|
||||
for line in inputs:
|
||||
raw_text_ids = [x[0] for x in line['source_ids']]
|
||||
raw_text = converter(raw_text_ids)
|
||||
new_line = [[t] + x[1:] for t, x in zip(raw_text, line['source_ids'])][: max_src_length]
|
||||
ret.append(new_line)
|
||||
return ret
|
||||
|
||||
|
||||
def convert_tgt_layout_inputs_to_tokens(inputs, converter, max_tgt_length, layout_flag=True):
|
||||
ret = []
|
||||
if not layout_flag:
|
||||
for line in inputs:
|
||||
ret.append(converter(line["target_ids"])[: max_tgt_length])
|
||||
else:
|
||||
for line in inputs:
|
||||
raw_text_ids = [x[0] for x in line['target_ids']]
|
||||
ret.append(converter(raw_text_ids)[: max_tgt_length])
|
||||
return ret
|
||||
|
||||
|
||||
def get_tokens_from_src_and_index(src, index, modifier=None):
|
||||
result = []
|
||||
for i in index:
|
||||
i = modifier(i)
|
||||
i = min(i, len(src) - 1)
|
||||
if isinstance(src[i], list):
|
||||
result.append(src[i][0])
|
||||
else:
|
||||
result.append(src[i])
|
||||
return result
|
||||
|
||||
|
||||
def get_layout_from_src_and_index(src, index, modifier=None):
|
||||
result = []
|
||||
s = set()
|
||||
for i in index:
|
||||
i = modifier(i)
|
||||
i = min(i, len(src) - 1)
|
||||
layout = src[i][1:]
|
||||
if repr(layout) not in s:
|
||||
result.append(layout)
|
||||
s.add(repr(layout))
|
||||
return result
|
||||
|
||||
|
||||
def get_everything_from_src_and_index(src, index, modifier=None):
|
||||
result = []
|
||||
for i in index:
|
||||
i = modifier(i)
|
||||
i = min(i, len(src) - 1)
|
||||
result.append(src[i])
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user