chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .block_pair_dataset import BlockPairDataset
|
||||
from .masked_lm_dataset import MaskedLMDataset
|
||||
from .masked_lm_dictionary import BertDictionary, MaskedLMDictionary
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BertDictionary",
|
||||
"BlockPairDataset",
|
||||
"MaskedLMDataset",
|
||||
"MaskedLMDictionary",
|
||||
]
|
||||
@@ -0,0 +1,311 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq.data import FairseqDataset
|
||||
|
||||
|
||||
class BlockPairDataset(FairseqDataset):
|
||||
"""Break a Dataset of tokens into sentence pair blocks for next sentence
|
||||
prediction as well as masked language model.
|
||||
|
||||
High-level logics are:
|
||||
1. break input tensor to tensor blocks
|
||||
2. pair the blocks with 50% next sentence and 50% random sentence
|
||||
3. return paired blocks as well as related segment labels
|
||||
|
||||
Args:
|
||||
dataset (~torch.utils.data.Dataset): dataset to break into blocks
|
||||
sizes: array of sentence lengths
|
||||
dictionary: dictionary for the task
|
||||
block_size: maximum block size
|
||||
break_mode: mode for breaking copurs into block pairs. currently we support
|
||||
2 modes
|
||||
doc: respect document boundaries and each part of the pair should belong to on document
|
||||
none: don't respect any boundary and cut tokens evenly
|
||||
short_seq_prob: probability for generating shorter block pairs
|
||||
doc_break_size: Size for empty line separating documents. Typically 1 if
|
||||
the sentences have eos, 0 otherwise.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset,
|
||||
dictionary,
|
||||
sizes,
|
||||
block_size,
|
||||
break_mode="doc",
|
||||
short_seq_prob=0.1,
|
||||
doc_break_size=1,
|
||||
):
|
||||
super().__init__()
|
||||
self.dataset = dataset
|
||||
self.pad = dictionary.pad()
|
||||
self.eos = dictionary.eos()
|
||||
self.cls = dictionary.cls()
|
||||
self.mask = dictionary.mask()
|
||||
self.sep = dictionary.sep()
|
||||
self.break_mode = break_mode
|
||||
self.dictionary = dictionary
|
||||
self.short_seq_prob = short_seq_prob
|
||||
self.block_indices = []
|
||||
|
||||
assert len(dataset) == len(sizes)
|
||||
|
||||
if break_mode == "doc":
|
||||
cur_doc = []
|
||||
for sent_id, sz in enumerate(sizes):
|
||||
assert doc_break_size == 0 or sz != 0, (
|
||||
"when doc_break_size is non-zero, we expect documents to be"
|
||||
"separated by a blank line with a single eos."
|
||||
)
|
||||
# empty line as document separator
|
||||
if sz == doc_break_size:
|
||||
if len(cur_doc) == 0:
|
||||
continue
|
||||
self.block_indices.append(cur_doc)
|
||||
cur_doc = []
|
||||
else:
|
||||
cur_doc.append(sent_id)
|
||||
max_num_tokens = block_size - 3 # Account for [CLS], [SEP], [SEP]
|
||||
self.sent_pairs = []
|
||||
self.sizes = []
|
||||
for doc_id, doc in enumerate(self.block_indices):
|
||||
self._generate_sentence_pair(doc, doc_id, max_num_tokens, sizes)
|
||||
elif break_mode is None or break_mode == "none":
|
||||
# each block should have half of the block size since we are constructing block pair
|
||||
sent_length = (block_size - 3) // 2
|
||||
total_len = sum(dataset.sizes)
|
||||
length = math.ceil(total_len / sent_length)
|
||||
|
||||
def block_at(i):
|
||||
start = i * sent_length
|
||||
end = min(start + sent_length, total_len)
|
||||
return (start, end)
|
||||
|
||||
sent_indices = np.array([block_at(i) for i in range(length)])
|
||||
sent_sizes = np.array([e - s for s, e in sent_indices])
|
||||
dataset_index = self._sent_to_dataset_index(sent_sizes)
|
||||
|
||||
# pair sentences
|
||||
self._pair_sentences(dataset_index)
|
||||
else:
|
||||
raise ValueError("Invalid break_mode: " + break_mode)
|
||||
|
||||
def _pair_sentences(self, dataset_index):
|
||||
"""
|
||||
Give a list of evenly cut blocks/sentences, pair these sentences with 50%
|
||||
consecutive sentences and 50% random sentences.
|
||||
This is used for none break mode
|
||||
"""
|
||||
# pair sentences
|
||||
for sent_id, sent in enumerate(dataset_index):
|
||||
next_sent_label = (
|
||||
1 if np.random.rand() > 0.5 and sent_id != len(dataset_index) - 1 else 0
|
||||
)
|
||||
if next_sent_label:
|
||||
next_sent = dataset_index[sent_id + 1]
|
||||
else:
|
||||
next_sent = dataset_index[
|
||||
self._skip_sampling(len(dataset_index), [sent_id, sent_id + 1])
|
||||
]
|
||||
self.sent_pairs.append((sent, next_sent, next_sent_label))
|
||||
|
||||
# The current blocks don't include the special tokens but the
|
||||
# sizes already account for this
|
||||
self.sizes.append(3 + sent[3] + next_sent[3])
|
||||
|
||||
def _sent_to_dataset_index(self, sent_sizes):
|
||||
"""
|
||||
Build index mapping block indices to the underlying dataset indices
|
||||
"""
|
||||
dataset_index = []
|
||||
ds_idx, ds_remaining = -1, 0
|
||||
for to_consume in sent_sizes:
|
||||
sent_size = to_consume
|
||||
if ds_remaining == 0:
|
||||
ds_idx += 1
|
||||
ds_remaining = sent_sizes[ds_idx]
|
||||
start_ds_idx = ds_idx
|
||||
start_offset = sent_sizes[ds_idx] - ds_remaining
|
||||
while to_consume > ds_remaining:
|
||||
to_consume -= ds_remaining
|
||||
ds_idx += 1
|
||||
ds_remaining = sent_sizes[ds_idx]
|
||||
ds_remaining -= to_consume
|
||||
dataset_index.append(
|
||||
(
|
||||
start_ds_idx, # starting index in dataset
|
||||
start_offset, # starting offset within starting index
|
||||
ds_idx, # ending index in dataset
|
||||
sent_size, # sentence length
|
||||
)
|
||||
)
|
||||
assert ds_remaining == 0
|
||||
assert ds_idx == len(self.dataset) - 1
|
||||
return dataset_index
|
||||
|
||||
def _generate_sentence_pair(self, doc, doc_id, max_num_tokens, sizes):
|
||||
"""
|
||||
Go through a single document and genrate sentence paris from it
|
||||
"""
|
||||
current_chunk = []
|
||||
current_length = 0
|
||||
curr = 0
|
||||
# To provide more randomness, we decrease target seq length for parts of
|
||||
# samples (10% by default). Note that max_num_tokens is the hard threshold
|
||||
# for batching and will never be changed.
|
||||
target_seq_length = max_num_tokens
|
||||
if np.random.random() < self.short_seq_prob:
|
||||
target_seq_length = np.random.randint(2, max_num_tokens)
|
||||
# loop through all sentences in document
|
||||
while curr < len(doc):
|
||||
sent_id = doc[curr]
|
||||
current_chunk.append(sent_id)
|
||||
current_length = sum(sizes[current_chunk])
|
||||
# split chunk and generate pair when exceed target_seq_length or
|
||||
# finish the loop
|
||||
if curr == len(doc) - 1 or current_length >= target_seq_length:
|
||||
# split the chunk into 2 parts
|
||||
a_end = 1
|
||||
if len(current_chunk) > 2:
|
||||
a_end = np.random.randint(1, len(current_chunk) - 1)
|
||||
sent_a = current_chunk[:a_end]
|
||||
len_a = sum(sizes[sent_a])
|
||||
# generate next sentence label, note that if there is only 1 sentence
|
||||
# in current chunk, label is always 0
|
||||
next_sent_label = (
|
||||
1 if np.random.rand() > 0.5 and len(current_chunk) != 1 else 0
|
||||
)
|
||||
if not next_sent_label:
|
||||
# if next sentence label is 0, sample sent_b from a random doc
|
||||
target_b_length = target_seq_length - len_a
|
||||
rand_doc_id = self._skip_sampling(len(self.block_indices), [doc_id])
|
||||
random_doc = self.block_indices[rand_doc_id]
|
||||
random_start = np.random.randint(0, len(random_doc))
|
||||
sent_b = []
|
||||
len_b = 0
|
||||
for j in range(random_start, len(random_doc)):
|
||||
sent_b.append(random_doc[j])
|
||||
len_b = sum(sizes[sent_b])
|
||||
if len_b >= target_b_length:
|
||||
break
|
||||
# return the second part of the chunk since it's not used
|
||||
num_unused_segments = len(current_chunk) - a_end
|
||||
curr -= num_unused_segments
|
||||
else:
|
||||
# if next sentence label is 1, use the second part of chunk as sent_B
|
||||
sent_b = current_chunk[a_end:]
|
||||
len_b = sum(sizes[sent_b])
|
||||
# currently sent_a and sent_B may be longer than max_num_tokens,
|
||||
# truncate them and return block idx and offsets for them
|
||||
sent_a, sent_b = self._truncate_sentences(
|
||||
sent_a, sent_b, max_num_tokens
|
||||
)
|
||||
self.sent_pairs.append((sent_a, sent_b, next_sent_label))
|
||||
self.sizes.append(3 + sent_a[3] + sent_b[3])
|
||||
current_chunk = []
|
||||
curr += 1
|
||||
|
||||
def _skip_sampling(self, total, skip_ids):
|
||||
"""
|
||||
Generate a random integer which is not in skip_ids. Sample range is [0, total)
|
||||
TODO: ids in skip_ids should be consecutive, we can extend it to more generic version later
|
||||
"""
|
||||
rand_id = np.random.randint(total - len(skip_ids))
|
||||
return rand_id if rand_id < min(skip_ids) else rand_id + len(skip_ids)
|
||||
|
||||
def _truncate_sentences(self, sent_a, sent_b, max_num_tokens):
|
||||
"""
|
||||
Trancate a pair of sentence to limit total length under max_num_tokens
|
||||
Logics:
|
||||
1. Truncate longer sentence
|
||||
2. Tokens to be truncated could be at the beginning or the end of the sentnce
|
||||
Returns:
|
||||
Truncated sentences represented by dataset idx
|
||||
"""
|
||||
len_a, len_b = sum(self.dataset.sizes[sent_a]), sum(self.dataset.sizes[sent_b])
|
||||
front_cut_a = front_cut_b = end_cut_a = end_cut_b = 0
|
||||
|
||||
while True:
|
||||
total_length = (
|
||||
len_a + len_b - front_cut_a - front_cut_b - end_cut_a - end_cut_b
|
||||
)
|
||||
if total_length <= max_num_tokens:
|
||||
break
|
||||
|
||||
if len_a - front_cut_a - end_cut_a > len_b - front_cut_b - end_cut_b:
|
||||
if np.random.rand() < 0.5:
|
||||
front_cut_a += 1
|
||||
else:
|
||||
end_cut_a += 1
|
||||
else:
|
||||
if np.random.rand() < 0.5:
|
||||
front_cut_b += 1
|
||||
else:
|
||||
end_cut_b += 1
|
||||
|
||||
# calculate ds indices as well as offsets and return
|
||||
truncated_sent_a = self._cut_sentence(sent_a, front_cut_a, end_cut_a)
|
||||
truncated_sent_b = self._cut_sentence(sent_b, front_cut_b, end_cut_b)
|
||||
return truncated_sent_a, truncated_sent_b
|
||||
|
||||
def _cut_sentence(self, sent, front_cut, end_cut):
|
||||
"""
|
||||
Cut a sentence based on the numbers of tokens to be cut from beginning and end
|
||||
Represent the sentence as dataset idx and return
|
||||
"""
|
||||
start_ds_idx, end_ds_idx, offset = sent[0], sent[-1], 0
|
||||
target_len = sum(self.dataset.sizes[sent]) - front_cut - end_cut
|
||||
while front_cut > 0:
|
||||
if self.dataset.sizes[start_ds_idx] > front_cut:
|
||||
offset += front_cut
|
||||
break
|
||||
else:
|
||||
front_cut -= self.dataset.sizes[start_ds_idx]
|
||||
start_ds_idx += 1
|
||||
while end_cut > 0:
|
||||
if self.dataset.sizes[end_ds_idx] > end_cut:
|
||||
break
|
||||
else:
|
||||
end_cut -= self.dataset.sizes[end_ds_idx]
|
||||
end_ds_idx -= 1
|
||||
return start_ds_idx, offset, end_ds_idx, target_len
|
||||
|
||||
def _fetch_block(self, start_ds_idx, offset, end_ds_idx, length):
|
||||
"""
|
||||
Fetch a block of tokens based on its dataset idx
|
||||
"""
|
||||
buffer = torch.cat(
|
||||
[self.dataset[idx] for idx in range(start_ds_idx, end_ds_idx + 1)]
|
||||
)
|
||||
s, e = offset, offset + length
|
||||
return buffer[s:e]
|
||||
|
||||
def __getitem__(self, index):
|
||||
block1, block2, next_sent_label = self.sent_pairs[index]
|
||||
block1 = self._fetch_block(*block1)
|
||||
block2 = self._fetch_block(*block2)
|
||||
return block1, block2, next_sent_label
|
||||
|
||||
def __len__(self):
|
||||
return len(self.sizes)
|
||||
|
||||
@property
|
||||
def supports_prefetch(self):
|
||||
return getattr(self.dataset, "supports_prefetch", False)
|
||||
|
||||
def prefetch(self, indices):
|
||||
prefetch_idx = set()
|
||||
for index in indices:
|
||||
for block1, block2, _ in [self.sent_pairs[index]]:
|
||||
for ds_idx in range(block1[0], block1[2] + 1):
|
||||
prefetch_idx.add(ds_idx)
|
||||
for ds_idx in range(block2[0], block2[2] + 1):
|
||||
prefetch_idx.add(ds_idx)
|
||||
self.dataset.prefetch(prefetch_idx)
|
||||
@@ -0,0 +1,303 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq.data import Dictionary, FairseqDataset, data_utils
|
||||
from fairseq.data.concat_dataset import ConcatDataset
|
||||
from fairseq.data.legacy.block_pair_dataset import BlockPairDataset
|
||||
from fairseq.data.token_block_dataset import TokenBlockDataset
|
||||
|
||||
|
||||
class MaskedLMDataset(FairseqDataset):
|
||||
"""
|
||||
A wrapper Dataset for masked language modelling. The dataset
|
||||
wraps around TokenBlockDataset or BlockedPairDataset and creates a batch
|
||||
where the input blocks are masked according to the specified masking
|
||||
probability. Additionally the batch can also contain sentence level targets
|
||||
if this is specified.
|
||||
|
||||
Args:
|
||||
dataset: Dataset which generates blocks of data. Only BlockPairDataset
|
||||
and TokenBlockDataset are supported.
|
||||
sizes: Sentence lengths
|
||||
vocab: Dictionary with the vocabulary and special tokens.
|
||||
pad_idx: Id of padding token in dictionary
|
||||
mask_idx: Id of mask token in dictionary
|
||||
classif_token_idx: Id of classification token in dictionary. This is the
|
||||
token associated with the sentence embedding (Eg: CLS for BERT)
|
||||
sep_token_idx: Id of separator token in dictionary
|
||||
(Eg: SEP in BERT)
|
||||
seed: Seed for random number generator for reproducibility.
|
||||
shuffle: Shuffle the elements before batching.
|
||||
has_pairs: Specifies whether the underlying dataset
|
||||
generates a pair of blocks along with a sentence_target or not.
|
||||
Setting it to True assumes that the underlying dataset generates a
|
||||
label for the pair of sentences which is surfaced as
|
||||
sentence_target. The default value assumes a single block with no
|
||||
sentence target.
|
||||
segment_id: An optional segment id for filling in the segment labels
|
||||
when we are in the single block setting (Eg: XLM). Default is 0.
|
||||
masking_ratio: specifies what percentage of the blocks should be masked.
|
||||
masking_prob: specifies the probability of a given token being
|
||||
replaced with the "MASK" token.
|
||||
random_token_prob: specifies the probability of a given token being
|
||||
replaced by a random token from the vocabulary.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset: FairseqDataset,
|
||||
sizes: np.ndarray,
|
||||
vocab: Dictionary,
|
||||
pad_idx: int,
|
||||
mask_idx: int,
|
||||
classif_token_idx: int,
|
||||
sep_token_idx: int,
|
||||
seed: int = 1,
|
||||
shuffle: bool = True,
|
||||
has_pairs: bool = True,
|
||||
segment_id: int = 0,
|
||||
masking_ratio: float = 0.15,
|
||||
masking_prob: float = 0.8,
|
||||
random_token_prob: float = 0.1,
|
||||
):
|
||||
# Make sure the input datasets are the ones supported
|
||||
assert (
|
||||
isinstance(dataset, TokenBlockDataset)
|
||||
or isinstance(dataset, BlockPairDataset)
|
||||
or isinstance(dataset, ConcatDataset)
|
||||
), (
|
||||
"MaskedLMDataset only wraps TokenBlockDataset or BlockPairDataset or "
|
||||
"ConcatDataset"
|
||||
)
|
||||
|
||||
self.dataset = dataset
|
||||
self.sizes = np.array(sizes)
|
||||
self.vocab = vocab
|
||||
self.pad_idx = pad_idx
|
||||
self.mask_idx = mask_idx
|
||||
self.classif_token_idx = classif_token_idx
|
||||
self.sep_token_idx = sep_token_idx
|
||||
self.shuffle = shuffle
|
||||
self.seed = seed
|
||||
self.has_pairs = has_pairs
|
||||
self.segment_id = segment_id
|
||||
self.masking_ratio = masking_ratio
|
||||
self.masking_prob = masking_prob
|
||||
self.random_token_prob = random_token_prob
|
||||
|
||||
# If we have only one block then sizes needs to be updated to include
|
||||
# the classification token
|
||||
if not has_pairs:
|
||||
self.sizes = self.sizes + 1
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
# if has_pairs, then expect 2 blocks and a sentence target
|
||||
if self.has_pairs:
|
||||
(block_one, block_two, sentence_target) = self.dataset[index]
|
||||
else:
|
||||
block_one = self.dataset[index]
|
||||
|
||||
return {
|
||||
"id": index,
|
||||
"block_one": block_one,
|
||||
"block_two": block_two if self.has_pairs else None,
|
||||
"sentence_target": sentence_target if self.has_pairs else None,
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def _mask_block(
|
||||
self,
|
||||
sentence: np.ndarray,
|
||||
mask_idx: int,
|
||||
pad_idx: int,
|
||||
dictionary_token_range: Tuple,
|
||||
):
|
||||
"""
|
||||
Mask tokens for Masked Language Model training
|
||||
Samples mask_ratio tokens that will be predicted by LM.
|
||||
|
||||
Note:This function may not be efficient enough since we had multiple
|
||||
conversions between np and torch, we can replace them with torch
|
||||
operators later.
|
||||
|
||||
Args:
|
||||
sentence: 1d tensor to be masked
|
||||
mask_idx: index to use for masking the sentence
|
||||
pad_idx: index to use for masking the target for tokens we aren't
|
||||
predicting
|
||||
dictionary_token_range: range of indices in dictionary which can
|
||||
be used for random word replacement
|
||||
(e.g. without special characters)
|
||||
Return:
|
||||
masked_sent: masked sentence
|
||||
target: target with words which we are not predicting replaced
|
||||
by pad_idx
|
||||
"""
|
||||
masked_sent = np.copy(sentence)
|
||||
sent_length = len(sentence)
|
||||
mask_num = math.ceil(sent_length * self.masking_ratio)
|
||||
mask = np.random.choice(sent_length, mask_num, replace=False)
|
||||
target = np.copy(sentence)
|
||||
|
||||
for i in range(sent_length):
|
||||
if i in mask:
|
||||
rand = np.random.random()
|
||||
|
||||
# replace with mask if probability is less than masking_prob
|
||||
# (Eg: 0.8)
|
||||
if rand < self.masking_prob:
|
||||
masked_sent[i] = mask_idx
|
||||
|
||||
# replace with random token if probability is less than
|
||||
# masking_prob + random_token_prob (Eg: 0.9)
|
||||
elif rand < (self.masking_prob + self.random_token_prob):
|
||||
# sample random token from dictionary
|
||||
masked_sent[i] = np.random.randint(
|
||||
dictionary_token_range[0], dictionary_token_range[1]
|
||||
)
|
||||
else:
|
||||
target[i] = pad_idx
|
||||
|
||||
return masked_sent, target
|
||||
|
||||
def _collate(self, samples: List[Dict], pad_idx: int, eos_idx: int):
|
||||
"""
|
||||
Does the heavy lifting for creating a batch from the input list of
|
||||
examples. The logic is as follows:
|
||||
1. Mask the input blocks. In case has_pair is True then we have 2
|
||||
blocks to mask.
|
||||
2. Prepend the first masked block tensor with the special token
|
||||
used as sentence embedding. Eg: CLS in BERT. This happens
|
||||
irrespective of the value of has_pair.
|
||||
3. If has_pair is True, then append the first masked block with the
|
||||
special separator token (eg: SEP for BERT) and compute segment
|
||||
label accordingly. In this case, also append the second masked
|
||||
block with this special separator token and compute its segment
|
||||
label.
|
||||
4. For the targets tensor, prepend and append with padding index
|
||||
accordingly.
|
||||
5. Concatenate all tensors.
|
||||
"""
|
||||
if len(samples) == 0:
|
||||
return {}
|
||||
# To ensure determinism, we reset the state of the PRNG after every
|
||||
# batch based on the seed and the first id of the batch. This ensures
|
||||
# that across epochs we get the same mask for the same example. This
|
||||
# is needed for reproducibility and is how BERT does masking
|
||||
# TODO: Can we add deteminism without this constraint?
|
||||
with data_utils.numpy_seed(self.seed + samples[0]["id"]):
|
||||
for s in samples:
|
||||
|
||||
# token range is needed for replacing with random token during
|
||||
# masking
|
||||
token_range = (self.vocab.nspecial, len(self.vocab))
|
||||
|
||||
# mask according to specified probabilities.
|
||||
masked_blk_one, masked_tgt_one = self._mask_block(
|
||||
s["block_one"],
|
||||
self.mask_idx,
|
||||
self.pad_idx,
|
||||
token_range,
|
||||
)
|
||||
|
||||
tokens = np.concatenate([[self.classif_token_idx], masked_blk_one])
|
||||
targets = np.concatenate([[self.pad_idx], masked_tgt_one])
|
||||
segments = np.ones(len(tokens)) * self.segment_id
|
||||
|
||||
# if has_pairs is True then we need to add the SEP token to both
|
||||
# the blocks after masking and re-compute segments based on the new
|
||||
# lengths.
|
||||
if self.has_pairs:
|
||||
tokens_one = np.concatenate([tokens, [self.sep_token_idx]])
|
||||
targets_one = np.concatenate([targets, [self.pad_idx]])
|
||||
|
||||
masked_blk_two, masked_tgt_two = self._mask_block(
|
||||
s["block_two"], self.mask_idx, self.pad_idx, token_range
|
||||
)
|
||||
tokens_two = np.concatenate([masked_blk_two, [self.sep_token_idx]])
|
||||
targets_two = np.concatenate([masked_tgt_two, [self.pad_idx]])
|
||||
|
||||
# block + 1 sep + 1 special (CLS)
|
||||
segments_one = np.zeros(len(tokens_one))
|
||||
# block + 1 sep
|
||||
segments_two = np.ones(len(tokens_two))
|
||||
|
||||
tokens = np.concatenate([tokens_one, tokens_two])
|
||||
targets = np.concatenate([targets_one, targets_two])
|
||||
segments = np.concatenate([segments_one, segments_two])
|
||||
|
||||
s["source"] = torch.LongTensor(tokens)
|
||||
s["segment_labels"] = torch.LongTensor(segments)
|
||||
s["lm_target"] = torch.LongTensor(targets)
|
||||
|
||||
def merge(key):
|
||||
return data_utils.collate_tokens(
|
||||
[s[key] for s in samples], pad_idx, eos_idx, left_pad=False
|
||||
)
|
||||
|
||||
return {
|
||||
"id": torch.LongTensor([s["id"] for s in samples]),
|
||||
"ntokens": sum(len(s["source"]) for s in samples),
|
||||
"net_input": {
|
||||
"src_tokens": merge("source"),
|
||||
"segment_labels": merge("segment_labels"),
|
||||
},
|
||||
"lm_target": merge("lm_target"),
|
||||
"sentence_target": torch.LongTensor([s["sentence_target"] for s in samples])
|
||||
if self.has_pairs
|
||||
else None,
|
||||
"nsentences": len(samples),
|
||||
}
|
||||
|
||||
def collater(self, samples: List[Dict]):
|
||||
"""Merge a list of samples to form a mini-batch.
|
||||
|
||||
Args:
|
||||
samples (List[dict]): samples to collate
|
||||
|
||||
Returns:
|
||||
dict: a mini-batch of data
|
||||
"""
|
||||
return self._collate(samples, self.vocab.pad(), self.vocab.eos())
|
||||
|
||||
def num_tokens(self, index: int):
|
||||
"""
|
||||
Return the number of tokens in a sample. This value is used to
|
||||
enforce max-tokens during batching.
|
||||
"""
|
||||
return self.sizes[index]
|
||||
|
||||
def size(self, index: int):
|
||||
"""
|
||||
Return an example's size as a float or tuple. This value is used when
|
||||
filtering a dataset with max-positions.
|
||||
"""
|
||||
return self.sizes[index]
|
||||
|
||||
def ordered_indices(self):
|
||||
"""
|
||||
Return an ordered list of indices. Batches will be constructed based
|
||||
on this order.
|
||||
"""
|
||||
if self.shuffle:
|
||||
return np.random.permutation(len(self))
|
||||
else:
|
||||
order = [np.arange(len(self))]
|
||||
order.append(self.sizes)
|
||||
return np.lexsort(order)
|
||||
|
||||
@property
|
||||
def supports_prefetch(self):
|
||||
return getattr(self.dataset, "supports_prefetch", False)
|
||||
|
||||
def prefetch(self, indices):
|
||||
self.dataset.prefetch(indices)
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from fairseq.data import Dictionary
|
||||
|
||||
|
||||
class MaskedLMDictionary(Dictionary):
|
||||
"""
|
||||
Dictionary for Masked Language Modelling tasks. This extends Dictionary by
|
||||
adding the mask symbol.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pad="<pad>",
|
||||
eos="</s>",
|
||||
unk="<unk>",
|
||||
mask="<mask>",
|
||||
):
|
||||
super().__init__(pad=pad, eos=eos, unk=unk)
|
||||
self.mask_word = mask
|
||||
self.mask_index = self.add_symbol(mask)
|
||||
self.nspecial = len(self.symbols)
|
||||
|
||||
def mask(self):
|
||||
"""Helper to get index of mask symbol"""
|
||||
return self.mask_index
|
||||
|
||||
|
||||
class BertDictionary(MaskedLMDictionary):
|
||||
"""
|
||||
Dictionary for BERT task. This extends MaskedLMDictionary by adding support
|
||||
for cls and sep symbols.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pad="<pad>",
|
||||
eos="</s>",
|
||||
unk="<unk>",
|
||||
mask="<mask>",
|
||||
cls="<cls>",
|
||||
sep="<sep>",
|
||||
):
|
||||
super().__init__(pad=pad, eos=eos, unk=unk, mask=mask)
|
||||
self.cls_word = cls
|
||||
self.sep_word = sep
|
||||
self.cls_index = self.add_symbol(cls)
|
||||
self.sep_index = self.add_symbol(sep)
|
||||
self.nspecial = len(self.symbols)
|
||||
|
||||
def cls(self):
|
||||
"""Helper to get index of cls symbol"""
|
||||
return self.cls_index
|
||||
|
||||
def sep(self):
|
||||
"""Helper to get index of sep symbol"""
|
||||
return self.sep_index
|
||||
Reference in New Issue
Block a user