def create_input_data(config, seed=None):
if seed is not None:
np.random.seed(seed)
rand_mask_idx_list = create_bigbird_rand_mask_idx_list(
config["num_layers"], config["seq_len"], config["seq_len"],
config["nhead"], config["block_size"], config["window_size"],
config["num_global_blocks"], config["num_rand_blocks"], config["seed"])
input_ids = np.random.randint(
low=0,
high=config['vocab_size'],
size=(config["batch_size"], config["seq_len"]))
num_to_predict = int(config["seq_len"] * 0.15)
masked_lm_positions = np.random.choice(
config["seq_len"], (config["batch_size"], num_to_predict),
replace=False)
masked_lm_positions = np.sort(masked_lm_positions)
pred_padding_len = config["seq_len"] - num_to_predict
temp_masked_lm_positions = np.full(
masked_lm_positions.size, 0, dtype=np.int32)
mask_token_num = 0
for i, x in enumerate(masked_lm_positions):
for j, pos in enumerate(x):
temp_masked_lm_positions[mask_token_num] = i * config[
"seq_len"] + pos
mask_token_num += 1
masked_lm_positions = temp_masked_lm_positions
return rand_mask_idx_list, input_ids, masked_lm_positions
def _collate_data(data, stack_fn=Stack()):
# Data Fields: input_ids, segment_ids, masked_lm_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels
num_fields = len(data[0])
out = [None] * num_fields
for i in [0, 1, 5]:
out[i] = stack_fn([x[i] for x in data])
batch_size, seq_length = out[0].shape
size = num_mask = sum(len(x[2]) for x in data)
out[2] = np.full(size, 0, dtype=np.int32)
# masked_lm_labels
out[3] = np.full([size, 1], -1, dtype=np.int64)
# masked weight
out[4] = np.full([size], 0, dtype="float32")
# # Organize as a 1D tensor for gather or use gather_nd
mask_token_num = 0
for i, x in enumerate(data):
for j, pos in enumerate(x[2]):
out[2][mask_token_num] = i * seq_length + pos
out[3][mask_token_num] = x[3][j]
out[4][mask_token_num] = x[4][j]
mask_token_num += 1
out.append(np.asarray([mask_token_num], dtype=np.float32))
seq_len = len(out[0][0])
rand_mask_idx_list = create_bigbird_rand_mask_idx_list(
config["num_layers"], seq_len, seq_len, config["nhead"],
config["block_size"], config["window_size"],
config["num_global_blocks"], config["num_rand_blocks"],
config["seed"])
out.extend(rand_mask_idx_list)
return out
def _collate_data(data, stack_fn=Stack()):
num_fields = len(data[0])
out = [None] * num_fields
out[0] = stack_fn([_tokenize(x['text']) for x in data])
out[1] = stack_fn([x['label'] for x in data])
seq_len = len(out[0][0])
# Construct the random attention mask for the random attention
rand_mask_idx_list = create_bigbird_rand_mask_idx_list(
config["num_layers"], seq_len, seq_len, config["nhead"],
config["block_size"], config["window_size"],
config["num_global_blocks"], config["num_rand_blocks"],
config["seed"])
out.extend(rand_mask_idx_list)
return out
def collect_data(samples, dataset, config):
stack_fn = Stack(dtype="int64" if dataset.label_list else "float32")
stack_fn1 = Stack()
num_fields = len(samples[0])
out = [None] * num_fields
out[0] = stack_fn1([x[0] for x in samples]) # input_ids
out[1] = stack_fn1([x[1] for x in samples]) # token_type_ids
if num_fields >= 2:
out[2] = stack_fn(x[2] for x in samples) # labels
seq_len = len(out[0][0])
# Construct the random attention mask for the random attention
rand_mask_idx_list = create_bigbird_rand_mask_idx_list(
config["num_layers"], seq_len, seq_len, config["nhead"],
config["block_size"], config["window_size"],
config["num_global_blocks"], config["num_rand_blocks"], config["seed"])
out.extend(rand_mask_idx_list)
return out