mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-19 07:54:56 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
+466
@@ -0,0 +1,466 @@
|
||||
# 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.
|
||||
|
||||
__all__ = ['BertPunctuationCapitalizationDataset', 'BertPunctuationCapitalizationInferDataset']
|
||||
|
||||
import itertools
|
||||
import os
|
||||
import pickle
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
|
||||
from nemo.collections.nlp.data.data_utils.data_preprocessing import get_label_stats, get_stats
|
||||
from nemo.core.classes import Dataset
|
||||
from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def get_features(
|
||||
queries: List[str],
|
||||
max_seq_length: int,
|
||||
tokenizer: TokenizerSpec,
|
||||
punct_label_ids: dict = None,
|
||||
capit_label_ids: dict = None,
|
||||
pad_label: str = 'O',
|
||||
punct_labels_lines=None,
|
||||
capit_labels_lines=None,
|
||||
ignore_extra_tokens=False,
|
||||
ignore_start_end: Optional[bool] = False,
|
||||
):
|
||||
"""
|
||||
Processes the data and returns features.
|
||||
|
||||
Args:
|
||||
queries: text sequences
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
tokenizer: such as AutoTokenizer
|
||||
pad_label: pad value use for labels. By default, it's the neutral label.
|
||||
punct_label_ids: dict to map punctuation labels to label ids.
|
||||
Starts with pad_label->0 and then increases in alphabetical order.
|
||||
Required for training and evaluation, not needed for inference.
|
||||
capit_label_ids: dict to map labels to label ids. Starts
|
||||
with pad_label->0 and then increases in alphabetical order.
|
||||
Required for training and evaluation, not needed for inference.
|
||||
punct_labels: list of labels for every word in a sequence (str)
|
||||
capit_labels: list of labels for every word in a sequence (str)
|
||||
ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
|
||||
ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
|
||||
|
||||
Returns:
|
||||
all_input_ids: input ids for all tokens
|
||||
all_segment_ids: token type ids
|
||||
all_input_mask: attention mask to use for BERT model
|
||||
all_subtokens_mask: masks out all subwords besides the first one
|
||||
all_loss_mask: loss mask to mask out tokens during training
|
||||
punct_all_labels: all labels for punctuation task (ints)
|
||||
capit_all_labels: all labels for capitalization task (ints)
|
||||
punct_label_ids: label (str) to id (int) map for punctuation task
|
||||
capit_label_ids: label (str) to id (int) map for capitalization task
|
||||
"""
|
||||
all_subtokens = []
|
||||
all_loss_mask = []
|
||||
all_subtokens_mask = []
|
||||
all_segment_ids = []
|
||||
all_input_ids = []
|
||||
all_input_mask = []
|
||||
sent_lengths = []
|
||||
punct_all_labels = []
|
||||
capit_all_labels = []
|
||||
with_label = False
|
||||
|
||||
if punct_labels_lines and capit_labels_lines:
|
||||
with_label = True
|
||||
|
||||
for i, query in enumerate(queries):
|
||||
words = query.strip().split()
|
||||
|
||||
# add bos token
|
||||
subtokens = [tokenizer.cls_token]
|
||||
loss_mask = [1 - ignore_start_end]
|
||||
subtokens_mask = [0]
|
||||
if with_label:
|
||||
pad_id = punct_label_ids[pad_label]
|
||||
punct_labels = [pad_id]
|
||||
punct_query_labels = [punct_label_ids[lab] for lab in punct_labels_lines[i]]
|
||||
|
||||
capit_labels = [pad_id]
|
||||
capit_query_labels = [capit_label_ids[lab] for lab in capit_labels_lines[i]]
|
||||
|
||||
for j, word in enumerate(words):
|
||||
word_tokens = tokenizer.text_to_tokens(word)
|
||||
subtokens.extend(word_tokens)
|
||||
|
||||
loss_mask.append(1)
|
||||
loss_mask.extend([int(not ignore_extra_tokens)] * (len(word_tokens) - 1))
|
||||
|
||||
subtokens_mask.append(1)
|
||||
subtokens_mask.extend([0] * (len(word_tokens) - 1))
|
||||
|
||||
if with_label:
|
||||
punct_labels.extend([punct_query_labels[j]] * len(word_tokens))
|
||||
capit_labels.extend([capit_query_labels[j]] * len(word_tokens))
|
||||
|
||||
# add eos token
|
||||
subtokens.append(tokenizer.sep_token)
|
||||
loss_mask.append(1 - ignore_start_end)
|
||||
subtokens_mask.append(0)
|
||||
sent_lengths.append(len(subtokens))
|
||||
all_subtokens.append(subtokens)
|
||||
all_loss_mask.append(loss_mask)
|
||||
all_subtokens_mask.append(subtokens_mask)
|
||||
all_input_mask.append([1] * len(subtokens))
|
||||
|
||||
if with_label:
|
||||
punct_labels.append(pad_id)
|
||||
punct_all_labels.append(punct_labels)
|
||||
capit_labels.append(pad_id)
|
||||
capit_all_labels.append(capit_labels)
|
||||
|
||||
max_seq_length = min(max_seq_length, max(sent_lengths))
|
||||
logging.info(f'Max length: {max_seq_length}')
|
||||
get_stats(sent_lengths)
|
||||
too_long_count = 0
|
||||
|
||||
for i, subtokens in enumerate(all_subtokens):
|
||||
if len(subtokens) > max_seq_length:
|
||||
subtokens = [tokenizer.cls_token] + subtokens[-max_seq_length + 1 :]
|
||||
all_input_mask[i] = [1] + all_input_mask[i][-max_seq_length + 1 :]
|
||||
all_loss_mask[i] = [int(not ignore_start_end)] + all_loss_mask[i][-max_seq_length + 1 :]
|
||||
all_subtokens_mask[i] = [0] + all_subtokens_mask[i][-max_seq_length + 1 :]
|
||||
|
||||
if with_label:
|
||||
punct_all_labels[i] = [pad_id] + punct_all_labels[i][-max_seq_length + 1 :]
|
||||
capit_all_labels[i] = [pad_id] + capit_all_labels[i][-max_seq_length + 1 :]
|
||||
too_long_count += 1
|
||||
|
||||
all_input_ids.append(tokenizer.tokens_to_ids(subtokens))
|
||||
|
||||
if len(subtokens) < max_seq_length:
|
||||
extra = max_seq_length - len(subtokens)
|
||||
all_input_ids[i] = all_input_ids[i] + [0] * extra
|
||||
all_loss_mask[i] = all_loss_mask[i] + [0] * extra
|
||||
all_subtokens_mask[i] = all_subtokens_mask[i] + [0] * extra
|
||||
all_input_mask[i] = all_input_mask[i] + [0] * extra
|
||||
|
||||
if with_label:
|
||||
punct_all_labels[i] = punct_all_labels[i] + [pad_id] * extra
|
||||
capit_all_labels[i] = capit_all_labels[i] + [pad_id] * extra
|
||||
|
||||
all_segment_ids.append([0] * max_seq_length)
|
||||
|
||||
logging.info(f'{too_long_count} are longer than {max_seq_length}')
|
||||
|
||||
for i in range(min(len(all_input_ids), 5)):
|
||||
logging.info("*** Example ***")
|
||||
logging.info("i: %s" % (i))
|
||||
logging.info("subtokens: %s" % " ".join(list(map(str, all_subtokens[i]))))
|
||||
logging.info("loss_mask: %s" % " ".join(list(map(str, all_loss_mask[i]))))
|
||||
logging.info("input_mask: %s" % " ".join(list(map(str, all_input_mask[i]))))
|
||||
logging.info("subtokens_mask: %s" % " ".join(list(map(str, all_subtokens_mask[i]))))
|
||||
if with_label:
|
||||
logging.info("punct_labels: %s" % " ".join(list(map(str, punct_all_labels[i]))))
|
||||
logging.info("capit_labels: %s" % " ".join(list(map(str, capit_all_labels[i]))))
|
||||
|
||||
return (
|
||||
all_input_ids,
|
||||
all_segment_ids,
|
||||
all_input_mask,
|
||||
all_subtokens_mask,
|
||||
all_loss_mask,
|
||||
punct_all_labels,
|
||||
capit_all_labels,
|
||||
punct_label_ids,
|
||||
capit_label_ids,
|
||||
)
|
||||
|
||||
|
||||
class BertPunctuationCapitalizationDataset(Dataset):
|
||||
"""
|
||||
Creates dataset to use during training for punctuaion and capitalization tasks with a pretrained model.
|
||||
For dataset to use during inference without labels, see BertPunctuationCapitalizationInferDataset.
|
||||
|
||||
Args:
|
||||
text_file: file to sequences, each line should a sentence, no header.
|
||||
label_file: file to labels, each line corresponds to word labels for a sentence in the text_file. No header.
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
tokenizer: such as AutoTokenizer
|
||||
num_samples: number of samples you want to use for the dataset.
|
||||
If -1, use all dataset. Useful for testing.
|
||||
pad_label: pad value use for labels.
|
||||
by default, it's the neutral label.
|
||||
punct_label_ids and capit_label_ids (dict):
|
||||
dict to map labels to label ids.
|
||||
Starts with pad_label->0 and then increases in alphabetical order
|
||||
For dev set use label_ids generated during training to support
|
||||
cases when not all labels are present in the dev set.
|
||||
For training set label_ids should be None or loaded from cache
|
||||
ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
|
||||
ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
|
||||
use_cache: whether to use processed data cache or not
|
||||
get_label_frequencies: whether to generate label frequencies
|
||||
punct_label_ids_file and capit_label_ids_file: name of the files to save in .nemo
|
||||
"""
|
||||
|
||||
@property
|
||||
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
||||
"""Returns definitions of module output ports. """
|
||||
return {
|
||||
'input_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'segment_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'input_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'loss_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'punct_labels': NeuralType(('B', 'T'), LabelsType()),
|
||||
'capit_labels': NeuralType(('B', 'T'), LabelsType()),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_file: str,
|
||||
label_file: str,
|
||||
max_seq_length: int,
|
||||
tokenizer: TokenizerSpec,
|
||||
num_samples: int = -1,
|
||||
pad_label: str = 'O',
|
||||
punct_label_ids: Dict[str, int] = None,
|
||||
capit_label_ids: Dict[str, int] = None,
|
||||
ignore_extra_tokens: bool = False,
|
||||
ignore_start_end: bool = False,
|
||||
use_cache: bool = True,
|
||||
get_label_frequencies: bool = False,
|
||||
punct_label_ids_file: str = 'punct_label_ids.csv',
|
||||
capit_label_ids_file: str = 'capit_label_ids.csv',
|
||||
):
|
||||
""" Initializes BertPunctuationCapitalizationDataset. """
|
||||
|
||||
if not (os.path.exists(text_file) and os.path.exists(label_file)):
|
||||
raise FileNotFoundError(
|
||||
f'{text_file} or {label_file} not found. The data should be splitted into 2 files: text.txt and \
|
||||
labels.txt. Each line of the text.txt file contains text sequences, where words are separated with \
|
||||
spaces. The labels.txt file contains corresponding labels for each word in text.txt, the labels are \
|
||||
separated with spaces. Each line of the files should follow the format: \
|
||||
[WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and \
|
||||
[LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).'
|
||||
)
|
||||
|
||||
# Cache features
|
||||
data_dir = os.path.dirname(text_file)
|
||||
filename = os.path.basename(text_file)
|
||||
|
||||
if not filename.endswith('.txt'):
|
||||
raise ValueError("{text_file} should have extension .txt")
|
||||
|
||||
filename = filename[:-4]
|
||||
vocab_size = getattr(tokenizer, "vocab_size", 0)
|
||||
features_pkl = os.path.join(
|
||||
data_dir,
|
||||
"cached_{}_{}_{}_{}_{}".format(
|
||||
filename, tokenizer.name, str(max_seq_length), str(vocab_size), str(num_samples)
|
||||
),
|
||||
)
|
||||
|
||||
self.punct_label_ids_file = os.path.join(data_dir, punct_label_ids_file)
|
||||
self.capit_label_ids_file = os.path.join(data_dir, capit_label_ids_file)
|
||||
|
||||
master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
|
||||
cache_files_exist = (
|
||||
os.path.exists(features_pkl)
|
||||
and os.path.exists(self.punct_label_ids_file)
|
||||
and os.path.exists(self.capit_label_ids_file)
|
||||
)
|
||||
features = None
|
||||
if master_device and not (cache_files_exist and use_cache):
|
||||
if num_samples == 0:
|
||||
raise ValueError("num_samples has to be positive", num_samples)
|
||||
logging.info(f'Processing {text_file}')
|
||||
with open(text_file, 'r') as f:
|
||||
text_lines = f.readlines()
|
||||
|
||||
# Collect all possible labels
|
||||
punct_unique_labels = set()
|
||||
capit_unique_labels = set()
|
||||
punct_labels_lines = []
|
||||
capit_labels_lines = []
|
||||
with open(label_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip().split()
|
||||
|
||||
# extract punctuation and capitalization labels
|
||||
punct_line, capit_line = zip(*line)
|
||||
punct_labels_lines.append(punct_line)
|
||||
capit_labels_lines.append(capit_line)
|
||||
|
||||
punct_unique_labels.update(punct_line)
|
||||
capit_unique_labels.update(capit_line)
|
||||
|
||||
if len(punct_labels_lines) != len(text_lines):
|
||||
raise ValueError("Labels file should contain labels for every word")
|
||||
|
||||
dataset = list(zip(text_lines, punct_labels_lines, capit_labels_lines))
|
||||
|
||||
if num_samples > 0:
|
||||
dataset = dataset[:num_samples]
|
||||
|
||||
dataset = list(zip(*dataset))
|
||||
text_lines = dataset[0]
|
||||
punct_labels_lines = dataset[1]
|
||||
capit_labels_lines = dataset[2]
|
||||
|
||||
# for dev/test sets use label mapping from training set
|
||||
if punct_label_ids:
|
||||
if len(punct_label_ids) != len(punct_unique_labels):
|
||||
logging.info(
|
||||
'Not all labels from the specified'
|
||||
+ 'label_ids dictionary are present in the'
|
||||
+ 'current dataset. Using the provided'
|
||||
+ 'label_ids dictionary.'
|
||||
)
|
||||
else:
|
||||
logging.info('Using the provided label_ids dictionary.')
|
||||
else:
|
||||
logging.info(
|
||||
'Creating a new label to label_id dictionary.'
|
||||
+ ' It\'s recommended to use label_ids generated'
|
||||
+ ' during training for dev/test sets to avoid'
|
||||
+ ' errors if some labels are not'
|
||||
+ ' present in the dev/test sets.'
|
||||
+ ' For training set label_ids should be None.'
|
||||
)
|
||||
|
||||
def create_label_ids(unique_labels, pad_label=pad_label):
|
||||
label_ids = {pad_label: 0}
|
||||
if pad_label in unique_labels:
|
||||
unique_labels.remove(pad_label)
|
||||
for label in sorted(unique_labels):
|
||||
label_ids[label] = len(label_ids)
|
||||
return label_ids
|
||||
|
||||
punct_label_ids = create_label_ids(punct_unique_labels)
|
||||
capit_label_ids = create_label_ids(capit_unique_labels)
|
||||
|
||||
self._save_label_ids(punct_label_ids, self.punct_label_ids_file)
|
||||
self._save_label_ids(capit_label_ids, self.capit_label_ids_file)
|
||||
|
||||
features = get_features(
|
||||
text_lines,
|
||||
max_seq_length,
|
||||
tokenizer,
|
||||
pad_label=pad_label,
|
||||
punct_labels_lines=punct_labels_lines,
|
||||
capit_labels_lines=capit_labels_lines,
|
||||
punct_label_ids=punct_label_ids,
|
||||
capit_label_ids=capit_label_ids,
|
||||
ignore_extra_tokens=ignore_extra_tokens,
|
||||
ignore_start_end=ignore_start_end,
|
||||
)
|
||||
|
||||
pickle.dump(features, open(features_pkl, "wb"))
|
||||
logging.info(f'Features saved to {features_pkl}')
|
||||
|
||||
# wait until the master process writes to the processed data files
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
if features is None:
|
||||
features = pickle.load(open(features_pkl, 'rb'))
|
||||
logging.info(f'Features restored from {features_pkl}')
|
||||
|
||||
self.all_input_ids = features[0]
|
||||
self.all_segment_ids = features[1]
|
||||
self.all_input_mask = features[2]
|
||||
self.all_subtokens_mask = features[3]
|
||||
self.all_loss_mask = features[4]
|
||||
self.punct_all_labels = features[5]
|
||||
self.capit_all_labels = features[6]
|
||||
self.punct_label_ids = features[7]
|
||||
self.capit_label_ids = features[8]
|
||||
|
||||
if get_label_frequencies:
|
||||
self.punct_label_frequencies = self._calculate_label_frequencies(self.punct_all_labels, data_dir, 'punct')
|
||||
self.capit_label_frequencies = self._calculate_label_frequencies(self.capit_all_labels, data_dir, 'capit')
|
||||
|
||||
def _calculate_label_frequencies(self, all_labels: List[int], data_dir: str, name: str) -> Dict[str, float]:
|
||||
""" Calculates labels frequencies """
|
||||
merged_labels = itertools.chain.from_iterable(all_labels)
|
||||
logging.info('Three most popular labels')
|
||||
_, label_frequencies, _ = get_label_stats(merged_labels, data_dir + '/label_count_' + name + '.tsv')
|
||||
return label_frequencies
|
||||
|
||||
def _save_label_ids(self, label_ids: Dict[str, int], filename: str) -> None:
|
||||
""" Saves label ids map to a file """
|
||||
with open(filename, 'w') as out:
|
||||
labels, _ = zip(*sorted(label_ids.items(), key=lambda x: x[1]))
|
||||
out.write('\n'.join(labels))
|
||||
logging.info(f'Labels: {label_ids}')
|
||||
logging.info(f'Labels mapping saved to : {out.name}')
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_input_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return (
|
||||
np.array(self.all_input_ids[idx]),
|
||||
np.array(self.all_segment_ids[idx]),
|
||||
np.array(self.all_input_mask[idx], dtype=np.long),
|
||||
np.array(self.all_subtokens_mask[idx]),
|
||||
np.array(self.all_loss_mask[idx]),
|
||||
np.array(self.punct_all_labels[idx]),
|
||||
np.array(self.capit_all_labels[idx]),
|
||||
)
|
||||
|
||||
|
||||
class BertPunctuationCapitalizationInferDataset(Dataset):
|
||||
"""
|
||||
Creates dataset to use during inference for punctuation and capitalization tasks with a pretrained model.
|
||||
For dataset to use during training with labels, see BertPunctuationCapitalizationDataset.
|
||||
|
||||
Args:
|
||||
queries file to sequences, each line should a sentence, no header.
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
tokenizer: such as AutoTokenizer
|
||||
"""
|
||||
|
||||
@property
|
||||
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {
|
||||
'input_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'segment_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'input_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
}
|
||||
|
||||
def __init__(self, queries: List[str], max_seq_length: int, tokenizer: TokenizerSpec):
|
||||
""" Initializes BertPunctuationCapitalizationInferDataset. """
|
||||
features = get_features(queries=queries, max_seq_length=max_seq_length, tokenizer=tokenizer)
|
||||
|
||||
self.all_input_ids = features[0]
|
||||
self.all_segment_ids = features[1]
|
||||
self.all_input_mask = features[2]
|
||||
self.all_subtokens_mask = features[3]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_input_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return (
|
||||
np.array(self.all_input_ids[idx]),
|
||||
np.array(self.all_segment_ids[idx]),
|
||||
np.array(self.all_input_mask[idx], dtype=np.float32),
|
||||
np.array(self.all_subtokens_mask[idx]),
|
||||
)
|
||||
@@ -0,0 +1,338 @@
|
||||
# Copyright 2018 The Google AI Language Team Authors and
|
||||
# The HuggingFace Inc. team.
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Utility functions for Token Classification NLP tasks
|
||||
Some parts of this code were adapted from the HuggingFace library at
|
||||
https://github.com/huggingface/pytorch-pretrained-BERT
|
||||
"""
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
|
||||
from nemo.collections.nlp.data.data_utils.data_preprocessing import get_stats
|
||||
from nemo.core.classes import Dataset
|
||||
from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
|
||||
from nemo.utils import logging
|
||||
|
||||
__all__ = ['BertTokenClassificationDataset', 'BertTokenClassificationInferDataset']
|
||||
|
||||
|
||||
def get_features(
|
||||
queries: List[str],
|
||||
tokenizer: TokenizerSpec,
|
||||
max_seq_length: int = -1,
|
||||
label_ids: dict = None,
|
||||
pad_label: str = 'O',
|
||||
raw_labels: List[str] = None,
|
||||
ignore_extra_tokens: bool = False,
|
||||
ignore_start_end: bool = False,
|
||||
):
|
||||
"""
|
||||
Processes the data and returns features.
|
||||
Args:
|
||||
queries: text sequences
|
||||
tokenizer: such as AutoTokenizer
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP], when -1 - use the max len from the data
|
||||
pad_label: pad value use for labels. By default, it's the neutral label.
|
||||
raw_labels: list of labels for every word in a sequence
|
||||
label_ids: dict to map labels to label ids.
|
||||
Starts with pad_label->0 and then increases in alphabetical order.
|
||||
Required for training and evaluation, not needed for inference.
|
||||
ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
|
||||
ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
|
||||
"""
|
||||
all_subtokens = []
|
||||
all_loss_mask = []
|
||||
all_subtokens_mask = []
|
||||
all_segment_ids = []
|
||||
all_input_ids = []
|
||||
all_input_mask = []
|
||||
sent_lengths = []
|
||||
all_labels = []
|
||||
with_label = False
|
||||
|
||||
if raw_labels is not None:
|
||||
with_label = True
|
||||
|
||||
for i, query in enumerate(queries):
|
||||
words = query.strip().split()
|
||||
|
||||
# add bos token
|
||||
subtokens = [tokenizer.cls_token]
|
||||
loss_mask = [1 - ignore_start_end]
|
||||
subtokens_mask = [0]
|
||||
if with_label:
|
||||
pad_id = label_ids[pad_label]
|
||||
labels = [pad_id]
|
||||
query_labels = [label_ids[lab] for lab in raw_labels[i]]
|
||||
|
||||
for j, word in enumerate(words):
|
||||
word_tokens = tokenizer.text_to_tokens(word)
|
||||
subtokens.extend(word_tokens)
|
||||
|
||||
loss_mask.append(1)
|
||||
loss_mask.extend([int(not ignore_extra_tokens)] * (len(word_tokens) - 1))
|
||||
|
||||
subtokens_mask.append(1)
|
||||
subtokens_mask.extend([0] * (len(word_tokens) - 1))
|
||||
|
||||
if with_label:
|
||||
labels.extend([query_labels[j]] * len(word_tokens))
|
||||
# add eos token
|
||||
subtokens.append(tokenizer.sep_token)
|
||||
loss_mask.append(1 - ignore_start_end)
|
||||
subtokens_mask.append(0)
|
||||
sent_lengths.append(len(subtokens))
|
||||
all_subtokens.append(subtokens)
|
||||
all_loss_mask.append(loss_mask)
|
||||
all_subtokens_mask.append(subtokens_mask)
|
||||
all_input_mask.append([1] * len(subtokens))
|
||||
|
||||
if with_label:
|
||||
labels.append(pad_id)
|
||||
all_labels.append(labels)
|
||||
|
||||
max_seq_length_data = max(sent_lengths)
|
||||
max_seq_length = min(max_seq_length, max_seq_length_data) if max_seq_length > 0 else max_seq_length_data
|
||||
logging.info(f'Setting Max Seq length to: {max_seq_length}')
|
||||
get_stats(sent_lengths)
|
||||
too_long_count = 0
|
||||
|
||||
for i, subtokens in enumerate(all_subtokens):
|
||||
if len(subtokens) > max_seq_length:
|
||||
subtokens = [tokenizer.cls_token] + subtokens[-max_seq_length + 1 :]
|
||||
all_input_mask[i] = [1] + all_input_mask[i][-max_seq_length + 1 :]
|
||||
all_loss_mask[i] = [int(not ignore_start_end)] + all_loss_mask[i][-max_seq_length + 1 :]
|
||||
all_subtokens_mask[i] = [0] + all_subtokens_mask[i][-max_seq_length + 1 :]
|
||||
|
||||
if with_label:
|
||||
all_labels[i] = [pad_id] + all_labels[i][-max_seq_length + 1 :]
|
||||
too_long_count += 1
|
||||
|
||||
all_input_ids.append(tokenizer.tokens_to_ids(subtokens))
|
||||
|
||||
if len(subtokens) < max_seq_length:
|
||||
extra = max_seq_length - len(subtokens)
|
||||
all_input_ids[i] = all_input_ids[i] + [0] * extra
|
||||
all_loss_mask[i] = all_loss_mask[i] + [0] * extra
|
||||
all_subtokens_mask[i] = all_subtokens_mask[i] + [0] * extra
|
||||
all_input_mask[i] = all_input_mask[i] + [0] * extra
|
||||
|
||||
if with_label:
|
||||
all_labels[i] = all_labels[i] + [pad_id] * extra
|
||||
|
||||
all_segment_ids.append([0] * max_seq_length)
|
||||
|
||||
logging.warning(f'{too_long_count} are longer than {max_seq_length}')
|
||||
|
||||
for i in range(min(len(all_input_ids), 1)):
|
||||
logging.info("*** Example ***")
|
||||
logging.info("i: %s", i)
|
||||
logging.info("subtokens: %s", " ".join(list(map(str, all_subtokens[i]))))
|
||||
logging.info("loss_mask: %s", " ".join(list(map(str, all_loss_mask[i]))))
|
||||
logging.info("input_mask: %s", " ".join(list(map(str, all_input_mask[i]))))
|
||||
logging.info("subtokens_mask: %s", " ".join(list(map(str, all_subtokens_mask[i]))))
|
||||
if with_label:
|
||||
logging.info("labels: %s", " ".join(list(map(str, all_labels[i]))))
|
||||
return (all_input_ids, all_segment_ids, all_input_mask, all_subtokens_mask, all_loss_mask, all_labels)
|
||||
|
||||
|
||||
class BertTokenClassificationDataset(Dataset):
|
||||
"""
|
||||
Creates dataset to use during training for token classification tasks with a pretrained model.
|
||||
|
||||
Converts from raw data to an instance that can be used by Dataloader.
|
||||
For dataset to use during inference without labels, see BertTokenClassificationInferDataset.
|
||||
|
||||
Args:
|
||||
text_file: file to sequences, each line should a sentence, no header.
|
||||
label_file: file to labels, each line corresponds to word labels for a sentence in the text_file. No header.
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
tokenizer: such as AutoTokenizer
|
||||
num_samples: number of samples you want to use for the dataset.
|
||||
If -1, use all dataset. Useful for testing.
|
||||
pad_label: pad value use for labels. By default, it's the neutral label.
|
||||
label_ids: label_ids (dict): dict to map labels to label ids.
|
||||
Starts with pad_label->0 and then increases in alphabetical order
|
||||
For dev set use label_ids generated during training to support
|
||||
cases when not all labels are present in the dev set.
|
||||
For training set label_ids should be None.
|
||||
ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
|
||||
ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
|
||||
use_cache: whether to use processed data cache or not
|
||||
"""
|
||||
|
||||
@property
|
||||
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {
|
||||
'input_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'segment_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'input_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'loss_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'labels': NeuralType(('B', 'T'), LabelsType()),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_file: str,
|
||||
label_file: str,
|
||||
max_seq_length: int,
|
||||
tokenizer: TokenizerSpec,
|
||||
num_samples: int = -1,
|
||||
pad_label: str = 'O',
|
||||
label_ids: Dict[str, int] = None,
|
||||
ignore_extra_tokens: bool = False,
|
||||
ignore_start_end: bool = False,
|
||||
use_cache: bool = True,
|
||||
):
|
||||
""" Initializes BertTokenClassificationDataset. """
|
||||
|
||||
data_dir = os.path.dirname(text_file)
|
||||
filename = os.path.basename(text_file)
|
||||
|
||||
if not filename.endswith('.txt'):
|
||||
raise ValueError("{text_file} should have extension .txt")
|
||||
|
||||
vocab_size = getattr(tokenizer, "vocab_size", 0)
|
||||
features_pkl = os.path.join(
|
||||
data_dir,
|
||||
"cached_{}_{}_{}_{}_{}".format(
|
||||
filename, tokenizer.name, str(max_seq_length), str(vocab_size), str(num_samples)
|
||||
),
|
||||
)
|
||||
|
||||
master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
|
||||
features = None
|
||||
if master_device and (not use_cache or not os.path.exists(features_pkl)):
|
||||
if num_samples == 0:
|
||||
raise ValueError("num_samples has to be positive", num_samples)
|
||||
|
||||
with open(text_file, 'r') as f:
|
||||
text_lines = f.readlines()
|
||||
|
||||
labels_lines = []
|
||||
with open(label_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip().split()
|
||||
labels_lines.append(line)
|
||||
|
||||
if len(labels_lines) != len(text_lines):
|
||||
raise ValueError("Labels file should contain labels for every word")
|
||||
|
||||
if num_samples > 0:
|
||||
dataset = list(zip(text_lines, labels_lines))
|
||||
dataset = dataset[:num_samples]
|
||||
|
||||
dataset = list(zip(*dataset))
|
||||
text_lines = dataset[0]
|
||||
labels_lines = dataset[1]
|
||||
|
||||
features = get_features(
|
||||
queries=text_lines,
|
||||
max_seq_length=max_seq_length,
|
||||
tokenizer=tokenizer,
|
||||
pad_label=pad_label,
|
||||
raw_labels=labels_lines,
|
||||
label_ids=label_ids,
|
||||
ignore_extra_tokens=ignore_extra_tokens,
|
||||
ignore_start_end=ignore_start_end,
|
||||
)
|
||||
|
||||
pickle.dump(features, open(features_pkl, "wb"))
|
||||
logging.info(f'features saved to {features_pkl}')
|
||||
|
||||
# wait until the master process writes to the processed data files
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
if features is None:
|
||||
features = pickle.load(open(features_pkl, 'rb'))
|
||||
logging.info(f'features restored from {features_pkl}')
|
||||
|
||||
self.all_input_ids = features[0]
|
||||
self.all_segment_ids = features[1]
|
||||
self.all_input_mask = features[2]
|
||||
self.all_subtokens_mask = features[3]
|
||||
self.all_loss_mask = features[4]
|
||||
self.all_labels = features[5]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_input_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return (
|
||||
np.array(self.all_input_ids[idx]),
|
||||
np.array(self.all_segment_ids[idx]),
|
||||
np.array(self.all_input_mask[idx], dtype=np.long),
|
||||
np.array(self.all_subtokens_mask[idx]),
|
||||
np.array(self.all_loss_mask[idx]),
|
||||
np.array(self.all_labels[idx]),
|
||||
)
|
||||
|
||||
|
||||
class BertTokenClassificationInferDataset(Dataset):
|
||||
"""
|
||||
Creates dataset to use during inference for token classification tasks with a pretrained model.
|
||||
For dataset to use during training with labels, see BertTokenClassificationDataset.
|
||||
"""
|
||||
|
||||
@property
|
||||
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
||||
"""Returns definitions of module output ports.
|
||||
"""
|
||||
return {
|
||||
'input_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'segment_ids': NeuralType(('B', 'T'), ChannelType()),
|
||||
'input_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self, queries: List[str], max_seq_length: int, tokenizer: TokenizerSpec,
|
||||
):
|
||||
"""
|
||||
Initializes BertTokenClassificationInferDataset
|
||||
Args:
|
||||
queries: text sequences
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
tokenizer: such as AutoTokenizer
|
||||
"""
|
||||
features = get_features(queries=queries, max_seq_length=max_seq_length, tokenizer=tokenizer)
|
||||
|
||||
self.all_input_ids = features[0]
|
||||
self.all_segment_ids = features[1]
|
||||
self.all_input_mask = features[2]
|
||||
self.all_subtokens_mask = features[3]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_input_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return (
|
||||
np.array(self.all_input_ids[idx]),
|
||||
np.array(self.all_segment_ids[idx]),
|
||||
np.array(self.all_input_mask[idx], dtype=np.long),
|
||||
np.array(self.all_subtokens_mask[idx]),
|
||||
)
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import Dict
|
||||
|
||||
from nemo.collections.nlp.data.data_utils.data_preprocessing import (
|
||||
fill_class_weights,
|
||||
get_freq_weights,
|
||||
get_label_stats,
|
||||
)
|
||||
from nemo.utils import logging
|
||||
|
||||
__all__ = ['get_label_ids']
|
||||
|
||||
|
||||
def get_label_ids(
|
||||
label_file: str,
|
||||
is_training: bool = False,
|
||||
pad_label: str = 'O',
|
||||
label_ids_dict: Dict[str, int] = None,
|
||||
get_weights: bool = True,
|
||||
class_labels_file_artifact='label_ids.csv',
|
||||
):
|
||||
"""
|
||||
Generates str to int labels mapping for training data or checks correctness of the label_ids_dict
|
||||
file for non-training files or if label_ids_dict is specified
|
||||
|
||||
Args:
|
||||
label_file: the path of the label file to process
|
||||
is_training: indicates whether the label_file is used for training
|
||||
pad_label: token used for padding
|
||||
label_ids_dict: str label name to int ids mapping. Required for non-training data.
|
||||
If specified, the check that all labels from label_file are present in label_ids_dict will be performed.
|
||||
For training data, if label_ids_dict is None, a new mapping will be generated from label_file.
|
||||
get_weights: set to True to calculate class weights, required for Weighted Loss.
|
||||
class_labels_file_artifact: name of the file to save in .nemo
|
||||
"""
|
||||
if not os.path.exists(label_file):
|
||||
raise ValueError(f'File {label_file} was not found.')
|
||||
|
||||
logging.info(f'Processing {label_file}')
|
||||
if not is_training and label_ids_dict is None:
|
||||
raise ValueError(
|
||||
f'For non training data, label_ids_dict created during preprocessing of the training data '
|
||||
f'should be provided'
|
||||
)
|
||||
|
||||
# collect all labels from the label_file
|
||||
data_dir = os.path.dirname(label_file)
|
||||
unique_labels = set(pad_label)
|
||||
all_labels = []
|
||||
with open(label_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip().split()
|
||||
all_labels.extend(line)
|
||||
unique_labels.update(line)
|
||||
|
||||
# check that all labels from label_file are present in the specified label_ids_dict
|
||||
# or generate label_ids_dict from data (for training only)
|
||||
if label_ids_dict:
|
||||
logging.info(f'Using provided labels mapping {label_ids_dict}')
|
||||
for name in unique_labels:
|
||||
if name not in label_ids_dict:
|
||||
raise ValueError(f'{name} class from {label_file} not found in the provided mapping: {label_ids_dict}')
|
||||
else:
|
||||
label_ids_dict = {pad_label: 0}
|
||||
if pad_label in unique_labels:
|
||||
unique_labels.remove(pad_label)
|
||||
for label in sorted(unique_labels):
|
||||
label_ids_dict[label] = len(label_ids_dict)
|
||||
|
||||
label_ids_filename = os.path.join(data_dir, class_labels_file_artifact)
|
||||
if is_training:
|
||||
with open(label_ids_filename, 'w') as f:
|
||||
labels, _ = zip(*sorted(label_ids_dict.items(), key=lambda x: x[1]))
|
||||
f.write('\n'.join(labels))
|
||||
logging.info(f'Labels mapping {label_ids_dict} saved to : {label_ids_filename}')
|
||||
|
||||
# calculate label statistics
|
||||
base_name = os.path.splitext(os.path.basename(label_file))[0]
|
||||
stats_file = os.path.join(data_dir, f'{base_name}_label_stats.tsv')
|
||||
if os.path.exists(stats_file) and not is_training and not get_weights:
|
||||
logging.info(f'{stats_file} found, skipping stats calculation.')
|
||||
else:
|
||||
all_labels = [label_ids_dict[label] for label in all_labels]
|
||||
logging.info(f'Three most popular labels in {label_file}:')
|
||||
total_labels, label_frequencies, max_id = get_label_stats(all_labels, stats_file)
|
||||
logging.info(f'Total labels: {total_labels}. Label frequencies - {label_frequencies}')
|
||||
|
||||
if get_weights:
|
||||
class_weights_pkl = os.path.join(data_dir, f'{base_name}_weights.p')
|
||||
if os.path.exists(class_weights_pkl):
|
||||
class_weights = pickle.load(open(class_weights_pkl, 'rb'))
|
||||
logging.info(f'Class weights restored from {class_weights_pkl}')
|
||||
else:
|
||||
class_weights_dict = get_freq_weights(label_frequencies)
|
||||
logging.info(f'Class Weights: {class_weights_dict}')
|
||||
class_weights = fill_class_weights(class_weights_dict, max_id)
|
||||
|
||||
pickle.dump(class_weights, open(class_weights_pkl, "wb"))
|
||||
logging.info(f'Class weights saved to {class_weights_pkl}')
|
||||
else:
|
||||
class_weights = None
|
||||
|
||||
return label_ids_dict, label_ids_filename, class_weights
|
||||
Reference in New Issue
Block a user