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,19 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from nemo.collections.nlp.data.text_classification.text_classification_dataset import (
|
||||
TextClassificationDataset,
|
||||
calc_class_weights,
|
||||
)
|
||||
@@ -0,0 +1,297 @@
|
||||
# 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.
|
||||
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
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 (
|
||||
fill_class_weights,
|
||||
get_freq_weights,
|
||||
get_label_stats,
|
||||
get_stats,
|
||||
)
|
||||
from nemo.collections.nlp.parts.utils_funcs import list2str
|
||||
from nemo.core.classes import Dataset
|
||||
from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
|
||||
from nemo.utils import logging
|
||||
from nemo.utils.env_var_parsing import get_envint
|
||||
|
||||
__all__ = ['TextClassificationDataset', 'calc_class_weights']
|
||||
|
||||
|
||||
class TextClassificationDataset(Dataset):
|
||||
"""A dataset class that converts from raw data to
|
||||
a dataset that can be used by DataLayerNM.
|
||||
Args:
|
||||
input_file: file to sequence + label.
|
||||
the first line is header (sentence [tab] label)
|
||||
each line should be [sentence][tab][label]
|
||||
tokenizer: tokenizer object such as AutoTokenizer
|
||||
max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
|
||||
num_samples: number of samples you want to use for the dataset.
|
||||
If -1, use all dataset. Useful for testing.
|
||||
shuffle: Shuffles the dataset after loading.
|
||||
use_cache: Enables caching to use pickle format to store and read data from
|
||||
"""
|
||||
|
||||
@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()),
|
||||
'label': NeuralType(('B',), LabelsType()),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: TokenizerSpec,
|
||||
input_file: str = None,
|
||||
queries: List[str] = None,
|
||||
max_seq_length: int = -1,
|
||||
num_samples: int = -1,
|
||||
shuffle: bool = False,
|
||||
use_cache: bool = False,
|
||||
):
|
||||
if not input_file and not queries:
|
||||
raise ValueError("Either input_file or queries should be passed to the text classification dataset.")
|
||||
|
||||
if input_file and not os.path.exists(input_file):
|
||||
raise FileNotFoundError(
|
||||
f'Data file `{input_file}` not found! Each line of the data file should contain text sequences, where '
|
||||
f'words are separated with spaces and the label separated by [TAB] following this format: '
|
||||
f'[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]'
|
||||
)
|
||||
|
||||
self.input_file = input_file
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.num_samples = num_samples
|
||||
self.shuffle = shuffle
|
||||
self.use_cache = use_cache
|
||||
self.vocab_size = self.tokenizer.vocab_size
|
||||
self.pad_id = tokenizer.pad_id
|
||||
|
||||
self.features = None
|
||||
labels, all_sents = [], []
|
||||
if input_file:
|
||||
data_dir, filename = os.path.split(input_file)
|
||||
vocab_size = getattr(tokenizer, "vocab_size", 0)
|
||||
tokenizer_name = tokenizer.name
|
||||
cached_features_file = os.path.join(
|
||||
data_dir,
|
||||
f"cached_{filename}_{tokenizer_name}_{max_seq_length}_{vocab_size}_{num_samples}_{self.pad_id}_{shuffle}.pkl",
|
||||
)
|
||||
|
||||
if get_envint("LOCAL_RANK", 0) == 0:
|
||||
if use_cache and os.path.exists(cached_features_file):
|
||||
logging.warning(
|
||||
f"Processing of {input_file} is skipped as caching is enabled and a cache file "
|
||||
f"{cached_features_file} already exists."
|
||||
)
|
||||
logging.warning(
|
||||
f"You may need to delete the cache file if any of the processing parameters (eg. tokenizer) or "
|
||||
f"the data are updated."
|
||||
)
|
||||
else:
|
||||
with open(input_file, "r") as f:
|
||||
lines = f.readlines()
|
||||
logging.info(f'Read {len(lines)} examples from {input_file}.')
|
||||
if num_samples > 0:
|
||||
lines = lines[:num_samples]
|
||||
logging.warning(
|
||||
f"Parameter 'num_samples' is set, so just the first {len(lines)} examples are kept."
|
||||
)
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(lines)
|
||||
|
||||
for index, line in enumerate(lines):
|
||||
if index % 20000 == 0:
|
||||
logging.debug(f"Processing line {index}/{len(lines)}")
|
||||
line_splited = line.strip().split()
|
||||
try:
|
||||
label = int(line_splited[-1])
|
||||
except ValueError:
|
||||
logging.debug(f"Skipping line {line}")
|
||||
continue
|
||||
labels.append(label)
|
||||
sent_words = line_splited[:-1]
|
||||
all_sents.append(sent_words)
|
||||
verbose = True
|
||||
|
||||
self.features = self.get_features(
|
||||
all_sents=all_sents,
|
||||
tokenizer=tokenizer,
|
||||
max_seq_length=max_seq_length,
|
||||
labels=labels,
|
||||
verbose=verbose,
|
||||
)
|
||||
with open(cached_features_file, 'wb') as out_file:
|
||||
pickle.dump(self.features, out_file, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
else:
|
||||
for query in queries:
|
||||
all_sents.append(query.strip().split())
|
||||
labels = [-1] * len(all_sents)
|
||||
verbose = False
|
||||
self.features = self.get_features(
|
||||
all_sents=all_sents, tokenizer=tokenizer, max_seq_length=max_seq_length, labels=labels, verbose=verbose
|
||||
)
|
||||
|
||||
# wait until the master process writes to the processed data files
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
if input_file:
|
||||
with open(cached_features_file, "rb") as input_file:
|
||||
self.features = pickle.load(input_file)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.features)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.features[idx]
|
||||
|
||||
def _collate_fn(self, batch):
|
||||
"""collate batch of input_ids, segment_ids, input_mask, and label
|
||||
Args:
|
||||
batch: A list of tuples of (input_ids, segment_ids, input_mask, label).
|
||||
"""
|
||||
max_length = 0
|
||||
for input_ids, segment_ids, input_mask, label in batch:
|
||||
if len(input_ids) > max_length:
|
||||
max_length = len(input_ids)
|
||||
|
||||
padded_input_ids = []
|
||||
padded_segment_ids = []
|
||||
padded_input_mask = []
|
||||
labels = []
|
||||
for input_ids, segment_ids, input_mask, label in batch:
|
||||
if len(input_ids) < max_length:
|
||||
pad_width = max_length - len(input_ids)
|
||||
padded_input_ids.append(np.pad(input_ids, pad_width=[0, pad_width], constant_values=self.pad_id))
|
||||
padded_segment_ids.append(np.pad(segment_ids, pad_width=[0, pad_width], constant_values=self.pad_id))
|
||||
padded_input_mask.append(np.pad(input_mask, pad_width=[0, pad_width], constant_values=self.pad_id))
|
||||
else:
|
||||
padded_input_ids.append(input_ids)
|
||||
padded_segment_ids.append(segment_ids)
|
||||
padded_input_mask.append(input_mask)
|
||||
labels.append(label)
|
||||
|
||||
return (
|
||||
torch.LongTensor(padded_input_ids),
|
||||
torch.LongTensor(padded_segment_ids),
|
||||
torch.LongTensor(padded_input_mask),
|
||||
torch.LongTensor(labels),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_features(all_sents, tokenizer, max_seq_length, labels=None, verbose=True):
|
||||
"""Encode a list of sentences into a list of tuples of (input_ids, segment_ids, input_mask, label)."""
|
||||
features = []
|
||||
sent_lengths = []
|
||||
too_long_count = 0
|
||||
for sent_id, sent in enumerate(all_sents):
|
||||
if sent_id % 1000 == 0:
|
||||
logging.debug(f"Encoding sentence {sent_id}/{len(all_sents)}")
|
||||
sent_subtokens = [tokenizer.cls_token]
|
||||
for word in sent:
|
||||
word_tokens = tokenizer.text_to_tokens(word)
|
||||
sent_subtokens.extend(word_tokens)
|
||||
|
||||
if max_seq_length > 0 and len(sent_subtokens) + 1 > max_seq_length:
|
||||
sent_subtokens = sent_subtokens[:max_seq_length]
|
||||
too_long_count += 1
|
||||
|
||||
sent_subtokens.append(tokenizer.sep_token)
|
||||
sent_lengths.append(len(sent_subtokens))
|
||||
|
||||
input_ids = [tokenizer.tokens_to_ids(t) for t in sent_subtokens]
|
||||
|
||||
# The mask has 1 for real tokens and 0 for padding tokens.
|
||||
# Only real tokens are attended to.
|
||||
input_mask = [1] * len(input_ids)
|
||||
segment_ids = [0] * len(input_ids)
|
||||
|
||||
if verbose and sent_id < 2:
|
||||
logging.info("*** Example ***")
|
||||
logging.info(f"example {sent_id}: {sent}")
|
||||
logging.info("subtokens: %s" % " ".join(sent_subtokens))
|
||||
logging.info("input_ids: %s" % list2str(input_ids))
|
||||
logging.info("segment_ids: %s" % list2str(segment_ids))
|
||||
logging.info("input_mask: %s" % list2str(input_mask))
|
||||
logging.info("label: %s" % labels[sent_id] if labels else "**Not Provided**")
|
||||
|
||||
label = labels[sent_id] if labels else -1
|
||||
features.append([np.asarray(input_ids), np.asarray(segment_ids), np.asarray(input_mask), label])
|
||||
|
||||
if max_seq_length > -1 and too_long_count > 0:
|
||||
logging.warning(
|
||||
f'Found {too_long_count} out of {len(all_sents)} sentences with more than {max_seq_length} subtokens. '
|
||||
f'Truncated long sentences from the end.'
|
||||
)
|
||||
if verbose:
|
||||
get_stats(sent_lengths)
|
||||
return features
|
||||
|
||||
|
||||
def calc_class_weights(file_path: str, num_classes: int):
|
||||
"""
|
||||
iterates over a data file and calculate the weights of each class to be used for class_balancing
|
||||
Args:
|
||||
file_path: path to the data file
|
||||
num_classes: number of classes in the dataset
|
||||
"""
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"Could not find data file {file_path} to calculate the class weights!")
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
input_lines = f.readlines()
|
||||
|
||||
labels = []
|
||||
for input_line in input_lines:
|
||||
parts = input_line.strip().split()
|
||||
try:
|
||||
label = int(parts[-1])
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f'No numerical labels found for {file_path}. Labels should be integers and separated by [TAB] at the end of each line.'
|
||||
)
|
||||
labels.append(label)
|
||||
|
||||
logging.info(f'Calculating stats of {file_path}...')
|
||||
total_sents, sent_label_freq, max_id = get_label_stats(labels, f'{file_path}_sentence_stats.tsv', verbose=False)
|
||||
if max_id >= num_classes:
|
||||
raise ValueError(f'Found an invalid label in {file_path}! Labels should be from [0, num_classes-1].')
|
||||
|
||||
class_weights_dict = get_freq_weights(sent_label_freq)
|
||||
|
||||
logging.info(f'Total Sentence: {total_sents}')
|
||||
logging.info(f'Sentence class frequencies: {sent_label_freq}')
|
||||
|
||||
logging.info(f'Class Weights: {class_weights_dict}')
|
||||
class_weights = fill_class_weights(weights=class_weights_dict, max_id=num_classes - 1)
|
||||
|
||||
return class_weights
|
||||
Reference in New Issue
Block a user