mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-23 08:26:06 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from nemo.collections.nlp.data.language_modeling.l2r_lm_dataset import L2RLanguageModelingDataset
|
||||
from nemo.collections.nlp.data.language_modeling.lm_bert_dataset import (
|
||||
BertPretrainingDataset,
|
||||
BertPretrainingPreprocessedDataloader,
|
||||
)
|
||||
from nemo.collections.nlp.data.language_modeling.sentence_dataset import SentenceDataset, TarredSentenceDataset
|
||||
@@ -0,0 +1,240 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import io
|
||||
import json
|
||||
from typing import Optional
|
||||
|
||||
import braceexpand
|
||||
import numpy as np
|
||||
import webdataset as wd
|
||||
from torch.utils.data import Dataset, IterableDataset
|
||||
|
||||
from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
|
||||
from nemo.collections.nlp.data.data_utils import dataset_to_ids
|
||||
from nemo.utils import logging
|
||||
|
||||
__all__ = ['L2RLanguageModelingDataset', 'TarredL2RLanguageModelingDataset']
|
||||
|
||||
|
||||
class L2RLanguageModelingDataset(Dataset):
|
||||
"""
|
||||
Dataset for training and evaluating left-to-right language models.
|
||||
|
||||
Args:
|
||||
tokenizer: tokenizer, such as WordTokenizer or CharTokenizer
|
||||
dataset: path to data
|
||||
max_seq_length: maximum sequence length (in tokens) of input tensors
|
||||
batch_step: distance (in tokens) between two successive sequences of
|
||||
the text. By default, it is equal to max_seq_length which corresponds
|
||||
to splitting text into disjoint segments covering full dataset
|
||||
use_cache: bool value, defaults to False. Determines whether the preprocessed,
|
||||
tokenized dataset should be cached into a pickle file. If true, cache is saved
|
||||
at the path provided in `dataset`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: TokenizerSpec,
|
||||
dataset: str,
|
||||
max_seq_length: Optional[int] = 512,
|
||||
batch_step: Optional[int] = None,
|
||||
use_cache: bool = False,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.batch_step = batch_step or self.max_seq_length
|
||||
ids = dataset_to_ids(dataset, tokenizer, cache_ids=use_cache, add_bos_eos=False)
|
||||
self.ids = np.array([j for i in ids for j in i])
|
||||
|
||||
def __len__(self):
|
||||
return (len(self.ids) - self.max_seq_length) // self.batch_step
|
||||
|
||||
def __getitem__(self, idx):
|
||||
left = idx * self.batch_step
|
||||
right = left + self.max_seq_length
|
||||
src_ids = self.ids[left:right]
|
||||
labels = self.ids[left + 1 : right + 1]
|
||||
src_mask = (src_ids != self.tokenizer.pad_id).astype(np.float32)
|
||||
return src_ids, src_mask, labels
|
||||
|
||||
|
||||
class TarredL2RLanguageModelingDataset(IterableDataset):
|
||||
"""
|
||||
A similar Dataset to the L2RLanguageModelingDataset, but which loads tarred tokenized numpy files.
|
||||
Accepts a single JSON metadata manifest file as well as the path(s) to the tarball(s) containing the wav files.
|
||||
The manifest should contain information such as the number of shards, the number of tokens in the corpus,
|
||||
and the number of tokens contained within each shard of the tarfile(s).
|
||||
|
||||
Valid formats for the text_tar_filepaths argument include:
|
||||
(1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar.gz', or
|
||||
(2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
|
||||
|
||||
Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
|
||||
This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
|
||||
Supported opening braces - { <=> (, [, < and the special tag _OP_.
|
||||
Supported closing braces - } <=> ), ], > and the special tag _CL_.
|
||||
For SLURM based tasks, we suggest the use of the special tags for ease of use.
|
||||
See the WebDataset documentation for more information about accepted data and input formats.
|
||||
|
||||
If using multiple processes the number of shards should be divisible by the number of workers to ensure an
|
||||
even split among workers. If it is not divisible, logging will give a warning but training will proceed.
|
||||
|
||||
Additionally, please note that the len() of this DataLayer is assumed to be the number of tokens
|
||||
of the text data. An incorrect manifest length may lead to some DataLoader issues down the line.
|
||||
|
||||
Args:
|
||||
text_tar_filepaths: Either a list of tokenized text tarball filepaths, or a
|
||||
string (can be brace-expandable).
|
||||
metadata_path (str): Path to the metadata manifest.
|
||||
tokenizer: tokenizer, such as WordTokenizer or CharTokenizer
|
||||
dataset: path to data
|
||||
max_seq_length: maximum sequence length (in tokens) of input tensors
|
||||
batch_step: distance (in tokens) between two successive sequences of
|
||||
the text. By default, it is equal to max_seq_length which corresponds
|
||||
to splitting text into disjoint segments covering full dataset
|
||||
shuffle_n (int): How many samples to look ahead and load to be shuffled.
|
||||
See WebDataset documentation for more details.
|
||||
Defaults to 0.
|
||||
shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
|
||||
- `scatter`: The default shard strategy applied by WebDataset, where each node gets
|
||||
a unique set of shards, which are permanently pre-allocated and never changed at runtime.
|
||||
- `replicate`: Optional shard strategy, where each node gets all of the set of shards
|
||||
available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
|
||||
The benefit of replication is that it allows each node to sample data points from the entire
|
||||
dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
|
||||
Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
|
||||
and therefore more than one node may sample the same tarfile, and even sample the same
|
||||
data points! As such, there is no assured guarantee that all samples in the dataset will be
|
||||
sampled at least once during 1 epoch.
|
||||
global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
|
||||
world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_tar_filepaths: str,
|
||||
metadata_path: str,
|
||||
tokenizer,
|
||||
max_seq_length: int = 512,
|
||||
batch_step: int = None,
|
||||
shuffle_n: int = 1,
|
||||
shard_strategy: str = "scatter",
|
||||
global_rank: int = 0,
|
||||
world_size: int = 0,
|
||||
):
|
||||
super(TarredL2RLanguageModelingDataset, self).__init__()
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.batch_step = batch_step or self.max_seq_length
|
||||
|
||||
valid_shard_strategies = ['scatter', 'replicate']
|
||||
if shard_strategy not in valid_shard_strategies:
|
||||
raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
|
||||
|
||||
with open(metadata_path, 'r') as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
self.metadata = metadata
|
||||
|
||||
if isinstance(text_tar_filepaths, str):
|
||||
# Replace '(', '[', '<' and '_OP_' with '{'
|
||||
brace_keys_open = ['(', '[', '<', '_OP_']
|
||||
for bkey in brace_keys_open:
|
||||
if bkey in text_tar_filepaths:
|
||||
text_tar_filepaths = text_tar_filepaths.replace(bkey, "{")
|
||||
|
||||
# Replace ')', ']', '>' and '_CL_' with '}'
|
||||
brace_keys_close = [')', ']', '>', '_CL_']
|
||||
for bkey in brace_keys_close:
|
||||
if bkey in text_tar_filepaths:
|
||||
text_tar_filepaths = text_tar_filepaths.replace(bkey, "}")
|
||||
|
||||
if isinstance(text_tar_filepaths, str):
|
||||
# Brace expand
|
||||
text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
|
||||
|
||||
if shard_strategy == 'scatter':
|
||||
logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
|
||||
|
||||
if len(text_tar_filepaths) % world_size != 0:
|
||||
logging.warning(
|
||||
f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
|
||||
f"by number of distributed workers ({world_size})."
|
||||
)
|
||||
|
||||
begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
|
||||
end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
|
||||
text_tar_filepaths = text_tar_filepaths[begin_idx:end_idx]
|
||||
logging.info(
|
||||
"Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
|
||||
)
|
||||
|
||||
elif shard_strategy == 'replicate':
|
||||
logging.info("All tarred dataset shards will be replicated across all nodes.")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
|
||||
|
||||
self.tarpath = text_tar_filepaths
|
||||
|
||||
# Put together WebDataset
|
||||
self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
|
||||
|
||||
if shuffle_n > 0:
|
||||
self._dataset = self._dataset.shuffle(shuffle_n)
|
||||
else:
|
||||
logging.info("WebDataset will not shuffle files within the tar files.")
|
||||
|
||||
self._dataset = self._dataset.rename(npy='npy', key='__key__').to_tuple('npy', 'key').map(f=self._build_sample)
|
||||
|
||||
def _build_sample(self, tup):
|
||||
# Load file
|
||||
npy, filepath = tup
|
||||
npy = io.BytesIO(npy)
|
||||
data = np.load(npy) # loads np.int64 vector
|
||||
npy.close()
|
||||
|
||||
# Select random contiguous subsegment
|
||||
idx = np.random.randint(0, (len(data) - self.max_seq_length) // self.batch_step)
|
||||
|
||||
# Slice of data chunk
|
||||
left = idx * self.batch_step
|
||||
right = left + self.max_seq_length
|
||||
data = data[left : right + 1]
|
||||
|
||||
# Create batch
|
||||
src_ids = data[:-1]
|
||||
labels = data[1:]
|
||||
src_mask = (src_ids != self.tokenizer.pad_id).astype(np.float32)
|
||||
return src_ids, src_mask, labels
|
||||
|
||||
def __iter__(self):
|
||||
# We need to wrap an infinite generator since the actual files
|
||||
# within the tar files contains large chunks of contiguous data.
|
||||
# This prevents PTL from early exiting the train loop after exhausting
|
||||
# all of the files in one iteration (though the actual dataset is many
|
||||
# times larger due to each file containing a large chunk of data).
|
||||
dl_iter = iter(self._dataset)
|
||||
while True:
|
||||
try:
|
||||
batch = next(dl_iter)
|
||||
yield batch
|
||||
except StopIteration:
|
||||
dl_iter = iter(self._dataset)
|
||||
continue
|
||||
|
||||
def __len__(self):
|
||||
return (self.metadata['num_text'] - self.max_seq_length) // self.batch_step
|
||||
@@ -0,0 +1,406 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import array
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import h5py
|
||||
import numpy as np
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from tqdm import tqdm
|
||||
|
||||
from nemo.collections.nlp.data.data_utils.data_preprocessing import find_newlines, load_data_indices
|
||||
from nemo.core.classes import Dataset
|
||||
|
||||
__all__ = ['BertPretrainingDataset', 'BertPretrainingPreprocessedDataloader']
|
||||
|
||||
|
||||
def load_h5(input_file: str):
|
||||
return h5py.File(input_file, "r")
|
||||
|
||||
|
||||
class BertPretrainingDataset(Dataset):
|
||||
"""
|
||||
Dataset for bert pretraining when using data preprocessing including tokenization
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: object,
|
||||
data_file: str,
|
||||
max_seq_length: Optional[int] = 128,
|
||||
mask_prob: Optional[float] = 0.15,
|
||||
short_seq_prob: Optional[float] = 0.1,
|
||||
seq_a_ratio: Optional[float] = 0.6,
|
||||
sentence_idx_file: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
tokenizer: tokenizer
|
||||
data_file: path to data
|
||||
max_seq_length: maximum sequence length of input tensors
|
||||
mask_probability: proability to mask token
|
||||
short_seq_prob: probability to create a sequence shorter than max_seq_length
|
||||
seq_a_ratio: ratio between lengths of first and second sequence
|
||||
sentence_idx_file: sentence indices file for caching
|
||||
"""
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
# Loading enormous datasets into RAM isn't always feasible -- for
|
||||
# example, the pubmed corpus is 200+ GB, which doesn't fit into RAM on
|
||||
# most computers. To get around this, we store the indices of newlines
|
||||
# in each file so we can seek to and retrieve sentences immediately
|
||||
# from main memory when needed during training.
|
||||
|
||||
# Try and load sentence indices file if already exists
|
||||
sentence_indices, sentence_idx_file, data_dir = load_data_indices(
|
||||
sentence_idx_file, data_file, "sentence_indices"
|
||||
)
|
||||
|
||||
# If sentence indices file doesn't exists, generate and store sentence indices
|
||||
if sentence_indices is None:
|
||||
sentence_indices = {}
|
||||
filenames = [data_file]
|
||||
|
||||
for filename in tqdm(filenames):
|
||||
with open(filename, "rb") as f:
|
||||
contents = f.read()
|
||||
newline_indices = find_newlines(contents)
|
||||
|
||||
if os.path.isdir(data_dir):
|
||||
# Only keep the parts of the filepath that are invariant to
|
||||
# the dataset's location on disk
|
||||
filename = os.path.basename(filename)
|
||||
|
||||
# In python, arrays are much more space-efficient than lists
|
||||
sentence_indices[filename] = array.array("I", newline_indices)
|
||||
|
||||
# Save sentence indices so we don't have to do this again
|
||||
with open(sentence_idx_file, "wb") as f:
|
||||
pickle.dump(sentence_indices, f)
|
||||
|
||||
corpus_size = 0
|
||||
empty_files = []
|
||||
|
||||
# Find total number of newlines across entire corpus and remove files
|
||||
# without any newlines
|
||||
for filename in sentence_indices:
|
||||
if len(sentence_indices[filename]) <= 1:
|
||||
empty_files.append(filename)
|
||||
else:
|
||||
corpus_size += len(sentence_indices[filename])
|
||||
|
||||
for filename in empty_files:
|
||||
del sentence_indices[filename]
|
||||
|
||||
self.corpus_size = corpus_size
|
||||
self.dataset = data_dir
|
||||
self.filenames = list(sentence_indices.keys())
|
||||
self.mask_probability = mask_prob
|
||||
self.max_seq_length = max_seq_length
|
||||
self.sentence_indices = sentence_indices
|
||||
self.vocab_size = self.tokenizer.vocab_size
|
||||
self.short_seq_prob = short_seq_prob
|
||||
self.seq_a_ratio = seq_a_ratio
|
||||
|
||||
def __len__(self):
|
||||
return self.corpus_size
|
||||
|
||||
def __getitem__(self, idx: int, min_doc_length: Optional[int] = 16):
|
||||
# Each sequence has three special tokens, as follows:
|
||||
# tokenizer.cls_token <document a> tokenizer.sep_token <document b> tokenizer.eos_token
|
||||
num_special_tokens = 3
|
||||
|
||||
max_num_tokens = self.max_seq_length - num_special_tokens
|
||||
target_seq_length = max_num_tokens
|
||||
if random.random() < self.short_seq_prob:
|
||||
# TODO: maybe introduce an argument to control this.
|
||||
target_seq_length = random.randint(2, max_num_tokens)
|
||||
|
||||
# prefer the seq_a to be slightly longer than seq_b, 0.6 by default
|
||||
target_seq_length_a = int(round(target_seq_length * self.seq_a_ratio))
|
||||
target_seq_length_b = target_seq_length - target_seq_length_a
|
||||
|
||||
def get_document(filepath, offset):
|
||||
# Retrieve a specific line from a file and return as a document
|
||||
if os.path.isdir(self.dataset):
|
||||
filepath = os.path.join(self.dataset, filepath)
|
||||
|
||||
with open(filepath, "rb") as f:
|
||||
f.seek(offset)
|
||||
doc_text = f.readline()[:-1].decode("utf-8", errors="ignore")
|
||||
document = self.tokenizer.text_to_ids(doc_text)
|
||||
|
||||
return document
|
||||
|
||||
def match_target_seq_length(
|
||||
document: str, target_seq_length: int, filename: str, line_idx: int, sentence_indices: Dict[str, dict]
|
||||
):
|
||||
# If document is shorter than target sequence length,
|
||||
# append the next line or take a random line as replacement.
|
||||
num_lines = len(sentence_indices[filename])
|
||||
|
||||
while len(document) < target_seq_length:
|
||||
if line_idx < (num_lines - 1):
|
||||
# append the next line
|
||||
line_idx += 1
|
||||
else:
|
||||
# current line is the last line, take a random one
|
||||
line_idx = random.randrange(num_lines)
|
||||
document = []
|
||||
|
||||
offset = sentence_indices[filename][line_idx]
|
||||
document += get_document(filename, offset)
|
||||
|
||||
return document, line_idx
|
||||
|
||||
# Take sequence A from a random file and a random line
|
||||
a_filename = random.choice(self.filenames)
|
||||
a_line_idx = random.randrange(len(self.sentence_indices[a_filename]))
|
||||
a_line_offset = self.sentence_indices[a_filename][a_line_idx]
|
||||
a_document = get_document(a_filename, a_line_offset)
|
||||
a_document, a_line_idx = match_target_seq_length(
|
||||
a_document, target_seq_length_a, a_filename, a_line_idx, self.sentence_indices
|
||||
)
|
||||
|
||||
is_last_line = a_line_idx >= (len(self.sentence_indices[a_filename]) - 1)
|
||||
# About 50% of the time, B is a random sentence from the corpus
|
||||
take_random_b = (random.random() < 0.5) or is_last_line
|
||||
|
||||
if take_random_b:
|
||||
# This should rarely go for more than one iteration for large
|
||||
# corpora. However, just to be careful, we try to make sure that
|
||||
# the random document is not the same as the document
|
||||
# we're processing.
|
||||
for _ in range(10):
|
||||
b_filename = random.choice(self.filenames)
|
||||
b_line_idx = random.choice(range(len(self.sentence_indices[b_filename])))
|
||||
if b_filename != a_filename:
|
||||
break
|
||||
else:
|
||||
# Take another line from the same file
|
||||
b_line_pos = self.sentence_indices[b_filename][b_line_idx]
|
||||
a_line_pos = self.sentence_indices[a_filename][a_line_idx]
|
||||
# TODO unclear about the following check
|
||||
if abs(b_line_pos - a_line_pos) > max_num_tokens:
|
||||
break
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
b_filename = a_filename
|
||||
b_line_idx = a_line_idx + 1
|
||||
|
||||
is_next = int(not take_random_b)
|
||||
b_line_pos = self.sentence_indices[b_filename][b_line_idx]
|
||||
b_document = get_document(b_filename, b_line_pos)
|
||||
b_document, b_line_idx = match_target_seq_length(
|
||||
b_document, target_seq_length_b, b_filename, b_line_idx, self.sentence_indices
|
||||
)
|
||||
|
||||
def truncate_seq_pair(a, b, max_num_tokens):
|
||||
# Truncates a pair of sequences to a maximum sequence length
|
||||
while (len(a) + len(b)) > max_num_tokens:
|
||||
# Truncate the longer sequence
|
||||
if len(a) > len(b):
|
||||
trunc_document = a
|
||||
else:
|
||||
trunc_document = b
|
||||
|
||||
if len(trunc_document) <= 1:
|
||||
raise ValueError(
|
||||
"Input text corpora probably too small. "
|
||||
"Failed to truncate sequence pair to "
|
||||
"maximum sequence legnth."
|
||||
)
|
||||
|
||||
# Randomly truncate from the front or the back
|
||||
if random.random() < 0.5:
|
||||
del trunc_document[0]
|
||||
else:
|
||||
trunc_document.pop()
|
||||
|
||||
truncate_seq_pair(a_document, b_document, max_num_tokens)
|
||||
|
||||
output_ids = (
|
||||
[self.tokenizer.cls_id] + a_document + [self.tokenizer.sep_id] + b_document + [self.tokenizer.eos_id]
|
||||
)
|
||||
|
||||
input_ids, output_mask = self.mask_ids(output_ids)
|
||||
|
||||
input_mask = np.zeros(self.max_seq_length, dtype=np.long)
|
||||
input_mask[: len(input_ids)] = 1
|
||||
|
||||
input_type_ids = np.zeros(self.max_seq_length, dtype=np.int)
|
||||
input_type_ids[len(a_document) + 2 : len(output_ids) + 1] = 1
|
||||
|
||||
padding_length = max(0, self.max_seq_length - len(input_ids))
|
||||
if padding_length > 0:
|
||||
input_ids.extend([self.tokenizer.pad_id] * padding_length)
|
||||
output_ids.extend([self.tokenizer.pad_id] * padding_length)
|
||||
output_mask.extend([0] * padding_length)
|
||||
|
||||
# TODO: wrap the return value with () for consistent style.
|
||||
return (
|
||||
np.array(input_ids),
|
||||
input_type_ids,
|
||||
np.array(input_mask, dtype=np.long),
|
||||
np.array(output_ids),
|
||||
np.array(output_mask, dtype=np.float32),
|
||||
is_next,
|
||||
)
|
||||
|
||||
def mask_ids(self, ids: List[int]):
|
||||
"""
|
||||
Args:
|
||||
ids: list of token ids representing a chunk of text
|
||||
Returns:
|
||||
masked_ids: list of input tokens with some of the entries masked
|
||||
according to the following protocol from the original BERT paper:
|
||||
each token is masked with a probability of 15% and is replaced with
|
||||
1) the [MASK] token 80% of the time,
|
||||
2) random token 10% of the time,
|
||||
3) the same token 10% of the time.
|
||||
output_mask: list of binary variables which indicate what tokens has
|
||||
been masked (to calculate the loss function for these tokens only)
|
||||
"""
|
||||
|
||||
# Whole-word masking by default, as it gives better performance.
|
||||
cand_indexes = [[ids[0]]]
|
||||
for tid in ids[1:]:
|
||||
token = self.tokenizer.ids_to_tokens([tid])[0]
|
||||
is_suffix = token.startswith('\u2581')
|
||||
if is_suffix:
|
||||
# group together with its previous token to form a whole-word
|
||||
cand_indexes[-1].append(tid)
|
||||
else:
|
||||
cand_indexes.append([tid])
|
||||
|
||||
masked_ids, output_mask = [], []
|
||||
mask_id = self.tokenizer.token_to_id("[MASK]")
|
||||
|
||||
for word_ids in cand_indexes:
|
||||
is_special = (word_ids[0] == self.tokenizer.cls_id) or (word_ids[0] == self.tokenizer.sep_id)
|
||||
if is_special or (random.random() > self.mask_probability):
|
||||
output_mask.extend([0] * len(word_ids))
|
||||
masked_ids.extend(word_ids)
|
||||
else:
|
||||
output_mask.extend([1] * len(word_ids))
|
||||
p = random.random()
|
||||
# for 80%, replace with mask
|
||||
if p < 0.8:
|
||||
masked_ids.extend([mask_id] * len(word_ids))
|
||||
# for 10%, replace by a random token
|
||||
elif p < 0.9:
|
||||
for _ in word_ids:
|
||||
# randomly select a valid word
|
||||
random_word = random.randrange(self.vocab_size)
|
||||
while random_word in (self.tokenizer.cls_id, self.tokenizer.sep_id):
|
||||
random_word = random.randrange(self.vocab_size)
|
||||
masked_ids.append(random_word)
|
||||
# for 10%, use same token
|
||||
else:
|
||||
masked_ids.extend(word_ids)
|
||||
|
||||
return masked_ids, output_mask
|
||||
|
||||
|
||||
class BertPretrainingPreprocessedDataset(Dataset):
|
||||
"""
|
||||
Dataset for already preprocessed data.
|
||||
"""
|
||||
|
||||
def __init__(self, input_file: str, max_predictions_per_seq: int):
|
||||
"""
|
||||
Args:
|
||||
input_file: data file in hdf5 format with preprocessed data in array format
|
||||
max_predictions_per_seq: maximum number of masked tokens per sequence. Need to be consistent with data in input file.
|
||||
"""
|
||||
self.input_file = input_file
|
||||
self.max_predictions_per_seq = max_predictions_per_seq
|
||||
f = load_h5(input_file)
|
||||
keys = [
|
||||
'input_ids',
|
||||
'input_mask',
|
||||
'segment_ids',
|
||||
'masked_lm_positions',
|
||||
'masked_lm_ids',
|
||||
'next_sentence_labels',
|
||||
]
|
||||
self.inputs = [np.asarray(f[key][:]) for key in keys]
|
||||
f.close()
|
||||
|
||||
def __len__(self):
|
||||
'Denotes the total number of samples'
|
||||
return len(self.inputs[0])
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
[input_ids, input_mask, segment_ids, masked_lm_positions, masked_lm_ids, next_sentence_labels] = [
|
||||
input[index].astype(np.int64) for input in self.inputs
|
||||
]
|
||||
|
||||
output_mask = np.zeros_like(input_ids)
|
||||
output_ids = input_ids.copy()
|
||||
|
||||
index = self.max_predictions_per_seq
|
||||
padded_mask_indices = (masked_lm_positions == 0).nonzero()
|
||||
if len(padded_mask_indices[0]) != 0:
|
||||
index = padded_mask_indices[0][0]
|
||||
|
||||
output_mask[masked_lm_positions[:index]] = 1.0
|
||||
output_ids[masked_lm_positions[:index]] = masked_lm_ids[:index]
|
||||
|
||||
# input_mask = np.asarray(input_mask, dtype=np.float32)
|
||||
# output_mask = np.asarray(output_mask, dtype=np.float32)
|
||||
return (input_ids, segment_ids, input_mask, output_ids, output_mask, next_sentence_labels)
|
||||
|
||||
|
||||
class BertPretrainingPreprocessedDataloader(DataLoader):
|
||||
"""
|
||||
Dataloader for already preprocessed data in hdf5 files that is already in the format expected by BERT model.
|
||||
"""
|
||||
|
||||
def __init__(self, data_files: List[str], max_predictions_per_seq: int, batch_size: int, seed: Optional[int] = 42):
|
||||
"""
|
||||
Args:
|
||||
data_files: list of data files in hdf5 format with preprocessed data in array format
|
||||
max_predictions_per_seq: maximum number of masked tokens per sequence. Need to be consistent with data in input file.
|
||||
batch_size: batch size per gpu per forward pass
|
||||
seed: seed to ensure each gpu process opens the same data file in each iteration
|
||||
"""
|
||||
super().__init__(None, batch_size=batch_size)
|
||||
self.random = random.Random(seed)
|
||||
self.data_files = data_files
|
||||
self.max_predictions_per_seq = max_predictions_per_seq
|
||||
|
||||
# def __len__(self):
|
||||
# return sum([len(load_h5(data_file)['input_ids']) for data_file in self.data_files])//(self.batch_size)
|
||||
|
||||
def __iter__(self):
|
||||
self.random.shuffle(self.data_files)
|
||||
for data_file in self.data_files:
|
||||
train_data = BertPretrainingPreprocessedDataset(
|
||||
input_file=data_file, max_predictions_per_seq=self.max_predictions_per_seq
|
||||
)
|
||||
train_sampler = DistributedSampler(train_data)
|
||||
# print("---")
|
||||
# print(os.getpid(), train_sampler.rank, train_sampler.num_replicas, train_sampler.num_samples)
|
||||
# print("---")
|
||||
train_dataloader = DataLoader(
|
||||
dataset=train_data, sampler=train_sampler, batch_size=self.batch_size, shuffle=False,
|
||||
)
|
||||
for x in train_dataloader:
|
||||
yield x
|
||||
@@ -0,0 +1,273 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Pytorch Dataset with sentences packed into batches by length."""
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import pickle
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
import braceexpand
|
||||
import numpy as np
|
||||
import webdataset as wd
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
from nemo.collections.nlp.data.data_utils.data_preprocessing import dataset_to_ids
|
||||
from nemo.core import Dataset
|
||||
|
||||
__all__ = ['SentenceDataset', 'TarredSentenceDataset']
|
||||
|
||||
|
||||
class SentenceDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: Any,
|
||||
dataset: Any,
|
||||
tokens_in_batch: int = 1024,
|
||||
clean: bool = False,
|
||||
cache_ids: bool = False,
|
||||
max_seq_length: int = 512,
|
||||
min_seq_length: int = 1,
|
||||
):
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.tokens_in_batch = tokens_in_batch
|
||||
|
||||
ids = dataset_to_ids(dataset, tokenizer, cache_ids=cache_ids)
|
||||
if clean:
|
||||
ids = self.clean(ids, max_tokens=max_seq_length, min_tokens=min_seq_length)
|
||||
self.batch_sent_ids, self.batch_elem_lengths = self.pack_data_into_batches(ids)
|
||||
self.batches = self.pad_batches(ids)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.batches)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
ids = self.batches[idx]
|
||||
mask = (ids != self.tokenizer.pad_id).astype(np.int32)
|
||||
return ids, mask
|
||||
|
||||
def pad_batches(self, ids):
|
||||
"""
|
||||
Augments source and target ids in the batches with padding symbol
|
||||
to make the lengths of all sentences in the batches equal.
|
||||
"""
|
||||
|
||||
batches = []
|
||||
for batch_elem_len, batch_sent_ids in zip(self.batch_elem_lengths, self.batch_sent_ids):
|
||||
batch = self.tokenizer.pad_id * np.ones((len(batch_sent_ids), batch_elem_len), dtype=np.int)
|
||||
for i, sentence_idx in enumerate(batch_sent_ids):
|
||||
batch[i][: len(ids[sentence_idx])] = ids[sentence_idx]
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
def pack_data_into_batches(self, ids):
|
||||
"""
|
||||
Takes two lists of source and target sentences, sorts them, and packs
|
||||
into batches to minimize the use of padding tokens. Returns a list of
|
||||
batches where each batch contains indices of sentences included into it
|
||||
"""
|
||||
|
||||
# create buckets sorted by the number of src tokens
|
||||
# each bucket is also sorted by the number of tgt tokens
|
||||
buckets = {}
|
||||
for i, line_ids in enumerate(ids):
|
||||
len_ = len(line_ids)
|
||||
if len_ not in buckets:
|
||||
buckets[len_] = [i]
|
||||
else:
|
||||
buckets[len_].append(i)
|
||||
|
||||
for b_idx in buckets:
|
||||
buckets[b_idx] = sorted(buckets[b_idx])
|
||||
|
||||
buckets = OrderedDict(sorted(buckets.items()))
|
||||
|
||||
batches = []
|
||||
batch_elem_lengths = []
|
||||
curr_batch = []
|
||||
len_of_longest_sent = 0
|
||||
for sent_len, bucket in buckets.items():
|
||||
for sent_i in bucket:
|
||||
if sent_len * (len(curr_batch) + 1) > self.tokens_in_batch:
|
||||
if not curr_batch:
|
||||
raise ValueError(
|
||||
f"The limitation on number of tokens in batch {self.tokens_in_batch} is too strong."
|
||||
f"Several sentences contain {sent_len} tokens."
|
||||
)
|
||||
batches.append(curr_batch)
|
||||
batch_elem_lengths.append(sent_len)
|
||||
curr_batch = []
|
||||
curr_batch.append(sent_i)
|
||||
len_of_longest_sent = sent_len
|
||||
if curr_batch:
|
||||
batches.append(curr_batch)
|
||||
batch_elem_lengths.append(len_of_longest_sent)
|
||||
return batches, batch_elem_lengths
|
||||
|
||||
def clean(self, ids, max_tokens=None, min_tokens=None):
|
||||
"""
|
||||
Cleans source and target sentences to get rid of noisy data.
|
||||
Specifically, a pair of sentences is removed if
|
||||
-- either source or target is longer than *max_tokens*
|
||||
-- either source or target is shorter than *min_tokens*
|
||||
-- absolute difference between source and target is larger than
|
||||
*max_tokens_diff*
|
||||
-- one sentence is *max_tokens_ratio* times longer than the other
|
||||
"""
|
||||
|
||||
ids_ = []
|
||||
for i in range(len(ids)):
|
||||
len_ = len(ids[i])
|
||||
if (max_tokens is not None and len_ > max_tokens) or (min_tokens is not None and len_ < min_tokens):
|
||||
continue
|
||||
ids_.append(ids[i])
|
||||
return ids_
|
||||
|
||||
|
||||
class TarredSentenceDataset(IterableDataset):
|
||||
"""
|
||||
A similar Dataset to the SentenceDataset, but which loads tarred tokenized pickle files.
|
||||
Accepts a single JSON metadata file containing the total number of batches
|
||||
as well as the path(s) to the tarball(s) containing the wav files.
|
||||
Valid formats for the text_tar_filepaths argument include:
|
||||
(1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar.gz', or
|
||||
(2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
|
||||
Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
|
||||
This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
|
||||
Supported opening braces - { <=> (, [, < and the special tag _OP_.
|
||||
Supported closing braces - } <=> ), ], > and the special tag _CL_.
|
||||
For SLURM based tasks, we suggest the use of the special tags for ease of use.
|
||||
See the WebDataset documentation for more information about accepted data and input formats.
|
||||
If using multiple processes the number of shards should be divisible by the number of workers to ensure an
|
||||
even split among workers. If it is not divisible, logging will give a warning but training will proceed.
|
||||
Additionally, please note that the len() of this DataLayer is assumed to be the number of tokens
|
||||
of the text data. An incorrect manifest length may lead to some DataLoader issues down the line.
|
||||
Args:
|
||||
text_tar_filepaths: Either a list of tokenized text tarball filepaths, or a
|
||||
string (can be brace-expandable).
|
||||
metadata_path (str): Path to the metadata manifest.
|
||||
encoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
|
||||
decoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
|
||||
shuffle_n (int): How many samples to look ahead and load to be shuffled.
|
||||
See WebDataset documentation for more details.
|
||||
Defaults to 0.
|
||||
shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
|
||||
- `scatter`: The default shard strategy applied by WebDataset, where each node gets
|
||||
a unique set of shards, which are permanently pre-allocated and never changed at runtime.
|
||||
- `replicate`: Optional shard strategy, where each node gets all of the set of shards
|
||||
available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
|
||||
The benefit of replication is that it allows each node to sample data points from the entire
|
||||
dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
|
||||
Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
|
||||
and therefore more than one node may sample the same tarfile, and even sample the same
|
||||
data points! As such, there is no assured guarantee that all samples in the dataset will be
|
||||
sampled at least once during 1 epoch.
|
||||
global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
|
||||
world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
|
||||
reverse_lang_direction (bool): When True, swaps the source and target directions when returning minibatches.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_tar_filepaths: str,
|
||||
metadata_path: str,
|
||||
tokenizer: str,
|
||||
shuffle_n: int = 1,
|
||||
shard_strategy: str = "scatter",
|
||||
global_rank: int = 0,
|
||||
world_size: int = 0,
|
||||
):
|
||||
super(TarredSentenceDataset, self).__init__()
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.pad_id = tokenizer.pad_id
|
||||
|
||||
valid_shard_strategies = ['scatter', 'replicate']
|
||||
if shard_strategy not in valid_shard_strategies:
|
||||
raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
|
||||
|
||||
with open(metadata_path, 'r') as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
self.metadata = metadata
|
||||
|
||||
if isinstance(text_tar_filepaths, str):
|
||||
# Replace '(', '[', '<' and '_OP_' with '{'
|
||||
brace_keys_open = ['(', '[', '<', '_OP_']
|
||||
for bkey in brace_keys_open:
|
||||
if bkey in text_tar_filepaths:
|
||||
text_tar_filepaths = text_tar_filepaths.replace(bkey, "{")
|
||||
|
||||
# Replace ')', ']', '>' and '_CL_' with '}'
|
||||
brace_keys_close = [')', ']', '>', '_CL_']
|
||||
for bkey in brace_keys_close:
|
||||
if bkey in text_tar_filepaths:
|
||||
text_tar_filepaths = text_tar_filepaths.replace(bkey, "}")
|
||||
|
||||
if isinstance(text_tar_filepaths, str):
|
||||
# Brace expand
|
||||
text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
|
||||
|
||||
if shard_strategy == 'scatter':
|
||||
logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
|
||||
if len(text_tar_filepaths) % world_size != 0:
|
||||
logging.warning(
|
||||
f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
|
||||
f"by number of distributed workers ({world_size})."
|
||||
)
|
||||
begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
|
||||
end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
|
||||
logging.info('Begin Index : %d' % (begin_idx))
|
||||
logging.info('End Index : %d' % (end_idx))
|
||||
text_tar_filepaths = text_tar_filepaths[begin_idx:end_idx]
|
||||
logging.info(
|
||||
"Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
|
||||
)
|
||||
|
||||
elif shard_strategy == 'replicate':
|
||||
logging.info("All tarred dataset shards will be replicated across all nodes.")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
|
||||
|
||||
self.tarpath = text_tar_filepaths
|
||||
|
||||
# Put together WebDataset
|
||||
self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
|
||||
|
||||
if shuffle_n > 0:
|
||||
self._dataset = self._dataset.shuffle(shuffle_n)
|
||||
else:
|
||||
logging.info("WebDataset will not shuffle files within the tar files.")
|
||||
|
||||
self._dataset = self._dataset.rename(pkl='pkl', key='__key__').to_tuple('pkl', 'key').map(f=self._build_sample)
|
||||
|
||||
def _build_sample(self, fname):
|
||||
# Load file
|
||||
pkl_file, _ = fname
|
||||
pkl_file = io.BytesIO(pkl_file)
|
||||
data = pickle.load(pkl_file) # loads np.int64 vector
|
||||
pkl_file.close()
|
||||
ids = data["src"]
|
||||
mask = (ids != self.pad_id).astype(np.int32)
|
||||
return ids, mask
|
||||
|
||||
def __iter__(self):
|
||||
return self._dataset.__iter__()
|
||||
|
||||
def __len__(self):
|
||||
return self.metadata['num_batches']
|
||||
Reference in New Issue
Block a user