mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,242 @@
|
||||
# Copyright (c) 2021, 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 argparse
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import string
|
||||
import subprocess
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
URL = {'tatoeba': 'https://downloads.tatoeba.org/exports/sentences.csv'}
|
||||
|
||||
|
||||
def __maybe_download_file(destination: str, source: str):
|
||||
"""
|
||||
Downloads source to destination if not exists.
|
||||
If exists, skips download
|
||||
Args:
|
||||
destination: local filepath
|
||||
source: url of resource
|
||||
"""
|
||||
source = URL[source]
|
||||
if not os.path.exists(destination):
|
||||
logging.info(f'Downloading {source} to {destination}')
|
||||
subprocess.run(['wget', '-O', destination, source])
|
||||
else:
|
||||
logging.info(f'{destination} found. Skipping download')
|
||||
|
||||
|
||||
def __process_english_sentences(
|
||||
in_file: str, out_file: str, percent_to_cut: float = 0, num_to_combine: int = 1, num_samples: int = -1
|
||||
):
|
||||
"""
|
||||
Extract English sentences from the Tatoeba dataset.
|
||||
|
||||
Expected in_file format
|
||||
that
|
||||
contrain letters and punctuation marks (,.?).
|
||||
Chop and combine sentences.
|
||||
Args:
|
||||
in_file: local filepath to the tatoeba dataset.
|
||||
Format: id [TAB] region_name [TAB] sentence,
|
||||
for example: "1276\teng\tLet's try something.\n"
|
||||
out_file: local filepath to the clean dataset
|
||||
percent_to_cut: Percent of sentences to cut in the middle
|
||||
to get examples of incomplete sentences.
|
||||
This could be useful since ASR output not always
|
||||
represents a complete sentence
|
||||
num_to_combine: Number of sentences to combine into
|
||||
a single example
|
||||
num_samples: Number of samples in the final dataset
|
||||
"""
|
||||
if not os.path.exists(in_file):
|
||||
raise FileNotFoundError(f'{in_file} not found.')
|
||||
|
||||
in_file = open(in_file, 'r')
|
||||
out_file = open(out_file, 'w')
|
||||
lines_to_combine = []
|
||||
samples_count = 0
|
||||
|
||||
for line in in_file:
|
||||
line = line.split('\t')
|
||||
# use only English sentences
|
||||
if line[1] == 'eng':
|
||||
line = line[2].strip()
|
||||
if re.match("^[A-Z][A-Za-z.,'?\s]+$", line): # nopep8
|
||||
# chop some sentences in the middle
|
||||
if percent_to_cut > 0:
|
||||
line = line.split()
|
||||
if random.random() < percent_to_cut:
|
||||
line = line[: len(line) // 2]
|
||||
line = ' '.join(line)
|
||||
|
||||
# combine multiple sentences into a single example
|
||||
# to make it harder for the model to learn eos punctuation
|
||||
if len(lines_to_combine) >= num_to_combine:
|
||||
if samples_count == num_samples:
|
||||
return
|
||||
out_file.write(' '.join(lines_to_combine) + '\n')
|
||||
lines_to_combine = []
|
||||
samples_count += 1
|
||||
lines_to_combine.append(line)
|
||||
|
||||
if len(lines_to_combine) > 0 and (samples_count < num_samples or num_samples < 0):
|
||||
out_file.write(' '.join(lines_to_combine) + '\n')
|
||||
|
||||
|
||||
def __split_into_train_dev(in_file: str, train_file: str, dev_file: str, percent_dev: float):
|
||||
"""
|
||||
Create train and dev split of the dataset.
|
||||
Args:
|
||||
in_file: local filepath to the dataset
|
||||
train_file: local filepath to the train dataset
|
||||
dev_file: local filepath to the dev dataset
|
||||
percent_dev: Percent of the sentences in the dev set
|
||||
"""
|
||||
if not os.path.exists(in_file):
|
||||
raise FileNotFoundError(f'{in_file} not found.')
|
||||
|
||||
lines = open(in_file, 'r').readlines()
|
||||
train_file = open(train_file, 'w')
|
||||
dev_file = open(dev_file, 'w')
|
||||
|
||||
dev_size = int(len(lines) * percent_dev)
|
||||
train_file.write(' '.join(lines[:-dev_size]))
|
||||
dev_file.write(' '.join(lines[-dev_size:]))
|
||||
|
||||
|
||||
def remove_punctuation(word: str):
|
||||
"""
|
||||
Removes all punctuation marks from a word except for '
|
||||
that is often a part of word: don't, it's, and so on
|
||||
"""
|
||||
all_punct_marks = string.punctuation.replace("'", '')
|
||||
return re.sub('[' + all_punct_marks + ']', '', word)
|
||||
|
||||
|
||||
def create_text_and_labels(output_dir: str, file_path: str, punct_marks: str = ',.?'):
|
||||
"""
|
||||
Create datasets for training and evaluation.
|
||||
|
||||
Args:
|
||||
output_dir: path to the output data directory
|
||||
file_path: path to file name
|
||||
punct_marks: supported punctuation marks
|
||||
|
||||
The data will be split 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).'
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
raise ValueError(f'{file_path} not found')
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
base_name = os.path.basename(file_path)
|
||||
labels_file = os.path.join(output_dir, 'labels_' + base_name)
|
||||
text_file = os.path.join(output_dir, 'text_' + base_name)
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
with open(text_file, 'w') as text_f:
|
||||
with open(labels_file, 'w') as labels_f:
|
||||
for line in f:
|
||||
line = line.split()
|
||||
text = ''
|
||||
labels = ''
|
||||
for word in line:
|
||||
label = word[-1] if word[-1] in punct_marks else 'O'
|
||||
word = remove_punctuation(word)
|
||||
if len(word) > 0:
|
||||
if word[0].isupper():
|
||||
label += 'U'
|
||||
else:
|
||||
label += 'O'
|
||||
|
||||
word = word.lower()
|
||||
text += word + ' '
|
||||
labels += label + ' '
|
||||
|
||||
text_f.write(text.strip() + '\n')
|
||||
labels_f.write(labels.strip() + '\n')
|
||||
|
||||
print(f'{text_file} and {labels_file} created from {file_path}.')
|
||||
|
||||
|
||||
def __delete_file(file_to_del: str):
|
||||
"""
|
||||
Deletes the file
|
||||
Args:
|
||||
file_to_del: local filepath to the file to delete
|
||||
"""
|
||||
if os.path.exists(file_to_del):
|
||||
os.remove(file_to_del)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Prepare tatoeba dataset')
|
||||
parser.add_argument("--data_dir", required=True, type=str)
|
||||
parser.add_argument("--dataset", default='tatoeba', type=str)
|
||||
parser.add_argument("--num_samples", default=-1, type=int, help='-1 to use the whole dataset')
|
||||
parser.add_argument("--percent_to_cut", default=0, type=float, help='Percent of sentences to cut in the middle')
|
||||
parser.add_argument(
|
||||
"--num_lines_to_combine", default=1, type=int, help='Number of lines to combine into single example'
|
||||
)
|
||||
parser.add_argument("--percent_dev", default=0.2, type=float, help='Size of the dev set, float')
|
||||
parser.add_argument("--clean_dir", action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.data_dir):
|
||||
os.makedirs(args.data_dir)
|
||||
|
||||
if args.dataset != 'tatoeba':
|
||||
raise ValueError("Unsupported dataset.")
|
||||
|
||||
logging.info(f'Downloading tatoeba dataset')
|
||||
tatoeba_dataset = os.path.join(args.data_dir, 'sentences.csv')
|
||||
__maybe_download_file(tatoeba_dataset, args.dataset)
|
||||
|
||||
logging.info(f'Processing English sentences...')
|
||||
clean_eng_sentences = os.path.join(args.data_dir, 'clean_eng_sentences.txt')
|
||||
__process_english_sentences(
|
||||
tatoeba_dataset, clean_eng_sentences, args.percent_to_cut, args.num_lines_to_combine, args.num_samples
|
||||
)
|
||||
|
||||
train_file = os.path.join(args.data_dir, 'train.txt')
|
||||
dev_file = os.path.join(args.data_dir, 'dev.txt')
|
||||
|
||||
logging.info(
|
||||
f'Splitting the {args.dataset} dataset into train and dev sets' + ' and creating labels and text files'
|
||||
)
|
||||
__split_into_train_dev(clean_eng_sentences, train_file, dev_file, args.percent_dev)
|
||||
|
||||
logging.info(f'Creating text and label files for training')
|
||||
create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'train.txt'))
|
||||
create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'dev.txt'))
|
||||
|
||||
if args.clean_dir:
|
||||
logging.info(f'Cleaning up {args.data_dir}')
|
||||
__delete_file(clean_eng_sentences)
|
||||
__delete_file(tatoeba_dataset)
|
||||
__delete_file(train_file)
|
||||
__delete_file(dev_file)
|
||||
logging.info(f'Processing of the {args.dataset} is complete')
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) 2021, 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 argparse
|
||||
import os
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def __convert_data(in_file: str, out_text_f: str, out_labels_f: str, max_length: int):
|
||||
"""
|
||||
Convert data from the IOB format to NeMo accepted format described below.
|
||||
in_file should be in the IOB format, see example here:
|
||||
https://www.clips.uantwerpen.be/conll2003/ner/.
|
||||
|
||||
Args:
|
||||
in_file: input file name
|
||||
out_text_f: output file with text
|
||||
out_labels_f: output file with labels
|
||||
max_length: use -1 to leave the examples' length as is, otherwise long examples will be split into multiple
|
||||
examples
|
||||
After the conversion, the dataset is split 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).
|
||||
|
||||
"""
|
||||
in_file = open(in_file, 'r')
|
||||
|
||||
if max_length == -1:
|
||||
with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
|
||||
for line in in_file:
|
||||
if line == '\n':
|
||||
out_text.write(line)
|
||||
out_labels.write(line)
|
||||
else:
|
||||
line = line.split()
|
||||
out_text.write(line[0] + ' ')
|
||||
out_labels.write(line[-1] + ' ')
|
||||
|
||||
else:
|
||||
words = []
|
||||
labels = []
|
||||
with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
|
||||
lines = in_file.readlines()
|
||||
for line_id, line in enumerate(lines):
|
||||
logging.info(f"{line_id} {len(lines)}")
|
||||
contends = line.strip()
|
||||
if len(contends) == 0:
|
||||
assert len(words) == len(labels)
|
||||
if len(words) > max_length:
|
||||
# split if the sentence is longer than max_length
|
||||
while len(words) > max_length:
|
||||
tmplabel = labels[:max_length]
|
||||
for iidx in range(len(tmplabel)):
|
||||
if tmplabel.pop() == 'O':
|
||||
break
|
||||
l = ' '.join([label for label in labels[: len(tmplabel) + 1] if len(label) > 0])
|
||||
w = ' '.join([word for word in words[: len(tmplabel) + 1] if len(word) > 0])
|
||||
|
||||
out_text.write(w + "\n")
|
||||
out_labels.write(l + "\n")
|
||||
words = words[len(tmplabel) + 1 :]
|
||||
labels = labels[len(tmplabel) + 1 :]
|
||||
|
||||
if len(words) == 0:
|
||||
continue
|
||||
l = ' '.join([label for label in labels if len(label) > 0])
|
||||
w = ' '.join([word for word in words if len(word) > 0])
|
||||
|
||||
out_text.write(w + "\n")
|
||||
out_labels.write(l + "\n")
|
||||
words = []
|
||||
labels = []
|
||||
continue
|
||||
|
||||
word = line.strip().split()[0]
|
||||
label = line.strip().split()[-1]
|
||||
words.append(word)
|
||||
labels.append(label)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Convert data from IOB format to the format compatible with \
|
||||
nlp/examples/token_classification/scripts/token_classification_train.py and \
|
||||
token_classification_evaluate.py'
|
||||
)
|
||||
parser.add_argument("--data_file", required=True, type=str, help='path to a file in IOB format')
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
default=-1,
|
||||
type=int,
|
||||
help='use -1 to leave the examples\'s length as is, '
|
||||
'otherwise long examples will be split into multiple examples',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
data_dir, basename = os.path.split(args.data_file)
|
||||
prefix = os.path.splitext(basename)[0]
|
||||
if not os.path.exists(args.data_file):
|
||||
raise FileNotFoundError(f"{args.data_file} not found")
|
||||
|
||||
logging.info(f'Processing {args.data_file}')
|
||||
out_text = os.path.join(data_dir, 'text_' + prefix + '.txt')
|
||||
out_labels = os.path.join(data_dir, 'labels_' + prefix + '.txt')
|
||||
|
||||
__convert_data(args.data_file, out_text, out_labels, args.max_length)
|
||||
logging.info(f'Processing of the {args.data_file} is complete')
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) 2021, 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.
|
||||
|
||||
"""
|
||||
The script converts raw text to the NeMo format for punctuation and capitalization task.
|
||||
|
||||
Raw Data Format
|
||||
---------------
|
||||
|
||||
The Punctuation and Capitalization model can work with any text dataset, although it is recommended to balance the data, especially for the punctuation task.
|
||||
Before pre-processing the data to the format expected by the model, the data should be split into train.txt and dev.txt (and optionally test.txt).
|
||||
Each line in the **train.txt/dev.txt/test.txt** should represent one or more full and/or truncated sentences.
|
||||
|
||||
Example of the train.txt/dev.txt file:
|
||||
When is the next flight to New York?
|
||||
The next flight is ...
|
||||
....
|
||||
|
||||
|
||||
The `source_data_dir` structure should look like this:
|
||||
.
|
||||
|--sourced_data_dir
|
||||
|-- dev.txt
|
||||
|-- train.txt
|
||||
|
||||
|
||||
|
||||
NeMo Data Format for training the model
|
||||
---------------------------------------
|
||||
|
||||
The punctuation and capitalization model expects the data in the following format:
|
||||
|
||||
The training and evaluation data is divided into 2 files: text.txt and labels.txt. \
|
||||
Each line of the **text.txt** file contains text sequences, where words are separated with spaces, i.e.
|
||||
|
||||
[WORD] [SPACE] [WORD] [SPACE] [WORD], for example:
|
||||
when is the next flight to new york
|
||||
the next flight is ...
|
||||
...
|
||||
|
||||
The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. \
|
||||
Each label in labels.txt file consists of 2 symbols:
|
||||
|
||||
* the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed);
|
||||
* the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.)
|
||||
|
||||
By default the following punctuation marks are considered: commas, periods, and question marks; the rest punctuation marks were removed from the data.
|
||||
This can be changed by introducing new labels in the labels.txt files
|
||||
|
||||
Each line of the labels.txt should follow the format: [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). \
|
||||
For example, labels for the above text.txt file should be:
|
||||
|
||||
OU OO OO OO OO OO OU ?U
|
||||
OU OO OO OO ...
|
||||
...
|
||||
|
||||
The complete list of all possible labels for this task used in this tutorial is: OO, ,O, .O, ?O, OU, ,U, .U, ?U.
|
||||
|
||||
Converting Raw data to NeMo format
|
||||
----------------------------------
|
||||
|
||||
To pre-process the raw text data, stored under :code:`sourced_data_dir` (see the :ref:`raw_data_format_punct`
|
||||
section), run the following command:
|
||||
|
||||
python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \
|
||||
-s <PATH_TO_THE_SOURCE_FILE>
|
||||
-o <PATH_TO_THE_OUTPUT_DIRECTORY>
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from get_tatoeba_data import create_text_and_labels
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Prepare data for punctuation and capitalization tasks')
|
||||
parser.add_argument("-s", "--source_file", required=True, type=str, help="Path to the source file")
|
||||
parser.add_argument("-o", "--output_dir", required=True, type=str, help="Path to the output directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.source_file):
|
||||
raise ValueError(f'{args.source_file} was not found')
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
create_text_and_labels(args.output_dir, args.source_file)
|
||||
|
||||
print(f'Processing of the {args.source_file} is complete')
|
||||
Reference in New Issue
Block a user