wehub-resource-sync
609 lines
23 KiB
Python
609 lines
23 KiB
Python
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
|
|
|