Skip to content

GitLab

Commit 22c0e300 authored by Jared Casper's avatar Jared Casper
Browse files

Merge branch 'staging' into 'master' 

April 2020 Release

See merge request ADLR/megatron-lm!69
parents 70174ae3 acf8780b
Expand all Show whitespace changes
Inline Side-by-side
Showing with 1164 additions and 886 deletions
megatron/data/albert_dataset.py → megatron/data/bert_dataset.py
View file @ 22c0e300
# coding=utf-8 # coding=utf-8
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -13,7 +13,7 @@ ...@@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""ALBERT Style dataset.""" """BERT Style dataset."""
import os import os
import time import time
...@@ -22,24 +22,18 @@ import numpy as np ...@@ -22,24 +22,18 @@ import numpy as np
import torch import torch
from torch.utils.data import Dataset from torch.utils.data import Dataset
from megatron import get_tokenizer
from megatron import mpu from megatron import mpu
from megatron.data import helpers
from megatron.data import FullBertTokenizer
from megatron.data.dataset_utils import build_training_sample from megatron.data.dataset_utils import build_training_sample
from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset
from megatron.utils import print_rank_0 from megatron import print_rank_0
def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl, def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
splits_string, train_valid_test_num_samples, train_valid_test_num_samples,
max_seq_length, masked_lm_prob, max_seq_length, masked_lm_prob,
short_seq_prob, seed, skip_warmup): short_seq_prob, seed, skip_warmup):
# Tokenizer is the same
tokenizer = FullBertTokenizer(vocab_file, do_lower_case=True)
print_rank_0(' > using full BERT tokenizer with vocabulary size: {}'.format(
tokenizer.vocab_size()))
# Indexed dataset. # Indexed dataset.
indexed_dataset = get_indexed_dataset_(data_prefix, indexed_dataset = get_indexed_dataset_(data_prefix,
data_impl, data_impl,
...@@ -53,6 +47,7 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl, ...@@ -53,6 +47,7 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl,
# Print stats about the splits. # Print stats about the splits.
print_rank_0(' > dataset split:') print_rank_0(' > dataset split:')
def print_split_stats(name, index): def print_split_stats(name, index):
print_rank_0(' {}:'.format(name)) print_rank_0(' {}:'.format(name))
print_rank_0(' document indices in [{}, {}) total of {} ' print_rank_0(' document indices in [{}, {}) total of {} '
...@@ -79,10 +74,9 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl, ...@@ -79,10 +74,9 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl,
# New doc_idx view. # New doc_idx view.
indexed_dataset.set_doc_idx(doc_idx_ptr[start_index:end_index]) indexed_dataset.set_doc_idx(doc_idx_ptr[start_index:end_index])
# Build the dataset accordingly. # Build the dataset accordingly.
dataset = AlbertDataset( dataset = BertDataset(
name=name, name=name,
indexed_dataset=indexed_dataset, indexed_dataset=indexed_dataset,
tokenizer=tokenizer,
data_prefix=data_prefix, data_prefix=data_prefix,
num_epochs=None, num_epochs=None,
max_num_samples=train_valid_test_num_samples[index], max_num_samples=train_valid_test_num_samples[index],
...@@ -105,9 +99,9 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl, ...@@ -105,9 +99,9 @@ def build_train_valid_test_datasets(vocab_file, data_prefix, data_impl,
return (train_dataset, valid_dataset, test_dataset) return (train_dataset, valid_dataset, test_dataset)
class AlbertDataset(Dataset): class BertDataset(Dataset):
def __init__(self, name, indexed_dataset, tokenizer, data_prefix, def __init__(self, name, indexed_dataset, data_prefix,
num_epochs, max_num_samples, masked_lm_prob, num_epochs, max_num_samples, masked_lm_prob,
max_seq_length, short_seq_prob, seed): max_seq_length, short_seq_prob, seed):
...@@ -117,11 +111,9 @@ class AlbertDataset(Dataset): ...@@ -117,11 +111,9 @@ class AlbertDataset(Dataset):
self.masked_lm_prob = masked_lm_prob self.masked_lm_prob = masked_lm_prob
self.max_seq_length = max_seq_length self.max_seq_length = max_seq_length
# Tokenizer and dataset. # Dataset.
self.tokenizer = tokenizer
self.indexed_dataset = indexed_dataset self.indexed_dataset = indexed_dataset
# Build the samples mapping. # Build the samples mapping.
self.samples_mapping = get_samples_mapping_(self.indexed_dataset, self.samples_mapping = get_samples_mapping_(self.indexed_dataset,
data_prefix, data_prefix,
...@@ -133,22 +125,17 @@ class AlbertDataset(Dataset): ...@@ -133,22 +125,17 @@ class AlbertDataset(Dataset):
self.name) self.name)
# Vocab stuff. # Vocab stuff.
self.vocab_id_list = list(self.tokenizer.inv_vocab.keys()) tokenizer = get_tokenizer()
self.vocab_id_to_token_dict = self.tokenizer.inv_vocab self.vocab_id_list = list(tokenizer.inv_vocab.keys())
self.cls_id = self.tokenizer.vocab['[CLS]'] self.vocab_id_to_token_dict = tokenizer.inv_vocab
self.sep_id = self.tokenizer.vocab['[SEP]'] self.cls_id = tokenizer.cls
self.mask_id = self.tokenizer.vocab['[MASK]'] self.sep_id = tokenizer.sep
self.pad_id = self.tokenizer.vocab['[PAD]'] self.mask_id = tokenizer.mask
self.pad_id = tokenizer.pad
def num_tokens(self):
return self.tokenizer.vocab_size()
def __len__(self): def __len__(self):
return self.samples_mapping.shape[0] return self.samples_mapping.shape[0]
def __getitem__(self, idx): def __getitem__(self, idx):
start_index, end_index, seq_length = self.samples_mapping[idx] start_index, end_index, seq_length = self.samples_mapping[idx]
...@@ -203,7 +190,7 @@ def get_train_valid_test_split_(splits_string, size): ...@@ -203,7 +190,7 @@ def get_train_valid_test_split_(splits_string, size):
splits = splits[:3] splits = splits[:3]
splits_sum = sum(splits) splits_sum = sum(splits)
assert splits_sum > 0.0 assert splits_sum > 0.0
splits = [split/splits_sum for split in splits] splits = [split / splits_sum for split in splits]
splits_index = [0] splits_index = [0]
for index, split in enumerate(splits): for index, split in enumerate(splits):
splits_index.append(splits_index[index] + splits_index.append(splits_index[index] +
...@@ -259,12 +246,16 @@ def get_samples_mapping_(indexed_dataset, ...@@ -259,12 +246,16 @@ def get_samples_mapping_(indexed_dataset,
start_time = time.time() start_time = time.time()
print_rank_0(' > building sapmles index mapping for {} ...'.format( print_rank_0(' > building sapmles index mapping for {} ...'.format(
name)) name))
# First compile and then import.
from megatron.data.dataset_utils import compile_helper
compile_helper()
from megatron.data import helpers
samples_mapping = helpers.build_mapping( samples_mapping = helpers.build_mapping(
indexed_dataset.doc_idx, indexed_dataset.doc_idx,
indexed_dataset.sizes, indexed_dataset.sizes,
num_epochs, num_epochs,
max_num_samples, max_num_samples,
max_seq_length-3, # account for added tokens max_seq_length - 3, # account for added tokens
short_seq_prob, short_seq_prob,
seed, seed,
verbose) verbose)
......
megatron/data/bert_tokenization.py deleted 100644 → 0
View file @ 70174ae3
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.
"""Tokenization classes."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import re
import unicodedata
import six
def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
"""Checks whether the casing config is consistent with the checkpoint name."""
# The casing has to be passed in by the user and there is no explicit check
# as to whether it matches the checkpoint. The casing information probably
# should have been stored in the bert_config.json file, but it's not, so
# we have to heuristically detect it to validate.
if not init_checkpoint:
return
m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
if m is None:
return
model_name = m.group(1)
lower_models = [
"uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
"multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
]
cased_models = [
"cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
"multi_cased_L-12_H-768_A-12"
]
is_bad_config = False
if model_name in lower_models and not do_lower_case:
is_bad_config = True
actual_flag = "False"
case_name = "lowercased"
opposite_flag = "True"
if model_name in cased_models and do_lower_case:
is_bad_config = True
actual_flag = "True"
case_name = "cased"
opposite_flag = "False"
if is_bad_config:
raise ValueError(
"You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
"However, `%s` seems to be a %s model, so you "
"should pass in `--do_lower_case=%s` so that the fine-tuning matches "
"how the model was pre-training. If this error is wrong, please "
"just comment out this check." % (actual_flag, init_checkpoint,
model_name, case_name, opposite_flag))
def convert_to_unicode(text):
"""Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text.decode("utf-8", "ignore")
elif isinstance(text, unicode):
return text
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def printable_text(text):
"""Returns text encoded in a way suitable for print or `tf.logging`."""
# These functions want `str` for both Python2 and Python3, but in one case
# it's a Unicode string and in the other it's a byte string.
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text
elif isinstance(text, unicode):
return text.encode("utf-8")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def load_vocab(vocab_file):
"""Loads a vocabulary file into a dictionary."""
vocab = collections.OrderedDict()
index = 0
with open(vocab_file, "r") as reader:
while True:
token = convert_to_unicode(reader.readline())
if not token:
break
token = token.strip()
vocab[token] = index
index += 1
return vocab
def convert_by_vocab(vocab, items):
"""Converts a sequence of [tokens|ids] using the vocab."""
output = []
for item in items:
output.append(vocab[item])
return output
def convert_tokens_to_ids(vocab, tokens):
return convert_by_vocab(vocab, tokens)
def convert_ids_to_tokens(inv_vocab, ids):
return convert_by_vocab(inv_vocab, ids)
def whitespace_tokenize(text):
"""Runs basic whitespace cleaning and splitting on a piece of text."""
text = text.strip()
if not text:
return []
tokens = text.split()
return tokens
class FullTokenizer(object):
"""Runs end-to-end tokenziation."""
def __init__(self, vocab_file, do_lower_case=True):
self.vocab = load_vocab(vocab_file)
self.inv_vocab = {v: k for k, v in self.vocab.items()}
self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
def tokenize(self, text):
split_tokens = []
for token in self.basic_tokenizer.tokenize(text):
for sub_token in self.wordpiece_tokenizer.tokenize(token):
split_tokens.append(sub_token)
return split_tokens
def convert_tokens_to_ids(self, tokens):
return convert_by_vocab(self.vocab, tokens)
def convert_ids_to_tokens(self, ids):
return convert_by_vocab(self.inv_vocab, ids)
def vocab_size(self):
return len(self.vocab)
class BasicTokenizer(object):
"""Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
def __init__(self, do_lower_case=True):
"""Constructs a BasicTokenizer.
Args:
do_lower_case: Whether to lower case the input.
"""
self.do_lower_case = do_lower_case
def tokenize(self, text):
"""Tokenizes a piece of text."""
text = convert_to_unicode(text)
text = self._clean_text(text)
# This was added on November 1st, 2018 for the multilingual and Chinese
# models. This is also applied to the English models now, but it doesn't
# matter since the English models were not trained on any Chinese data
# and generally don't have any Chinese data in them (there are Chinese
# characters in the vocabulary because Wikipedia does have some Chinese
# words in the English Wikipedia.).
text = self._tokenize_chinese_chars(text)
orig_tokens = whitespace_tokenize(text)
split_tokens = []
for token in orig_tokens:
if self.do_lower_case:
token = token.lower()
token = self._run_strip_accents(token)
split_tokens.extend(self._run_split_on_punc(token))
output_tokens = whitespace_tokenize(" ".join(split_tokens))
return output_tokens
def _run_strip_accents(self, text):
"""Strips accents from a piece of text."""
text = unicodedata.normalize("NFD", text)
output = []
for char in text:
cat = unicodedata.category(char)
if cat == "Mn":
continue
output.append(char)
return "".join(output)
def _run_split_on_punc(self, text):
"""Splits punctuation on a piece of text."""
chars = list(text)
i = 0
start_new_word = True
output = []
while i < len(chars):
char = chars[i]
if _is_punctuation(char):
output.append([char])
start_new_word = True
else:
if start_new_word:
output.append([])
start_new_word = False
output[-1].append(char)
i += 1
return ["".join(x) for x in output]
def _tokenize_chinese_chars(self, text):
"""Adds whitespace around any CJK character."""
output = []
for char in text:
cp = ord(char)
if self._is_chinese_char(cp):
output.append(" ")
output.append(char)
output.append(" ")
else:
output.append(char)
return "".join(output)
def _is_chinese_char(self, cp):
"""Checks whether CP is the codepoint of a CJK character."""
# This defines a "chinese character" as anything in the CJK Unicode block:
# https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
#
# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
# despite its name. The modern Korean Hangul alphabet is a different block,
# as is Japanese Hiragana and Katakana. Those alphabets are used to write
# space-separated words, so they are not treated specially and handled
# like the all of the other languages.
if ((cp >= 0x4E00 and cp <= 0x9FFF) or #
(cp >= 0x3400 and cp <= 0x4DBF) or #
(cp >= 0x20000 and cp <= 0x2A6DF) or #
(cp >= 0x2A700 and cp <= 0x2B73F) or #
(cp >= 0x2B740 and cp <= 0x2B81F) or #
(cp >= 0x2B820 and cp <= 0x2CEAF) or
(cp >= 0xF900 and cp <= 0xFAFF) or #
(cp >= 0x2F800 and cp <= 0x2FA1F)): #
return True
return False
def _clean_text(self, text):
"""Performs invalid character removal and whitespace cleanup on text."""
output = []
for char in text:
cp = ord(char)
if cp == 0 or cp == 0xfffd or _is_control(char):
continue
if _is_whitespace(char):
output.append(" ")
else:
output.append(char)
return "".join(output)
class WordpieceTokenizer(object):
"""Runs WordPiece tokenziation."""
def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
self.vocab = vocab
self.unk_token = unk_token
self.max_input_chars_per_word = max_input_chars_per_word
def tokenize(self, text):
"""Tokenizes a piece of text into its word pieces.
This uses a greedy longest-match-first algorithm to perform tokenization
using the given vocabulary.
For example:
input = "unaffable"
output = ["un", "##aff", "##able"]
Args:
text: A single token or whitespace separated tokens. This should have
already been passed through `BasicTokenizer.
Returns:
A list of wordpiece tokens.
"""
text = convert_to_unicode(text)
output_tokens = []
for token in whitespace_tokenize(text):
chars = list(token)
if len(chars) > self.max_input_chars_per_word:
output_tokens.append(self.unk_token)
continue
is_bad = False
start = 0
sub_tokens = []
while start < len(chars):
end = len(chars)
cur_substr = None
while start < end:
substr = "".join(chars[start:end])
if start > 0:
substr = "##" + substr
if substr in self.vocab:
cur_substr = substr
break
end -= 1
if cur_substr is None:
is_bad = True
break
sub_tokens.append(cur_substr)
start = end
if is_bad:
output_tokens.append(self.unk_token)
else:
output_tokens.extend(sub_tokens)
return output_tokens
def _is_whitespace(char):
"""Checks whether `chars` is a whitespace character."""
# \t, \n, and \r are technically contorl characters but we treat them
# as whitespace since they are generally considered as such.
if char == " " or char == "\t" or char == "\n" or char == "\r":
return True
cat = unicodedata.category(char)
if cat == "Zs":
return True
return False
def _is_control(char):
"""Checks whether `chars` is a control character."""
# These are technically control characters but we count them as whitespace
# characters.
if char == "\t" or char == "\n" or char == "\r":
return False
cat = unicodedata.category(char)
if cat in ("Cc", "Cf"):
return True
return False
def _is_punctuation(char):
"""Checks whether `chars` is a punctuation character."""
cp = ord(char)
# We treat all non-letter/number ASCII as punctuation.
# Characters such as "^", "$", and "`" are not in the Unicode
# Punctuation class but we treat them as punctuation anyways, for
# consistency.
if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
(cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
return True
cat = unicodedata.category(char)
if cat.startswith("P"):
return True
return False
megatron/data/dataset_utils.py
View file @ 22c0e300
...@@ -14,10 +14,27 @@ ...@@ -14,10 +14,27 @@
# limitations under the License. # limitations under the License.
# Most of the code here has been copied from:
# https://github.com/google-research/albert/blob/master/create_pretraining_data.py
# with some modifications.
import collections import collections
import numpy as np import numpy as np
def compile_helper():
"""Compile helper function ar runtime. Make sure this
is invoked on a single process."""
import os
import subprocess
path = os.path.abspath(os.path.dirname(__file__))
ret = subprocess.run(['make', '-C', path])
if ret.returncode != 0:
print("Making C++ dataset helpers module failed, exiting.")
import sys
sys.exit(1)
def build_training_sample(sample, def build_training_sample(sample,
target_seq_length, max_seq_length, target_seq_length, max_seq_length,
vocab_id_list, vocab_id_to_token_dict, vocab_id_list, vocab_id_to_token_dict,
...@@ -132,6 +149,7 @@ def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng): ...@@ -132,6 +149,7 @@ def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng):
tokens.pop() tokens.pop()
return True return True
def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id): def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id):
"""Merge segments A and B, add [CLS] and [SEP] and build tokentypes.""" """Merge segments A and B, add [CLS] and [SEP] and build tokentypes."""
...@@ -233,7 +251,7 @@ def create_masked_lm_predictions(tokens, ...@@ -233,7 +251,7 @@ def create_masked_lm_predictions(tokens,
for idx in range(len(cand_indexes)): for idx in range(len(cand_indexes)):
ngram_index = [] ngram_index = []
for n in ngrams: for n in ngrams:
ngram_index.append(cand_indexes[idx:idx+n]) ngram_index.append(cand_indexes[idx:idx + n])
ngram_indexes.append(ngram_index) ngram_indexes.append(ngram_index)
np_rng.shuffle(ngram_indexes) np_rng.shuffle(ngram_indexes)
...@@ -367,12 +385,12 @@ def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions, ...@@ -367,12 +385,12 @@ def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
assert len(masked_positions) == len(masked_labels) assert len(masked_positions) == len(masked_labels)
# Tokens and token types. # Tokens and token types.
filler = [pad_id]*padding_length filler = [pad_id] * padding_length
tokens_np = np.array(tokens + filler, dtype=np.int64) tokens_np = np.array(tokens + filler, dtype=np.int64)
tokentypes_np = np.array(tokentypes + filler, dtype=np.int64) tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)
# Padding mask. # Padding mask.
padding_mask_np = np.array([1]*num_tokens + [0]*padding_length, padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
dtype=np.int64) dtype=np.int64)
# Lables and loss mask. # Lables and loss mask.
......
megatron/data/gpt2_dataset.py 0 → 100644
View file @ 22c0e300
# coding=utf-8
# 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.
"""GPT2 style dataset."""
import os
import time
import numpy as np
import torch
from megatron import print_rank_0
from megatron import mpu
from megatron.data.bert_dataset import get_train_valid_test_split_
from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset
def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
train_valid_test_num_samples,
seq_length, seed, skip_warmup):
"""Build train, valid, and test datasets."""
# Indexed dataset.
indexed_dataset = get_indexed_dataset_(data_prefix,
data_impl,
skip_warmup)
total_num_of_documents = indexed_dataset.sizes.shape[0]
splits = get_train_valid_test_split_(splits_string, total_num_of_documents)
# Print stats about the splits.
print_rank_0(' > dataset split:')
def print_split_stats(name, index):
print_rank_0(' {}:'.format(name))
print_rank_0(' document indices in [{}, {}) total of {} '
'documents'.format(splits[index], splits[index + 1],
splits[index + 1] - splits[index]))
print_split_stats('train', 0)
print_split_stats('validation', 1)
print_split_stats('test', 2)
def build_dataset(index, name):
dataset = None
if splits[index + 1] > splits[index]:
documents = np.arange(start=splits[index], stop=splits[index + 1],
step=1, dtype=np.int32)
dataset = GPT2Dataset(name, data_prefix,
documents, indexed_dataset,
train_valid_test_num_samples[index],
seq_length, seed)
return dataset
train_dataset = build_dataset(0, 'train')
valid_dataset = build_dataset(1, 'valid')
test_dataset = build_dataset(2, 'test')
return (train_dataset, valid_dataset, test_dataset)
def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
"""Build indexed dataset."""
print_rank_0(' > building dataset index ...')
start_time = time.time()
indexed_dataset = make_indexed_dataset(data_prefix,
data_impl,
skip_warmup)
print_rank_0(' > finished creating indexed dataset in {:4f} '
'seconds'.format(time.time() - start_time))
print_rank_0(' number of documents: {}'.format(
indexed_dataset.sizes.shape[0]))
return indexed_dataset
class GPT2Dataset(torch.utils.data.Dataset):
def __init__(self, name, data_prefix, documents, indexed_dataset,
num_samples, seq_length, seed):
self.name = name
self.indexed_dataset = indexed_dataset
# Checks
assert np.min(documents) >= 0
assert np.max(documents) < indexed_dataset.sizes.shape[0]
# Build index mappings.
self.doc_idx, self.sample_idx, self.shuffle_idx = _build_index_mappings(
self.name, data_prefix, documents, self.indexed_dataset.sizes,
num_samples, seq_length, seed)
def __len__(self):
# -1 is due to data structure used to retieve the index:
# sample i --> [sample_idx[i], sample_idx[i+1])
return self.sample_idx.shape[0] - 1
def __getitem__(self, idx):
# Get the shuffled index.
idx = self.shuffle_idx[idx]
# Start and end documents and offsets.
doc_index_f = self.sample_idx[idx][0]
doc_index_l = self.sample_idx[idx + 1][0]
offset_f = self.sample_idx[idx][1]
offset_l = self.sample_idx[idx + 1][1]
# If we are within the same document, just extract the chunk.
if doc_index_f == doc_index_l:
sample = self.indexed_dataset.get(self.doc_idx[doc_index_f],
offset=offset_f,
length=offset_l - offset_f + 1)
else:
# Otherwise, get the rest of the initial document.
sample_list = [self.indexed_dataset.get(self.doc_idx[doc_index_f],
offset=offset_f)]
# Loop over all in between documents and add the entire document.
for i in range(doc_index_f + 1, doc_index_l):
sample_list.append(self.indexed_dataset.get(self.doc_idx[i]))
# And finally add the relevant portion of last document.
sample_list.append(self.indexed_dataset.get(
self.doc_idx[doc_index_l],
length=offset_l + 1))
sample = np.concatenate(sample_list)
return {'text': np.array(sample, dtype=np.int64)}
def _build_index_mappings(name, data_prefix, documents, sizes,
num_samples, seq_length, seed):
"""Build doc-idx, sample-idx, and shuffle-idx.
doc-idx: is an array (ordered) of documents to be used in training.
sample-idx: is the start document index and document offset for each
training sample.
shuffle-idx: maps the sample index into a random index into sample-idx.
"""
# Number of tokens in each epoch and number of required epochs.
tokens_per_epoch = _num_tokens(documents, sizes)
num_epochs = _num_epochs(tokens_per_epoch, seq_length, num_samples)
# rng state
np_rng = np.random.RandomState(seed=seed)
# Filename of the index mappings.
_filename = data_prefix
_filename += '_{}_indexmap'.format(name)
_filename += '_{}ns'.format(num_samples)
_filename += '_{}sl'.format(seq_length)
_filename += '_{}s'.format(seed)
doc_idx_filename = _filename + '_doc_idx.npy'
sample_idx_filename = _filename + '_sample_idx.npy'
shuffle_idx_filename = _filename + '_shuffle_idx.npy'
# Build the indexed mapping if not exist.
if torch.distributed.get_rank() == 0:
if (not os.path.isfile(doc_idx_filename)) or \
(not os.path.isfile(sample_idx_filename)) or \
(not os.path.isfile(shuffle_idx_filename)):
print_rank_0(' > WARNING: could not find index map files, building '
'the indices on rank 0 ...')
# doc-idx.
start_time = time.time()
doc_idx = _build_doc_idx(documents, num_epochs, np_rng)
np.save(doc_idx_filename, doc_idx, allow_pickle=True)
print_rank_0(' > elasped time to build and save doc-idx mapping '
'(seconds): {:4f}'.format(time.time() - start_time))
# sample-idx.
start_time = time.time()
# Use C++ implementation for speed.
# First compile and then import.
from megatron.data.dataset_utils import compile_helper
compile_helper()
from megatron.data import helpers
assert doc_idx.dtype == np.int32
assert sizes.dtype == np.int32
sample_idx = helpers.build_sample_idx(sizes, doc_idx, seq_length,
num_epochs, tokens_per_epoch)
# sample_idx = _build_sample_idx(sizes, doc_idx, seq_length,
# num_epochs, tokens_per_epoch)
np.save(sample_idx_filename, sample_idx, allow_pickle=True)
print_rank_0(' > elasped time to build and save sample-idx mapping '
'(seconds): {:4f}'.format(time.time() - start_time))
# shuffle-idx.
start_time = time.time()
# -1 is due to data structure used to retieve the index:
# sample i --> [sample_idx[i], sample_idx[i+1])
shuffle_idx = _build_shuffle_idx(sample_idx.shape[0] - 1, np_rng)
np.save(shuffle_idx_filename, shuffle_idx, allow_pickle=True)
print_rank_0(' > elasped time to build and save shuffle-idx mapping'
' (seconds): {:4f}'.format(time.time() - start_time))
# This should be a barrier but nccl barrier assumes
# device_index=rank which is not the case for model
# parallel case
counts = torch.cuda.LongTensor([1])
torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
assert counts[0].item() == torch.distributed.get_world_size(
group=mpu.get_data_parallel_group())
# Load mappings.
start_time = time.time()
print_rank_0(' > loading doc-idx mapping from {}'.format(
doc_idx_filename))
doc_idx = np.load(doc_idx_filename, allow_pickle=True)
print_rank_0(' > loading sample-idx mapping from {}'.format(
sample_idx_filename))
sample_idx = np.load(sample_idx_filename, allow_pickle=True)
print_rank_0(' > loading shuffle-idx mapping from {}'.format(
shuffle_idx_filename))
shuffle_idx = np.load(shuffle_idx_filename, allow_pickle=True)
print_rank_0(' loaded indexed file in {:3.3f} seconds'.format(
time.time() - start_time))
print_rank_0(' total number of samples: {}'.format(
sample_idx.shape[0]))
print_rank_0(' total number of epochs: {}'.format(num_epochs))
return doc_idx, sample_idx, shuffle_idx
def _num_tokens(documents, sizes):
"""Total number of tokens in the dataset."""
return np.sum(sizes[documents])
def _num_epochs(tokens_per_epoch, seq_length, num_samples):
"""Based on number of samples and sequence lenght, calculate how many
epochs will be needed."""
num_epochs = 0
total_tokens = 0
while True:
num_epochs += 1
total_tokens += tokens_per_epoch
# -1 is because we need to retrieve seq_length + 1 token each time
# but the last token will overlap with the first token of the next
# sample except for the last sample.
if ((total_tokens - 1) // seq_length) >= num_samples:
return num_epochs
def _build_doc_idx(documents, num_epochs, np_rng):
"""Build an array with length = number-of-epochs * number-of-dcuments.
Each index is mapped to a corresponding document."""
doc_idx = np.mgrid[0:num_epochs, 0:len(documents)][1]
doc_idx[:] = documents
doc_idx = doc_idx.reshape(-1)
doc_idx = doc_idx.astype(np.int32)
np_rng.shuffle(doc_idx)
return doc_idx
def _build_sample_idx(sizes, doc_idx, seq_length,
num_epochs, tokens_per_epoch):
"""Sample index mapping is a 2D array with sizes
[number-of-samples + 1, 2] where [..., 0] contains
the index into `doc_idx` and [..., 1] is the
starting offset in that document."""
# Total number of samples. For -1 see comments in `_num_epochs`.
num_samples = (num_epochs * tokens_per_epoch - 1) // seq_length
sample_idx = np.zeros([num_samples + 1, 2], dtype=np.int32)
# Index into sample_idx.
sample_index = 0
# Index into doc_idx.
doc_idx_index = 0
# Begining offset for each document.
doc_offset = 0
# Start with first document and no offset.
sample_idx[sample_index][0] = doc_idx_index
sample_idx[sample_index][1] = doc_offset
sample_index += 1
while sample_index <= num_samples:
# Start with a fresh sequence.
remaining_seq_length = seq_length + 1
while remaining_seq_length != 0:
# Get the document length.
doc_id = doc_idx[doc_idx_index]
doc_length = sizes[doc_id] - doc_offset
# And add it to the current sequence.
remaining_seq_length -= doc_length
# If we have more than a full sequence, adjust offset and set
# remaining length to zero so we return from the while loop.
# Note that -1 here is for the same reason we have -1 in
# `_num_epochs` calculations.
if remaining_seq_length <= 0:
doc_offset += (remaining_seq_length + doc_length - 1)
remaining_seq_length = 0
else:
# Otherwise, start from the begining of the next document.
doc_idx_index += 1
doc_offset = 0
# Record the sequence.
sample_idx[sample_index][0] = doc_idx_index
sample_idx[sample_index][1] = doc_offset
sample_index += 1
return sample_idx
def _build_shuffle_idx(size, np_rng):
"""Build the range [0, size) and shuffle."""
dtype_ = np.uint32
if size >= (np.iinfo(np.uint32).max - 1):
dtype_ = np.int64
shuffle_idx = np.arange(start=0, stop=size, step=1, dtype=dtype_)
np_rng.shuffle(shuffle_idx)
return shuffle_idx
megatron/data/helpers.cpp
View file @ 22c0e300
/* /*
coding=utf-8 coding=utf-8
Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License"); Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License. you may not use this file except in compliance with the License.
...@@ -33,6 +33,95 @@ using namespace std; ...@@ -33,6 +33,95 @@ using namespace std;
const int32_t LONG_SENTENCE_LEN = 512; const int32_t LONG_SENTENCE_LEN = 512;
py::array build_sample_idx(const py::array_t<int32_t>& sizes_,
const py::array_t<int32_t>& doc_idx_,
const int32_t seq_length,
const int32_t num_epochs,
const int64_t tokens_per_epoch) {
/* Sample index (sample_idx) is used for gpt2 like dataset for which
the documents are flattened and the samples are built based on this
1-D flatten array. It is a 2D array with sizes [number-of-samples + 1, 2]
where [..., 0] contains the index into `doc_idx` and [..., 1] is the
starting offset in that document.*/
// Consistency checks.
assert(seq_length > 1);
assert(num_epochs > 0);
assert(tokens_per_epoch > 1);
// Remove bound checks.
auto sizes = sizes_.unchecked<1>();
auto doc_idx = doc_idx_.unchecked<1>();
// Mapping and it's length (1D).
int64_t num_samples = (num_epochs * tokens_per_epoch - 1) / seq_length;
int32_t* sample_idx = new int32_t[2*(num_samples+1)];
cout << " using:" << endl << std::flush;
cout << " number of documents: " <<
doc_idx_.shape(0) / num_epochs << endl << std::flush;
cout << " number of epochs: " << num_epochs <<
endl << std::flush;
cout << " sequence length: " << seq_length <<
endl << std::flush;
cout << " total number of samples: " << num_samples <<
endl << std::flush;
// Index into sample_idx.
int64_t sample_index = 0;
// Index into doc_idx.
int64_t doc_idx_index = 0;
// Begining offset for each document.
int32_t doc_offset = 0;
// Start with first document and no offset.
sample_idx[2 * sample_index] = doc_idx_index;
sample_idx[2 * sample_index + 1] = doc_offset;
++sample_index;
while (sample_index <= num_samples) {
// Start with a fresh sequence.
int32_t remaining_seq_length = seq_length + 1;
while (remaining_seq_length != 0) {
// Get the document length.
auto doc_id = doc_idx[doc_idx_index];
auto doc_length = sizes[doc_id] - doc_offset;
// And add it to the current sequence.
remaining_seq_length -= doc_length;
// If we have more than a full sequence, adjust offset and set
// remaining length to zero so we return from the while loop.
// Note that -1 here is for the same reason we have -1 in
// `_num_epochs` calculations.
if (remaining_seq_length <= 0) {
doc_offset += (remaining_seq_length + doc_length - 1);
remaining_seq_length = 0;
} else {
// Otherwise, start from the begining of the next document.
++doc_idx_index;
doc_offset = 0;
}
}
// Record the sequence.
sample_idx[2 * sample_index] = doc_idx_index;
sample_idx[2 * sample_index + 1] = doc_offset;
++sample_index;
}
// Method to deallocate memory.
py::capsule free_when_done(sample_idx, [](void *mem_) {
int32_t *mem = reinterpret_cast<int32_t*>(mem_);
delete[] mem;
});
// Return the numpy array.
const auto byte_size = sizeof(int32_t);
return py::array(std::vector<int64_t>{num_samples+1, 2}, // shape
{2*byte_size, byte_size}, // C-style contiguous strides
sample_idx, // the data pointer
free_when_done); // numpy array references
}
inline int32_t get_target_sample_len(const int32_t short_seq_ratio, inline int32_t get_target_sample_len(const int32_t short_seq_ratio,
const int32_t max_length, const int32_t max_length,
std::mt19937& rand32_gen) { std::mt19937& rand32_gen) {
...@@ -307,4 +396,5 @@ py::array build_mapping(const py::array_t<int64_t>& docs_, ...@@ -307,4 +396,5 @@ py::array build_mapping(const py::array_t<int64_t>& docs_,
PYBIND11_MODULE(helpers, m) { PYBIND11_MODULE(helpers, m) {
m.def("build_mapping", &build_mapping); m.def("build_mapping", &build_mapping);
m.def("build_sample_idx", &build_sample_idx);
} }
megatron/data/indexed_dataset.py
View file @ 22c0e300
...@@ -18,7 +18,8 @@ from itertools import accumulate ...@@ -18,7 +18,8 @@ from itertools import accumulate
import numpy as np import numpy as np
import torch import torch
from megatron.utils import print_rank_0 from megatron import print_rank_0
def __best_fitting_dtype(vocab_size=None): def __best_fitting_dtype(vocab_size=None):
if vocab_size is not None and vocab_size < 65500: if vocab_size is not None and vocab_size < 65500:
...@@ -42,6 +43,8 @@ def infer_dataset_impl(path): ...@@ -42,6 +43,8 @@ def infer_dataset_impl(path):
else: else:
return None return None
else: else:
print(f"Dataset does not exist: {path}")
print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
return None return None
...@@ -53,6 +56,10 @@ def make_builder(out_file, impl, vocab_size=None): ...@@ -53,6 +56,10 @@ def make_builder(out_file, impl, vocab_size=None):
def make_dataset(path, impl, skip_warmup=False): def make_dataset(path, impl, skip_warmup=False):
if not IndexedDataset.exists(path):
print(f"Dataset does not exist: {path}")
print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
return None
if impl == 'infer': if impl == 'infer':
impl = infer_dataset_impl(path) impl = infer_dataset_impl(path)
if impl == 'lazy' and IndexedDataset.exists(path): if impl == 'lazy' and IndexedDataset.exists(path):
...@@ -61,6 +68,7 @@ def make_dataset(path, impl, skip_warmup=False): ...@@ -61,6 +68,7 @@ def make_dataset(path, impl, skip_warmup=False):
return IndexedCachedDataset(path) return IndexedCachedDataset(path)
elif impl == 'mmap' and MMapIndexedDataset.exists(path): elif impl == 'mmap' and MMapIndexedDataset.exists(path):
return MMapIndexedDataset(path, skip_warmup) return MMapIndexedDataset(path, skip_warmup)
print(f"Unknown dataset implementation: {impl}")
return None return None
...@@ -107,13 +115,15 @@ def index_file_path(prefix_path): ...@@ -107,13 +115,15 @@ def index_file_path(prefix_path):
def data_file_path(prefix_path): def data_file_path(prefix_path):
return prefix_path + '.bin' return prefix_path + '.bin'
def create_doc_idx(sizes): def create_doc_idx(sizes):
doc_idx = [0] doc_idx = [0]
for i, s in enumerate(sizes): for i, s in enumerate(sizes):
if s == 0: if s == 0:
doc_idx.append(i+1) doc_idx.append(i + 1)
return doc_idx return doc_idx
class IndexedDataset(torch.utils.data.Dataset): class IndexedDataset(torch.utils.data.Dataset):
"""Loader for IndexedDataset""" """Loader for IndexedDataset"""
_HDR_MAGIC = b'TNTIDX\x00\x00' _HDR_MAGIC = b'TNTIDX\x00\x00'
...@@ -153,7 +163,7 @@ class IndexedDataset(torch.utils.data.Dataset): ...@@ -153,7 +163,7 @@ class IndexedDataset(torch.utils.data.Dataset):
if self.data_file: if self.data_file:
self.data_file.close() self.data_file.close()
#@lru_cache(maxsize=8) # @lru_cache(maxsize=8)
def __getitem__(self, idx): def __getitem__(self, idx):
if not self.data_file: if not self.data_file:
self.read_data(self.path) self.read_data(self.path)
...@@ -233,7 +243,7 @@ class IndexedCachedDataset(IndexedDataset): ...@@ -233,7 +243,7 @@ class IndexedCachedDataset(IndexedDataset):
self.data_file.close() self.data_file.close()
self.data_file = None self.data_file = None
#@lru_cache(maxsize=8) # @lru_cache(maxsize=8)
def __getitem__(self, idx): def __getitem__(self, idx):
if isinstance(idx, int): if isinstance(idx, int):
i = idx i = idx
...@@ -397,13 +407,18 @@ class MMapIndexedDataset(torch.utils.data.Dataset): ...@@ -397,13 +407,18 @@ class MMapIndexedDataset(torch.utils.data.Dataset):
self._bin_buffer_mmap = np.memmap(path, mode='r', order='C') self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
self._bin_buffer = memoryview(self._bin_buffer_mmap) self._bin_buffer = memoryview(self._bin_buffer_mmap)
print_rank_0(" reading sizes...") print_rank_0(" reading sizes...")
self._sizes = np.frombuffer(self._bin_buffer, dtype=np.int32, count=self._len, offset=offset) self._sizes = np.frombuffer(
self._bin_buffer,
dtype=np.int32,
count=self._len,
offset=offset)
print_rank_0(" reading pointers...") print_rank_0(" reading pointers...")
self._pointers = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._len, self._pointers = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._len,
offset=offset + self._sizes.nbytes) offset=offset + self._sizes.nbytes)
print_rank_0(" reading document index...") print_rank_0(" reading document index...")
self._doc_idx = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._doc_count, self._doc_idx = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._doc_count,
offset=offset + self._sizes.nbytes + self._pointers.nbytes) offset=offset + self._sizes.nbytes + self._pointers.nbytes)
def __del__(self): def __del__(self):
self._bin_buffer_mmap._mmap.close() self._bin_buffer_mmap._mmap.close()
del self._bin_buffer_mmap del self._bin_buffer_mmap
...@@ -462,13 +477,12 @@ class MMapIndexedDataset(torch.utils.data.Dataset): ...@@ -462,13 +477,12 @@ class MMapIndexedDataset(torch.utils.data.Dataset):
def __len__(self): def __len__(self):
return len(self._index) return len(self._index)
#@lru_cache(maxsize=8) # @lru_cache(maxsize=8)
def __getitem__(self, idx): def __getitem__(self, idx):
if isinstance(idx, int): if isinstance(idx, int):
ptr, size = self._index[idx] ptr, size = self._index[idx]
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr) np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
if self._index.dtype != np.int64: count=size, offset=ptr)
np_array = np_array.astype(np.int64)
return np_array return np_array
elif isinstance(idx, slice): elif isinstance(idx, slice):
start, stop, step = idx.indices(len(self)) start, stop, step = idx.indices(len(self))
...@@ -478,10 +492,25 @@ class MMapIndexedDataset(torch.utils.data.Dataset): ...@@ -478,10 +492,25 @@ class MMapIndexedDataset(torch.utils.data.Dataset):
sizes = self._index._sizes[idx] sizes = self._index._sizes[idx]
offsets = list(accumulate(sizes)) offsets = list(accumulate(sizes))
total_size = sum(sizes) total_size = sum(sizes)
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=total_size, offset=ptr) np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
count=total_size, offset=ptr)
sents = np.split(np_array, offsets[:-1]) sents = np.split(np_array, offsets[:-1])
return sents return sents
def get(self, idx, offset=0, length=None):
""" Retrieves a single item from the dataset with the option to only
return a portion of the item.
get(idx) is the same as [idx] but get() does not support slicing.
"""
ptr, size = self._index[idx]
if length is None:
length = size - offset
ptr += offset * np.dtype(self._index.dtype).itemsize
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
count=length, offset=ptr)
return np_array
@property @property
def sizes(self): def sizes(self):
return self._index.sizes return self._index.sizes
......
megatron/data/preprocess_data.py deleted 100644 → 0
View file @ 70174ae3
import argparse
import json
import multiprocessing
import nltk
import sys
import time
import torch
from bert_tokenization import FullTokenizer
import indexed_dataset
class CustomLanguageVars(nltk.tokenize.punkt.PunktLanguageVars):
_period_context_fmt = r"""
\S* # some word material
%(SentEndChars)s # a potential sentence ending
\s* # <-- THIS is what I changed
(?=(?P<after_tok>
%(NonWord)s # either other punctuation
|
(?P<next_tok>\S+) # <-- Normally you would have \s+ here
))"""
class Encoder(object):
def __init__(self, args):
self.args = args
def initializer(self):
# Use Encoder class as a container for global data
Encoder.tokenizer = FullTokenizer(self.args.vocab, do_lower_case=True)
spliter = nltk.load("tokenizers/punkt/english.pickle")
if self.args.keep_newlines:
# this prevents punkt from eating newlines after sentences
Encoder.spliter = nltk.tokenize.punkt.PunktSentenceTokenizer(
train_text = spliter._params,
lang_vars = CustomLanguageVars())
else:
Encoder.splitter = spliter
def encode(self, json_line):
text = json.loads(json_line)[self.args.json_key]
doc_ids = []
for sentence in Encoder.splitter.tokenize(text):
tokens = Encoder.tokenizer.tokenize(sentence)
ids = Encoder.tokenizer.convert_tokens_to_ids(tokens)
if len(ids) > 0:
doc_ids.append(ids)
return doc_ids, len(json_line)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, help='Path to input JSON')
parser.add_argument('--vocab', type=str, help='Path to vocab.txt')
parser.add_argument('--json-key', type=str, default='text',
help='Key to extract from json')
parser.add_argument('--output-prefix', type=str, help='Path to binary output file without suffix')
parser.add_argument('--workers', type=int, default=20,
help='Number of worker processes to launch')
parser.add_argument('--log-interval', type=int, default=100,
help='Interval between progress updates')
parser.add_argument('--keep-newlines', action='store_true',
help='Keep newlines between sentences.')
parser.add_argument('--dataset-impl', type=str, default='mmap',
choices=['lazy', 'cached', 'mmap'])
args = parser.parse_args()
args.keep_empty = False
startup_start = time.time()
print("Opening", args.input)
fin = open(args.input, 'r', encoding='utf-8')
nltk.download("punkt", quiet=True)
encoder = Encoder(args)
tokenizer = FullTokenizer(args.vocab, do_lower_case=True)
pool = multiprocessing.Pool(args.workers, initializer=encoder.initializer)
encoded_docs = pool.imap(encoder.encode, fin, 25)
print(f"Vocab size: {tokenizer.vocab_size()}")
output_bin_file = "{}.bin".format(args.output_prefix)
output_idx_file = "{}.idx".format(args.output_prefix)
builder = indexed_dataset.make_builder(output_bin_file,
impl=args.dataset_impl,
vocab_size=tokenizer.vocab_size())
startup_end = time.time()
proc_start = time.time()
total_bytes_processed = 0
print("Time to startup:", startup_end - startup_start)
for i, (doc, bytes_processed) in enumerate(encoded_docs, start=1):
total_bytes_processed += bytes_processed
for sentence in doc:
#print(sentence)
#print(tokenizer.convert_ids_to_tokens(sentence))
builder.add_item(torch.IntTensor(sentence))
builder.end_document()
if i % args.log_interval == 0:
current = time.time()
elapsed = current - proc_start
mbs = total_bytes_processed/elapsed/1024/1024
print(f"Processed {i} documents",
f"({i/elapsed} docs/s, {mbs} MB/s).",
file=sys.stderr)
builder.finalize(output_idx_file)
if __name__ == '__main__':
main()
megatron/data/samplers.py 0 → 100644
View file @ 22c0e300
# coding=utf-8
# 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.
"""Batch samplers that work with either random or sequential data samplers."""
import torch
from torch.utils import data
class RandomSampler(data.sampler.Sampler):
"""Based off of pytorch RandomSampler and DistributedSampler. Essentially
a RandomSampler, but this class lets the user set an epoch like
DistributedSampler Samples elements randomly. If without replacement, then
sample from a shuffled dataset. If with replacement, then user can
specify ``num_samples`` to draw.
Arguments:
data_source (Dataset): dataset to sample from
num_samples (int): number of samples to draw, default=len(dataset)
replacement (bool): samples are drawn with replacement if ``True``,
default=False
"""
def __init__(self, data_source, replacement=False, num_samples=None):
self.data_source = data_source
self.replacement = replacement
self._num_samples = num_samples
self.epoch = -1
if self._num_samples is not None and replacement is False:
raise ValueError("With replacement=False, num_samples should not "
"be specified, since a random permute will be "
"performed.")
if not isinstance(self.num_samples, int) or self.num_samples <= 0:
raise ValueError("num_samples should be a positive integer "
"value, but got num_samples={}".format(
self.num_samples))
if not isinstance(self.replacement, bool):
raise ValueError("replacement should be a boolean value, but got "
"replacement={}".format(self.replacement))
@property
def num_samples(self):
# dataset size might change at runtime
if self._num_samples is None:
return len(self.data_source)
return self._num_samples
def __iter__(self):
n = len(self.data_source)
g = torch.Generator()
if self.epoch >= 0:
g.manual_seed(self.epoch)
if self.replacement:
return iter(torch.randint(high=n, size=(self.num_samples,),
dtype=torch.int64, generator=g).tolist())
return iter(torch.randperm(n, generator=g).tolist())
def __len__(self):
return self.num_samples
def set_epoch(self, epoch):
self.epoch = epoch
class DistributedBatchSampler(data.sampler.BatchSampler):
"""Similar to normal implementation of distributed sampler, except
implementation is at the batch sampler level, instead of just the
sampler level. This allows wrapping of arbitrary data samplers
(sequential, random, WeightedRandomSampler, etc.) with this batch
sampler."""
def __init__(self, sampler, batch_size, drop_last, rank=-1,
world_size=2, wrap_last=False):
super(DistributedBatchSampler, self).__init__(sampler, batch_size,
drop_last)
if rank == -1:
assert False, 'should not be here'
rank = torch.distributed.get_rank()
self.rank = rank
self.world_size = world_size
self.sampler.wrap_around = 0
self.wrap_around = 0
self.wrap_last = wrap_last
self.start_iter = 0
def __iter__(self):
batch = []
i = 0
for idx in self.data_iterator(self.sampler, wrap_around=False):
batch.append(idx)
if len(batch) == self.batch_size:
tbatch = self._batch(batch)
if i >= self.start_iter:
yield tbatch
self.start_iter = 0
i += 1
batch = []
batch_len = len(batch)
if batch_len > 0 and not self.drop_last:
if self.wrap_last:
self.sampler.wrap_around -= (self.batch_size)
self.wrap_around += (len(batch))
self.wrap_around %= self.batch_size
yield self._batch(batch)
if self.wrap_last:
self.sampler.wrap_around += self.batch_size
def data_iterator(self, _iter, wrap_around=False):
"""iterates through data and handles wrap around"""
for i, idx in enumerate(_iter):
if i < self.wrap_around % self.batch_size:
continue
if wrap_around:
self.wrap_around += 1
self.wrap_around %= self.batch_size
yield idx
def _batch(self, batch):
"""extracts samples only pertaining to this worker's batch"""
start = self.rank * self.batch_size // self.world_size
end = (self.rank + 1) * self.batch_size // self.world_size
return batch[start:end]
megatron/data/test/test_indexed_dataset.py
View file @ 22c0e300
# This file isn't really a formal automated test, it's just a place to
# put some code used during development and manual testing of
# indexed_dataset.
from megatron.data import indexed_dataset
from megatron.tokenizer import build_tokenizer
import argparse import argparse
import os import os
import sys import sys
...@@ -7,52 +13,90 @@ import torch ...@@ -7,52 +13,90 @@ import torch
script_dir = os.path.dirname(os.path.realpath(__file__)) script_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(script_dir, "../../../")) sys.path.append(os.path.join(script_dir, "../../../"))
from megatron.data import indexed_dataset, FullBertTokenizer, AlbertDataset
def test_indexed_dataset(args): def test_indexed_dataset(args):
ds = indexed_dataset.make_dataset(args.data, args.dataset_impl) ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True) tokenizer = build_tokenizer(args)
print(len(ds.doc_idx)) print(len(ds.doc_idx))
print(len(ds)) print(len(ds))
print(ds.doc_idx[-1]) print(ds.doc_idx[-1])
if ds.supports_prefetch: if ds.supports_prefetch:
# just prefetch the whole thing in test (so assume it is small) # just prefetch the whole thing in test (so assume it is small)
ds.prefetch(range(len(ds))) ds.prefetch(range(len(ds)))
for i in range(len(ds.doc_idx)-1): if args.count > len(ds.doc_idx) - 1:
args.count = len(ds.doc_idx) - 1
for i in range(args.count):
start = ds.doc_idx[i] start = ds.doc_idx[i]
end = ds.doc_idx[i+1] end = ds.doc_idx[i + 1]
ids = ds[start:end] ids = ds[start:end]
print(f"Document {i}:")
print("--------------")
for s in ids: for s in ids:
assert len(s) > 0 assert len(s) > 0
l = s.data.tolist() l = s.data.tolist()
tokens = tokenizer.convert_ids_to_tokens(l) text = tokenizer.detokenize(l)
for t in tokens: print(text)
if '\n' in t: print("---")
print("Newline in string!")
print(i)
def test_indexed_dataset_get(args):
def test_albert_dataset(args): ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
# tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True) tokenizer = build_tokenizer(args)
# idataset = indexed_dataset.make_dataset(args.data, args.dataset_impl) size = ds.sizes[0]
# ds = AlbertDataset(idataset, tokenizer) print(f"size: {size}")
ds = AlbertDataset.from_paths(args.vocab, args.data, args.dataset_impl, full = ds.get(0)
args.epochs, args.max_num_samples, print(full)
args.masked_lm_prob, args.seq_length, # print(tokenizer.detokenize(full.data.tolist()))
args.short_seq_prob, args.seed) print("---")
truncated = 0 end = ds.get(0, offset=size - 10)
total = 0 print(end)
for s in ds: # print(tokenizer.detokenize(end.data.tolist()))
ids = s['text']
tokens = ds.tokenizer.convert_ids_to_tokens(ids) start = ds.get(0, length=10)
print(tokens) print(start)
exit() # print(tokenizer.detokenize(start.data.tolist()))
part = ds.get(0, offset=2, length=8)
print(part)
# print(tokenizer.detokenize(part.data.tolist()))
# def test_albert_dataset(args):
# # tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True)
# # idataset = indexed_dataset.make_dataset(args.data, args.dataset_impl)
# # ds = AlbertDataset(idataset, tokenizer)
# ds = AlbertDataset.from_paths(args.vocab, args.data, args.dataset_impl,
# args.epochs, args.max_num_samples,
# args.masked_lm_prob, args.seq_length,
# args.short_seq_prob, args.seed)
# truncated = 0
# total = 0
# for i, s in enumerate(ds):
# ids = s['text']
# tokens = ds.tokenizer.convert_ids_to_tokens(ids)
# print(tokens)
# if i >= args.count-1:
# exit()
def main(): def main():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, help='prefix to data files') parser.add_argument('--data', type=str, help='prefix to data files')
parser.add_argument('--vocab', type=str, help='Path to vocab.txt')
parser.add_argument('--dataset-impl', type=str, default='infer', parser.add_argument('--dataset-impl', type=str, default='infer',
choices=['lazy', 'cached', 'mmap', 'infer']) choices=['lazy', 'cached', 'mmap', 'infer'])
parser.add_argument('--count', type=int, default=10,
help='Number of samples/documents to print')
group = parser.add_argument_group(title='tokenizer')
group.add_argument('--tokenizer-type', type=str, required=True,
choices=['BertWordPieceLowerCase',
'GPT2BPETokenizer'],
help='What type of tokenizer to use.')
group.add_argument('--vocab-file', type=str, default=None,
help='Path to the vocab file')
group.add_argument('--merge-file', type=str, default=None,
help='Path to the BPE merge file (if necessary).')
parser.add_argument('--epochs', type=int, default=5, parser.add_argument('--epochs', type=int, default=5,
help='Number of epochs to plan for') help='Number of epochs to plan for')
parser.add_argument('--max-num-samples', type=int, default=None, parser.add_argument('--max-num-samples', type=int, default=None,
...@@ -66,12 +110,16 @@ def main(): ...@@ -66,12 +110,16 @@ def main():
parser.add_argument('--seed', type=int, default=1234, parser.add_argument('--seed', type=int, default=1234,
help='random seed') help='random seed')
args = parser.parse_args() args = parser.parse_args()
args.rank = 0
args.make_vocab_size_divisible_by = 128
args.model_parallel_size = 1
if args.dataset_impl == "infer": if args.dataset_impl == "infer":
args.dataset_impl = indexed_dataset.infer_dataset_impl(args.data) args.dataset_impl = indexed_dataset.infer_dataset_impl(args.data)
test_albert_dataset(args) # test_albert_dataset(args)
# test_indexed_dataset(args) test_indexed_dataset_get(args)
if __name__ == "__main__": if __name__ == "__main__":
main() main()
megatron/data_utils/__init__.py → megatron/deprecated_data_utils/__init__.py
View file @ 22c0e300
# coding=utf-8 # coding=utf-8
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -28,6 +28,7 @@ TRAIN_DATA = 0 ...@@ -28,6 +28,7 @@ TRAIN_DATA = 0
VAL_DATA = 1 VAL_DATA = 1
TEST_DATA = 2 TEST_DATA = 2
def should_split(split): def should_split(split):
""" """
given split proportions checks if should split given split proportions checks if should split
...@@ -37,12 +38,14 @@ def should_split(split): ...@@ -37,12 +38,14 @@ def should_split(split):
>>> should_split([1,.1,.2]) >>> should_split([1,.1,.2])
True True
""" """
return max(split)/sum(split) != 1. return max(split) / sum(split) != 1.
def get_ext(path): def get_ext(path):
"""gets path extension""" """gets path extension"""
return os.path.splitext(path)[1] return os.path.splitext(path)[1]
def get_dataset(path, **kwargs): def get_dataset(path, **kwargs):
"""gets dataset object based on keyword args and file at `path`""" """gets dataset object based on keyword args and file at `path`"""
if supported_corpus(path): if supported_corpus(path):
...@@ -53,13 +56,15 @@ def get_dataset(path, **kwargs): ...@@ -53,13 +56,15 @@ def get_dataset(path, **kwargs):
elif ext in ['.csv', '.tsv']: elif ext in ['.csv', '.tsv']:
text = csv_dataset(path, **kwargs) text = csv_dataset(path, **kwargs)
else: else:
raise NotImplementedError('data file type %s is not supported'%(ext)) raise NotImplementedError('data file type %s is not supported' % (ext))
return text return text
def supported_corpus(corpus_name): def supported_corpus(corpus_name):
"""checks if corpus name is defined in `corpora.py`""" """checks if corpus name is defined in `corpora.py`"""
return corpus_name in corpora.NAMED_CORPORA return corpus_name in corpora.NAMED_CORPORA
def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=None, split=[1.], def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=None, split=[1.],
delim=',', loose=False, binarize_sent=False, drop_unlabeled=False, tokenizer=None, delim=',', loose=False, binarize_sent=False, drop_unlabeled=False, tokenizer=None,
tokenizer_type='CharacterLevelTokenizer', tokenizer_model_path=None, vocab_size=None, tokenizer_type='CharacterLevelTokenizer', tokenizer_model_path=None, vocab_size=None,
...@@ -71,6 +76,7 @@ def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=N ...@@ -71,6 +76,7 @@ def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=N
if non_binary_cols is not None: if non_binary_cols is not None:
# multilabel dataset support (only for csvs) # multilabel dataset support (only for csvs)
label_key = non_binary_cols label_key = non_binary_cols
def get_dataset_from_path(path_): def get_dataset_from_path(path_):
if lazy: if lazy:
# get lazily loaded dataset # get lazily loaded dataset
...@@ -121,7 +127,8 @@ def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=N ...@@ -121,7 +127,8 @@ def make_dataset(path, seq_length, text_key, label_key, lazy=False, process_fn=N
if 'bert' in ds_type.lower(): if 'bert' in ds_type.lower():
presplit_sentences = kwargs['presplit_sentences'] if 'presplit_sentences' in kwargs else False presplit_sentences = kwargs['presplit_sentences'] if 'presplit_sentences' in kwargs else False
dstype = bert_sentencepair_dataset dstype = bert_sentencepair_dataset
ds = [dstype(d, max_seq_len=seq_length, presplit_sentences=presplit_sentences) if d is not None else None for d in ds] ds = [dstype(d, max_seq_len=seq_length, presplit_sentences=presplit_sentences)
if d is not None else None for d in ds]
elif ds_type.lower() == 'gpt2': elif ds_type.lower() == 'gpt2':
ds = [GPT2Dataset(d, max_seq_len=seq_length) if d is not None else None for d in ds] ds = [GPT2Dataset(d, max_seq_len=seq_length) if d is not None else None for d in ds]
else: else:
......
configure_data.py → megatron/deprecated_data_utils/configure_data.py
View file @ 22c0e300
# coding=utf-8 # coding=utf-8
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -21,6 +21,7 @@ import torch ...@@ -21,6 +21,7 @@ import torch
from megatron import data_utils from megatron import data_utils
from megatron import mpu from megatron import mpu
class DataConfig: class DataConfig:
def __init__(self, defaults={}): def __init__(self, defaults={}):
...@@ -48,7 +49,8 @@ def make_data_loader(dataset, batch_size, args): ...@@ -48,7 +49,8 @@ def make_data_loader(dataset, batch_size, args):
shuffle = args.shuffle shuffle = args.shuffle
if shuffle: if shuffle:
sampler = data_utils.samplers.RandomSampler(dataset, replacement=True, num_samples=batch_size*args.train_iters) sampler = data_utils.samplers.RandomSampler(
dataset, replacement=True, num_samples=batch_size * args.train_iters)
else: else:
sampler = torch.utils.data.SequentialSampler(dataset) sampler = torch.utils.data.SequentialSampler(dataset)
world_size = torch.distributed.get_world_size( world_size = torch.distributed.get_world_size(
...@@ -204,6 +206,7 @@ def make_loaders(args): ...@@ -204,6 +206,7 @@ def make_loaders(args):
return (train, valid, test), tokenizer return (train, valid, test), tokenizer
def get_split(args): def get_split(args):
""" """
Get dataset splits from comma separated string list Get dataset splits from comma separated string list
...@@ -217,7 +220,7 @@ def get_split(args): ...@@ -217,7 +220,7 @@ def get_split(args):
splits = [float(args.split)] splits = [float(args.split)]
split_total = sum(splits) split_total = sum(splits)
if split_total < 1.: if split_total < 1.:
splits.append(1-split_total) splits.append(1 - split_total)
while len(splits) < 3: while len(splits) < 3:
splits.append(0.) splits.append(0.)
splits = splits[:3] splits = splits[:3]
...@@ -226,10 +229,10 @@ def get_split(args): ...@@ -226,10 +229,10 @@ def get_split(args):
if args.test_data is not None: if args.test_data is not None:
splits[2] = 0. splits[2] = 0.
final_sum = sum(splits) final_sum = sum(splits)
return [s/final_sum for s in splits] return [s / final_sum for s in splits]
def configure_data():
def configure_data():
"""add cmdline flags for configuring datasets""" """add cmdline flags for configuring datasets"""
# These are options that are used by data_utils, but are either # These are options that are used by data_utils, but are either
# deprecated or not meant to be exposed to the command line user. # deprecated or not meant to be exposed to the command line user.
......
megatron/data_utils/corpora.py → megatron/deprecated_data_utils/corpora.py
View file @ 22c0e300
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -16,6 +16,7 @@ ...@@ -16,6 +16,7 @@
from .datasets import json_dataset, csv_dataset from .datasets import json_dataset, csv_dataset
import os import os
class wikipedia(json_dataset): class wikipedia(json_dataset):
""" """
dataset for wikipedia with arguments configured for convenience dataset for wikipedia with arguments configured for convenience
...@@ -24,6 +25,7 @@ class wikipedia(json_dataset): ...@@ -24,6 +25,7 @@ class wikipedia(json_dataset):
""" """
PATH = 'data/wikipedia/wikidump_lines.json' PATH = 'data/wikipedia/wikidump_lines.json'
assert_str = "make sure to set PATH for wikipedia data_utils/corpora.py" assert_str = "make sure to set PATH for wikipedia data_utils/corpora.py"
def __init__(self, **kwargs): def __init__(self, **kwargs):
assert os.path.exists(wikipedia.PATH), \ assert os.path.exists(wikipedia.PATH), \
wikipedia.assert_str wikipedia.assert_str
...@@ -42,6 +44,7 @@ class webtext(json_dataset): ...@@ -42,6 +44,7 @@ class webtext(json_dataset):
""" """
PATH = 'data/webtext/data.json' PATH = 'data/webtext/data.json'
assert_str = "make sure to set PATH for webtext data_utils/corpora.py" assert_str = "make sure to set PATH for webtext data_utils/corpora.py"
def __init__(self, **kwargs): def __init__(self, **kwargs):
assert os.path.exists(webtext.PATH), \ assert os.path.exists(webtext.PATH), \
webtext.assert_str webtext.assert_str
......
megatron/data_utils/datasets.py → megatron/deprecated_data_utils/datasets.py
View file @ 22c0e300
# coding=utf-8 # coding=utf-8
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -34,6 +34,7 @@ from nltk import tokenize ...@@ -34,6 +34,7 @@ from nltk import tokenize
from .lazy_loader import lazy_array_loader, exists_lazy, make_lazy from .lazy_loader import lazy_array_loader, exists_lazy, make_lazy
from .tokenization import Tokenization from .tokenization import Tokenization
class ConcatDataset(data.Dataset): class ConcatDataset(data.Dataset):
""" """
Dataset to concatenate multiple datasets. Dataset to concatenate multiple datasets.
...@@ -57,7 +58,8 @@ class ConcatDataset(data.Dataset): ...@@ -57,7 +58,8 @@ class ConcatDataset(data.Dataset):
super(ConcatDataset, self).__init__() super(ConcatDataset, self).__init__()
assert len(datasets) > 0, 'datasets should not be an empty iterable' assert len(datasets) > 0, 'datasets should not be an empty iterable'
self.datasets = list(datasets) self.datasets = list(datasets)
self.is_lazy = sum([isinstance(ds, lazy_array_loader) for ds in self.datasets]) == len(self.datasets) self.is_lazy = sum([isinstance(ds, lazy_array_loader)
for ds in self.datasets]) == len(self.datasets)
self.cumulative_sizes = self.cumsum(self.datasets) self.cumulative_sizes = self.cumsum(self.datasets)
self._X = None self._X = None
self._Y = None self._Y = None
...@@ -90,7 +92,8 @@ class ConcatDataset(data.Dataset): ...@@ -90,7 +92,8 @@ class ConcatDataset(data.Dataset):
self._lens.extend(data.lens) self._lens.extend(data.lens)
else: else:
for data in self.datasets: for data in self.datasets:
self._lens.extend([len(d['text']) if isinstance(d, dict) else len(d) for d in data]) self._lens.extend([len(d['text']) if isinstance(
d, dict) else len(d) for d in data])
return self._lens return self._lens
@property @property
...@@ -116,6 +119,7 @@ class ConcatDataset(data.Dataset): ...@@ -116,6 +119,7 @@ class ConcatDataset(data.Dataset):
"cumulative_sizes", DeprecationWarning, stacklevel=2) "cumulative_sizes", DeprecationWarning, stacklevel=2)
return self.cumulative_sizes return self.cumulative_sizes
class SplitDataset(data.Dataset): class SplitDataset(data.Dataset):
""" """
Dataset wrapper to access a subset of another dataset. Dataset wrapper to access a subset of another dataset.
...@@ -126,6 +130,7 @@ class SplitDataset(data.Dataset): ...@@ -126,6 +130,7 @@ class SplitDataset(data.Dataset):
ds (Dataset or array-like): List of datasets to be subindexed ds (Dataset or array-like): List of datasets to be subindexed
split_inds (1D array-like): List of indices part of subset split_inds (1D array-like): List of indices part of subset
""" """
def __init__(self, ds, split_inds, **kwargs): def __init__(self, ds, split_inds, **kwargs):
self.split_inds = list(split_inds) self.split_inds = list(split_inds)
self.wrapped_data = ds self.wrapped_data = ds
...@@ -163,7 +168,8 @@ class SplitDataset(data.Dataset): ...@@ -163,7 +168,8 @@ class SplitDataset(data.Dataset):
for idx in self.split_inds: for idx in self.split_inds:
yield self.wrapped_data[idx] yield self.wrapped_data[idx]
def split_ds(ds, split=[.8,.2,.0], shuffle=True):
def split_ds(ds, split=[.8, .2, .0], shuffle=True):
""" """
Split a dataset into subsets given proportions of how Split a dataset into subsets given proportions of how
much to allocate per split. If a split is 0% returns None for that split. much to allocate per split. If a split is 0% returns None for that split.
...@@ -184,18 +190,19 @@ def split_ds(ds, split=[.8,.2,.0], shuffle=True): ...@@ -184,18 +190,19 @@ def split_ds(ds, split=[.8,.2,.0], shuffle=True):
np.random.shuffle(inds) np.random.shuffle(inds)
start_idx = 0 start_idx = 0
residual_idx = 0 residual_idx = 0
rtn_ds = [None]*len(split) rtn_ds = [None] * len(split)
for i, f in enumerate(split): for i, f in enumerate(split):
if f != 0: if f != 0:
proportion = ds_len*split[i] proportion = ds_len * split[i]
residual_idx += proportion % 1 residual_idx += proportion % 1
split_ = int(int(proportion) + residual_idx) split_ = int(int(proportion) + residual_idx)
split_inds = inds[start_idx:start_idx+max(split_, 1)] split_inds = inds[start_idx:start_idx + max(split_, 1)]
rtn_ds[i] = SplitDataset(ds, split_inds) rtn_ds[i] = SplitDataset(ds, split_inds)
start_idx += split_ start_idx += split_
residual_idx %= 1 residual_idx %= 1
return rtn_ds return rtn_ds
class csv_dataset(data.Dataset): class csv_dataset(data.Dataset):
""" """
Class for loading datasets from csv files. Class for loading datasets from csv files.
...@@ -214,6 +221,7 @@ class csv_dataset(data.Dataset): ...@@ -214,6 +221,7 @@ class csv_dataset(data.Dataset):
X (list): all strings from the csv file X (list): all strings from the csv file
Y (np.ndarray): labels to train with Y (np.ndarray): labels to train with
""" """
def __init__(self, path, tokenizer=None, preprocess_fn=None, delim=',', def __init__(self, path, tokenizer=None, preprocess_fn=None, delim=',',
binarize_sent=False, drop_unlabeled=False, text_key='sentence', label_key='label', binarize_sent=False, drop_unlabeled=False, text_key='sentence', label_key='label',
**kwargs): **kwargs):
...@@ -229,7 +237,6 @@ class csv_dataset(data.Dataset): ...@@ -229,7 +237,6 @@ class csv_dataset(data.Dataset):
if '.tsv' in self.path: if '.tsv' in self.path:
self.delim = '\t' self.delim = '\t'
self.X = [] self.X = []
self.Y = [] self.Y = []
try: try:
...@@ -239,7 +246,7 @@ class csv_dataset(data.Dataset): ...@@ -239,7 +246,7 @@ class csv_dataset(data.Dataset):
else: else:
cols += [label_key] cols += [label_key]
data = pd.read_csv(self.path, sep=self.delim, usecols=cols, encoding='latin-1') data = pd.read_csv(self.path, sep=self.delim, usecols=cols, encoding='latin-1')
except: except BaseException:
data = pd.read_csv(self.path, sep=self.delim, usecols=[text_key], encoding='latin-1') data = pd.read_csv(self.path, sep=self.delim, usecols=[text_key], encoding='latin-1')
data = data.dropna(axis=0) data = data.dropna(axis=0)
...@@ -248,7 +255,7 @@ class csv_dataset(data.Dataset): ...@@ -248,7 +255,7 @@ class csv_dataset(data.Dataset):
try: try:
self.Y = data[label_key].values self.Y = data[label_key].values
except Exception as e: except Exception as e:
self.Y = np.ones(len(self.X))*-1 self.Y = np.ones(len(self.X)) * -1
if binarize_sent: if binarize_sent:
self.Y = binarize_labels(self.Y, hard=binarize_sent) self.Y = binarize_labels(self.Y, hard=binarize_sent)
...@@ -295,23 +302,25 @@ class csv_dataset(data.Dataset): ...@@ -295,23 +302,25 @@ class csv_dataset(data.Dataset):
write the metrics, text, and labels to a csv file write the metrics, text, and labels to a csv file
""" """
if path is None: if path is None:
path = self.path+'.results' path = self.path + '.results'
print('generating csv at ' + path) print('generating csv at ' + path)
with open(path, 'w') as csvfile: with open(path, 'w') as csvfile:
c = csv.writer(csvfile, delimiter=self.delim) c = csv.writer(csvfile, delimiter=self.delim)
if writer_gen is not None: if writer_gen is not None:
#if first item of generator is a header of what the metrics mean then write header to csv file # if first item of generator is a header of what the metrics mean then
# write header to csv file
if not skip_header: if not skip_header:
header = (self.label_key,)+tuple(next(writer_gen))+(self.text_key,) header = (self.label_key,) + tuple(next(writer_gen)) + (self.text_key,)
c.writerow(header) c.writerow(header)
for i, row in enumerate(writer_gen): for i, row in enumerate(writer_gen):
row = (self.Y[i],)+tuple(row)+(self.X[i],) row = (self.Y[i],) + tuple(row) + (self.X[i],)
c.writerow(row) c.writerow(row)
else: else:
c.writerow([self.label_key, self.text_key]) c.writerow([self.label_key, self.text_key])
for row in zip(self.Y, self.X): for row in zip(self.Y, self.X):
c.writerow(row) c.writerow(row)
class json_dataset(data.Dataset): class json_dataset(data.Dataset):
""" """
Class for loading datasets from a json dump. Class for loading datasets from a json dump.
...@@ -327,6 +336,7 @@ class json_dataset(data.Dataset): ...@@ -327,6 +336,7 @@ class json_dataset(data.Dataset):
all_strs (list): list of all strings from the dataset all_strs (list): list of all strings from the dataset
all_labels (list): list of all labels from the dataset (if they have it) all_labels (list): list of all labels from the dataset (if they have it)
""" """
def __init__(self, path, tokenizer=None, preprocess_fn=None, binarize_sent=False, def __init__(self, path, tokenizer=None, preprocess_fn=None, binarize_sent=False,
text_key='sentence', label_key='label', loose_json=False, **kwargs): text_key='sentence', label_key='label', loose_json=False, **kwargs):
self.is_lazy = False self.is_lazy = False
...@@ -389,24 +399,25 @@ class json_dataset(data.Dataset): ...@@ -389,24 +399,25 @@ class json_dataset(data.Dataset):
write the metrics, text, and labels to a json file write the metrics, text, and labels to a json file
""" """
if path is None: if path is None:
path = self.path+'.results' path = self.path + '.results'
jsons = [] jsons = []
if writer_gen is not None: if writer_gen is not None:
#if first item of generator is a header of what the metrics mean then write header to csv file # if first item of generator is a header of what the metrics mean then
# write header to csv file
def gen_helper(): def gen_helper():
keys = {} keys = {}
keys[0] = self.label_key keys[0] = self.label_key
if not skip_header: if not skip_header:
for idx, k in enumerate(tuple(next(writer_gen))): for idx, k in enumerate(tuple(next(writer_gen))):
keys[idx+1] = k keys[idx + 1] = k
for i, row in enumerate(writer_gen): for i, row in enumerate(writer_gen):
if i == 0 and skip_header: if i == 0 and skip_header:
for idx, _ in enumerate(row): for idx, _ in enumerate(row):
keys[idx+1] = 'metric_%d'%(idx,) keys[idx + 1] = 'metric_%d' % (idx,)
j = {} j = {}
for idx, v in enumerate((self.Y[i],)+tuple(row)): for idx, v in enumerate((self.Y[i],) + tuple(row)):
k = keys[idx] k = keys[idx]
j[k] = v j[k] = v
yield j yield j
...@@ -453,6 +464,7 @@ class json_dataset(data.Dataset): ...@@ -453,6 +464,7 @@ class json_dataset(data.Dataset):
j[self.label_key] = -1 j[self.label_key] = -1
yield j yield j
class GPT2Dataset(data.Dataset): class GPT2Dataset(data.Dataset):
def __init__(self, ds, def __init__(self, ds,
...@@ -503,7 +515,7 @@ class GPT2Dataset(data.Dataset): ...@@ -503,7 +515,7 @@ class GPT2Dataset(data.Dataset):
def __getitem__(self, idx): def __getitem__(self, idx):
# init rng # init rng
rng = random.Random(idx) rng = random.Random(idx)
rng = np.random.RandomState(seed=[rng.randint(0, 2**32-1) for _ in range(16)]) rng = np.random.RandomState(seed=[rng.randint(0, 2**32 - 1) for _ in range(16)])
# get possibly weighted random index from dataset # get possibly weighted random index from dataset
data_idx = self.get_weighted_samples(rng) data_idx = self.get_weighted_samples(rng)
...@@ -538,10 +550,10 @@ class GPT2Dataset(data.Dataset): ...@@ -538,10 +550,10 @@ class GPT2Dataset(data.Dataset):
else: else:
data_idx = (data_idx + 1) % self.ds_len data_idx = (data_idx + 1) % self.ds_len
tokens += self.getidx(data_idx) tokens += self.getidx(data_idx)
tokens = tokens[:(self.max_seq_len+1)] tokens = tokens[:(self.max_seq_len + 1)]
tokens = self.pad_seq(tokens) tokens = self.pad_seq(tokens)
return {'text': np.array(tokens),} return {'text': np.array(tokens), }
def getidx(self, data_idx): def getidx(self, data_idx):
data = self.ds[data_idx] data = self.ds[data_idx]
...@@ -556,7 +568,7 @@ class GPT2Dataset(data.Dataset): ...@@ -556,7 +568,7 @@ class GPT2Dataset(data.Dataset):
def pad_seq(self, seq): def pad_seq(self, seq):
total_tokens = self.max_seq_len + 1 total_tokens = self.max_seq_len + 1
num_pad_tokens = max(0, total_tokens - len(seq)) num_pad_tokens = max(0, total_tokens - len(seq))
seq += [self.tokenizer.get_command('pad').Id]*(num_pad_tokens) seq += [self.tokenizer.get_command('pad').Id] * (num_pad_tokens)
return seq return seq
def contains_sentence_end(self, tok): def contains_sentence_end(self, tok):
...@@ -569,6 +581,7 @@ class GPT2Dataset(data.Dataset): ...@@ -569,6 +581,7 @@ class GPT2Dataset(data.Dataset):
return True return True
return False return False
class bert_sentencepair_dataset(data.Dataset): class bert_sentencepair_dataset(data.Dataset):
""" """
Dataset containing sentencepairs for BERT training. Each index corresponds to a randomly generated sentence pair. Dataset containing sentencepairs for BERT training. Each index corresponds to a randomly generated sentence pair.
...@@ -581,7 +594,9 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -581,7 +594,9 @@ class bert_sentencepair_dataset(data.Dataset):
dataset_size (int): number of random sentencepairs in the dataset. Default: len(ds)*(len(ds)-1) dataset_size (int): number of random sentencepairs in the dataset. Default: len(ds)*(len(ds)-1)
""" """
def __init__(self, ds, max_seq_len=512, mask_lm_prob=.15, max_preds_per_seq=None, short_seq_prob=.01, dataset_size=None, presplit_sentences=False, weighted=True, **kwargs):
def __init__(self, ds, max_seq_len=512, mask_lm_prob=.15, max_preds_per_seq=None,
short_seq_prob=.01, dataset_size=None, presplit_sentences=False, weighted=True, **kwargs):
self.ds = ds self.ds = ds
self.ds_len = len(self.ds) self.ds_len = len(self.ds)
self.tokenizer = self.ds.GetTokenizer() self.tokenizer = self.ds.GetTokenizer()
...@@ -590,12 +605,12 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -590,12 +605,12 @@ class bert_sentencepair_dataset(data.Dataset):
self.max_seq_len = max_seq_len self.max_seq_len = max_seq_len
self.mask_lm_prob = mask_lm_prob self.mask_lm_prob = mask_lm_prob
if max_preds_per_seq is None: if max_preds_per_seq is None:
max_preds_per_seq = math.ceil(max_seq_len*mask_lm_prob /10)*10 max_preds_per_seq = math.ceil(max_seq_len * mask_lm_prob / 10) * 10
self.max_preds_per_seq = max_preds_per_seq self.max_preds_per_seq = max_preds_per_seq
self.short_seq_prob = short_seq_prob self.short_seq_prob = short_seq_prob
self.dataset_size = dataset_size self.dataset_size = dataset_size
if self.dataset_size is None: if self.dataset_size is None:
self.dataset_size = self.ds_len * (self.ds_len-1) self.dataset_size = self.ds_len * (self.ds_len - 1)
self.presplit_sentences = presplit_sentences self.presplit_sentences = presplit_sentences
if not self.presplit_sentences: if not self.presplit_sentences:
nltk.download('punkt', download_dir="./nltk") nltk.download('punkt', download_dir="./nltk")
...@@ -607,7 +622,8 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -607,7 +622,8 @@ class bert_sentencepair_dataset(data.Dataset):
if hasattr(self.ds, 'is_lazy') and self.ds.is_lazy: if hasattr(self.ds, 'is_lazy') and self.ds.is_lazy:
lens = np.array(self.ds.lens) lens = np.array(self.ds.lens)
else: else:
lens = np.array([len(d['text']) if isinstance(d, dict) else len(d) for d in self.ds]) lens = np.array([len(d['text']) if isinstance(d, dict) else len(d)
for d in self.ds])
self.total_len = np.sum(lens) self.total_len = np.sum(lens)
self.weighting = list(accumulate(lens)) self.weighting = list(accumulate(lens))
else: else:
...@@ -626,7 +642,7 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -626,7 +642,7 @@ class bert_sentencepair_dataset(data.Dataset):
def __getitem__(self, idx): def __getitem__(self, idx):
# get rng state corresponding to index (allows deterministic random pair) # get rng state corresponding to index (allows deterministic random pair)
rng = random.Random(idx) rng = random.Random(idx)
np_rng = np.random.RandomState(seed=[rng.randint(0, 2**32-1) for _ in range(16)]) np_rng = np.random.RandomState(seed=[rng.randint(0, 2**32 - 1) for _ in range(16)])
# get seq length # get seq length
target_seq_length = self.max_seq_len target_seq_length = self.max_seq_len
short_seq = False short_seq = False
...@@ -639,15 +655,25 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -639,15 +655,25 @@ class bert_sentencepair_dataset(data.Dataset):
lena = 0 lena = 0
lenb = 0 lenb = 0
while (is_random_next is None) or (lena < 1) or (lenb < 1): while (is_random_next is None) or (lena < 1) or (lenb < 1):
tokensa, tokensb, is_random_next = self.create_random_sentencepair(target_seq_length, rng, np_rng) tokensa, tokensb, is_random_next = self.create_random_sentencepair(
target_seq_length, rng, np_rng)
lena = len(tokensa[0]) lena = len(tokensa[0])
lenb = len(tokensb[0]) lenb = len(tokensb[0])
# truncate sentence pair to max_seq_len # truncate sentence pair to max_seq_len
tokensa, tokensb = self.truncate_seq_pair(tokensa, tokensb, self.max_seq_len, rng) tokensa, tokensb = self.truncate_seq_pair(tokensa, tokensb, self.max_seq_len, rng)
# join sentence pair, mask, and pad # join sentence pair, mask, and pad
tokens, mask, mask_labels, pad_mask = self.create_masked_lm_predictions(tokensa, tokensb, self.mask_lm_prob, self.max_preds_per_seq, self.vocab_words, rng) tokens, mask, mask_labels, pad_mask = self.create_masked_lm_predictions(
sample = {'text': np.array(tokens[0]), 'types': np.array(tokens[1]), 'is_random': int(is_random_next), 'mask': np.array(mask), 'mask_labels': np.array(mask_labels), 'pad_mask': np.array(pad_mask)} tokensa, tokensb, self.mask_lm_prob, self.max_preds_per_seq, self.vocab_words, rng)
sample = {
'text': np.array(
tokens[0]),
'types': np.array(
tokens[1]),
'is_random': int(is_random_next),
'mask': np.array(mask),
'mask_labels': np.array(mask_labels),
'pad_mask': np.array(pad_mask)}
return sample return sample
def sentence_split(self, document): def sentence_split(self, document):
...@@ -665,7 +691,7 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -665,7 +691,7 @@ class bert_sentencepair_dataset(data.Dataset):
"""tokenize sentence and get token types""" """tokenize sentence and get token types"""
tokens = self.tokenizer.EncodeAsIds(sent).tokenization tokens = self.tokenizer.EncodeAsIds(sent).tokenization
str_type = 'str' + str(sentence_num) str_type = 'str' + str(sentence_num)
token_types = [self.tokenizer.get_type(str_type).Id]*len(tokens) token_types = [self.tokenizer.get_type(str_type).Id] * len(tokens)
return tokens, token_types return tokens, token_types
def get_doc(self, idx): def get_doc(self, idx):
...@@ -694,21 +720,22 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -694,21 +720,22 @@ class bert_sentencepair_dataset(data.Dataset):
# doc_a_idx = np_rng.choice(self.ds_len, p=self.weighting) # doc_a_idx = np_rng.choice(self.ds_len, p=self.weighting)
doc_a_idx = self.get_weighted_samples(np_rng) doc_a_idx = self.get_weighted_samples(np_rng)
else: else:
doc_a_idx = rng.randint(0, self.ds_len-1) doc_a_idx = rng.randint(0, self.ds_len - 1)
doc_a = self.sentence_split(self.get_doc(doc_a_idx)) doc_a = self.sentence_split(self.get_doc(doc_a_idx))
if not doc_a: if not doc_a:
doc_a = None doc_a = None
random_start_a = rng.randint(0, len(doc_a)-1) random_start_a = rng.randint(0, len(doc_a) - 1)
while random_start_a < len(doc_a): while random_start_a < len(doc_a):
sentence = doc_a[random_start_a] sentence = doc_a[random_start_a]
sentence, sentence_types = self.sentence_tokenize(sentence, 0, random_start_a == 0, random_start_a == len(doc_a)) sentence, sentence_types = self.sentence_tokenize(
sentence, 0, random_start_a == 0, random_start_a == len(doc_a))
curr_strs.append(sentence) curr_strs.append(sentence)
curr_str_types.append(sentence_types) curr_str_types.append(sentence_types)
curr_len += len(sentence) curr_len += len(sentence)
if random_start_a == len(doc_a) - 1 or curr_len >= target_seq_length: if random_start_a == len(doc_a) - 1 or curr_len >= target_seq_length:
break break
random_start_a = (random_start_a+1) random_start_a = (random_start_a + 1)
if curr_strs: if curr_strs:
num_a = 1 num_a = 1
...@@ -738,16 +765,17 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -738,16 +765,17 @@ class bert_sentencepair_dataset(data.Dataset):
if not doc_b: if not doc_b:
doc_b = None doc_b = None
random_start_b = rng.randint(0, len(doc_b)-1) random_start_b = rng.randint(0, len(doc_b) - 1)
while random_start_b < len(doc_b): while random_start_b < len(doc_b):
sentence_b = doc_b[random_start_b] sentence_b = doc_b[random_start_b]
new_b_tokens, new_b_types = self.sentence_tokenize(sentence_b, 1, random_start_b == 0, random_start_b == len(doc_b)) new_b_tokens, new_b_types = self.sentence_tokenize(
sentence_b, 1, random_start_b == 0, random_start_b == len(doc_b))
b_len += len(new_b_tokens) b_len += len(new_b_tokens)
tokens_b.extend(new_b_tokens) tokens_b.extend(new_b_tokens)
token_types_b.extend(new_b_types) token_types_b.extend(new_b_types)
if len(tokens_b) >= target_b_length: if len(tokens_b) >= target_b_length:
break break
random_start_b = (random_start_b+1) random_start_b = (random_start_b + 1)
else: else:
is_random_next = False is_random_next = False
for j in range(num_a, len(curr_strs)): for j in range(num_a, len(curr_strs)):
...@@ -817,8 +845,10 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -817,8 +845,10 @@ class bert_sentencepair_dataset(data.Dataset):
return seq, pad_mask return seq, pad_mask
def concat_tokens(self, tokens_a, token_types_a, tokens_b, token_types_b): def concat_tokens(self, tokens_a, token_types_a, tokens_b, token_types_b):
tokens = [self.tokenizer.get_command('ENC').Id] + tokens_a + [self.tokenizer.get_command('sep').Id] + tokens_b + [self.tokenizer.get_command('sep').Id] tokens = [self.tokenizer.get_command('ENC').Id] + tokens_a + [self.tokenizer.get_command(
token_types = [token_types_a[0]] + token_types_a + [token_types_a[0]] + token_types_b + [token_types_b[0]] 'sep').Id] + tokens_b + [self.tokenizer.get_command('sep').Id]
token_types = [token_types_a[0]] + token_types_a + \
[token_types_a[0]] + token_types_b + [token_types_b[0]]
return tokens, token_types return tokens, token_types
def create_masked_lm_predictions(self, a, b, mask_lm_prob, max_preds_per_seq, vocab_words, rng): def create_masked_lm_predictions(self, a, b, mask_lm_prob, max_preds_per_seq, vocab_words, rng):
...@@ -833,7 +863,7 @@ class bert_sentencepair_dataset(data.Dataset): ...@@ -833,7 +863,7 @@ class bert_sentencepair_dataset(data.Dataset):
len_a = len(tokens_a) len_a = len(tokens_a)
len_b = len(tokens_b) len_b = len(tokens_b)
cand_indices = [idx+1 for idx in range(len_a)] + [idx+2+len_a for idx in range(len_b)] cand_indices = [idx + 1 for idx in range(len_a)] + [idx + 2 + len_a for idx in range(len_b)]
rng.shuffle(cand_indices) rng.shuffle(cand_indices)
......
megatron/data_utils/samplers.py → megatron/deprecated_data_utils/samplers.py
View file @ 22c0e300
# coding=utf-8 # coding=utf-8
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
...@@ -21,6 +21,7 @@ import torch ...@@ -21,6 +21,7 @@ import torch
from torch.utils import data from torch.utils import data
import numpy as np import numpy as np
class RandomSampler(data.sampler.Sampler): class RandomSampler(data.sampler.Sampler):
r""" r"""
Based off of pytorch RandomSampler and DistributedSampler. Essentially a RandomSampler, Based off of pytorch RandomSampler and DistributedSampler. Essentially a RandomSampler,
...@@ -63,7 +64,8 @@ class RandomSampler(data.sampler.Sampler): ...@@ -63,7 +64,8 @@ class RandomSampler(data.sampler.Sampler):
if self.epoch >= 0: if self.epoch >= 0:
g.manual_seed(self.epoch) g.manual_seed(self.epoch)
if self.replacement: if self.replacement:
return iter(torch.randint(high=n, size=(self.num_samples,), dtype=torch.int64, generator=g).tolist()) return iter(torch.randint(high=n, size=(self.num_samples,),
dtype=torch.int64, generator=g).tolist())
return iter(torch.randperm(n, generator=g).tolist()) return iter(torch.randperm(n, generator=g).tolist())
def __len__(self): def __len__(self):
...@@ -72,12 +74,14 @@ class RandomSampler(data.sampler.Sampler): ...@@ -72,12 +74,14 @@ class RandomSampler(data.sampler.Sampler):
def set_epoch(self, epoch): def set_epoch(self, epoch):
self.epoch = epoch self.epoch = epoch
class DistributedBatchSampler(data.sampler.BatchSampler): class DistributedBatchSampler(data.sampler.BatchSampler):
""" """
similar to normal implementation of distributed sampler, except implementation is at the similar to normal implementation of distributed sampler, except implementation is at the
batch sampler level, instead of just the sampler level. This allows wrapping of arbitrary batch sampler level, instead of just the sampler level. This allows wrapping of arbitrary
data samplers (sequential, random, WeightedRandomSampler, etc.) with this batch sampler. data samplers (sequential, random, WeightedRandomSampler, etc.) with this batch sampler.
""" """
def __init__(self, sampler, batch_size, drop_last, rank=-1, world_size=2, wrap_last=False): def __init__(self, sampler, batch_size, drop_last, rank=-1, world_size=2, wrap_last=False):
super(DistributedBatchSampler, self).__init__(sampler, batch_size, drop_last) super(DistributedBatchSampler, self).__init__(sampler, batch_size, drop_last)
if rank == -1: if rank == -1:
...@@ -125,7 +129,7 @@ class DistributedBatchSampler(data.sampler.BatchSampler): ...@@ -125,7 +129,7 @@ class DistributedBatchSampler(data.sampler.BatchSampler):
def data_iterator(self, _iter, wrap_around=False): def data_iterator(self, _iter, wrap_around=False):
"""iterates through data and handles wrap around""" """iterates through data and handles wrap around"""
for i, idx in enumerate(_iter): for i, idx in enumerate(_iter):
if i < self.wrap_around%self.batch_size: if i < self.wrap_around % self.batch_size:
continue continue
if wrap_around: if wrap_around:
self.wrap_around += 1 self.wrap_around += 1
...@@ -134,6 +138,6 @@ class DistributedBatchSampler(data.sampler.BatchSampler): ...@@ -134,6 +138,6 @@ class DistributedBatchSampler(data.sampler.BatchSampler):
def _batch(self, batch): def _batch(self, batch):
"""extracts samples only pertaining to this worker's batch""" """extracts samples only pertaining to this worker's batch"""
start = self.rank*self.batch_size//self.world_size start = self.rank * self.batch_size // self.world_size
end = (self.rank+1)*self.batch_size//self.world_size end = (self.rank + 1) * self.batch_size // self.world_size
return batch[start:end] return batch[start:end]
scripts/presplit_sentences_json.py → megatron/deprecated_data_utils/scripts/presplit_sentences_json.py
View file @ 22c0e300
...@@ -24,4 +24,4 @@ with open(input_file, 'r') as ifile: ...@@ -24,4 +24,4 @@ with open(input_file, 'r') as ifile:
if line != '\n': if line != '\n':
sent_list.extend(nltk.tokenize.sent_tokenize(line)) sent_list.extend(nltk.tokenize.sent_tokenize(line))
parsed['text'] = line_seperator.join(sent_list) parsed['text'] = line_seperator.join(sent_list)
ofile.write(json.dumps(parsed)+'\n') ofile.write(json.dumps(parsed) + '\n')
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment