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,243 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
This script can be used to process and import IMDB, ChemProt, SST-2, and THUCnews datasets into NeMo's format.
|
||||
You may run it as the following:
|
||||
|
||||
python import_datasets.py \
|
||||
--dataset_name DATASET_NAME \
|
||||
--target_data_dir TARGET_PATH \
|
||||
--source_data_dir SOURCE_PATH
|
||||
|
||||
The dataset should be specified by "DATASET_NAME" which can be from ["sst-2", "chemprot", "imdb", "thucnews"].
|
||||
It reads the data from "SOURCE_PATH" folder, processes and converts the data into NeMo's format.
|
||||
Then writes the results into "TARGET_PATH" folder.
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import glob
|
||||
import os
|
||||
from os.path import exists
|
||||
|
||||
import tqdm
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def process_imdb(infold, outfold, uncased, modes=['train', 'test']):
|
||||
if not os.path.exists(infold):
|
||||
link = 'https://ai.stanford.edu/~amaas/data/sentiment/'
|
||||
raise ValueError(
|
||||
f'Data not found at {infold}. '
|
||||
f'Please download IMDB reviews dataset from {link} and '
|
||||
f'extract it into the folder specified by source_data_dir argument.'
|
||||
)
|
||||
|
||||
logging.info(f'Processing IMDB dataset and store at {outfold}')
|
||||
os.makedirs(outfold, exist_ok=True)
|
||||
|
||||
outfiles = {}
|
||||
for mode in modes:
|
||||
outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
|
||||
for sent in ['neg', 'pos']:
|
||||
if sent == 'neg':
|
||||
label = 0
|
||||
else:
|
||||
label = 1
|
||||
files = glob.glob(f'{infold}/{mode}/{sent}/*.txt')
|
||||
for file in files:
|
||||
with open(file, 'r') as f:
|
||||
review = f.read().strip()
|
||||
if uncased:
|
||||
review = review.lower()
|
||||
review = review.replace("<br />", "")
|
||||
outfiles[mode].write(f'{review}\t{label}\n')
|
||||
|
||||
for mode in modes:
|
||||
outfiles[mode].close()
|
||||
|
||||
class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
|
||||
class_labels_file.write('negative\npositive\n')
|
||||
class_labels_file.close()
|
||||
|
||||
|
||||
def process_sst2(infold, outfold, uncased, splits=['train', 'dev']):
|
||||
"""Process sst2 dataset."""
|
||||
# "test" split doesn't have labels, so it is skipped
|
||||
if not os.path.exists(infold):
|
||||
link = 'https://dl.fbaipublicfiles.com/glue/data/SST-2.zip'
|
||||
raise ValueError(
|
||||
f'Data not found at {infold}. Please download SST-2 dataset from `{link}` and '
|
||||
f'extract it into the folder specified by `source_data_dir` argument.'
|
||||
)
|
||||
|
||||
logging.info(f'Processing SST-2 dataset')
|
||||
os.makedirs(outfold, exist_ok=True)
|
||||
|
||||
def _read_tsv(input_file, quotechar=None):
|
||||
"""Read a tab separated value file."""
|
||||
with open(input_file, "r") as f:
|
||||
reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
|
||||
lines = []
|
||||
for line in reader:
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
||||
for split in splits:
|
||||
# Load input file.
|
||||
input_file = os.path.join(infold, split + '.tsv')
|
||||
lines = _read_tsv(input_file)
|
||||
# Create output.
|
||||
outfile = open(os.path.join(outfold, split + '.tsv'), 'w')
|
||||
|
||||
# Copy lines, skip the header (line 0).
|
||||
for line in lines[1:]:
|
||||
text = line[0]
|
||||
label = line[1]
|
||||
# Lowercase when required.
|
||||
if uncased:
|
||||
text = text.lower()
|
||||
# Write output.
|
||||
outfile.write(f'{text}\t{label}\n')
|
||||
# Close file.
|
||||
outfile.close()
|
||||
|
||||
class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
|
||||
class_labels_file.write('negative\npositive\n')
|
||||
class_labels_file.close()
|
||||
|
||||
logging.info(f'Result stored at {outfold}')
|
||||
|
||||
|
||||
def process_chemprot(source_dir, target_dir, uncased, modes=['train', 'test', 'dev']):
|
||||
if not os.path.exists(source_dir):
|
||||
link = 'https://github.com/arwhirang/recursive_chemprot/tree/master/Demo/tree_LSTM/data'
|
||||
raise ValueError(f'Data not found at {source_dir}. ' f'Please download ChemProt from {link}.')
|
||||
|
||||
logging.info(f'Processing Chemprot dataset and store at {target_dir}')
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
|
||||
naming_map = {'train': 'trainingPosit_chem', 'test': 'testPosit_chem', 'dev': 'developPosit_chem'}
|
||||
|
||||
def _read_tsv(input_file, quotechar=None):
|
||||
"""Reads a tab separated value file."""
|
||||
with open(input_file, "r") as f:
|
||||
reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
|
||||
lines = []
|
||||
for line in reader:
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
||||
outfiles = {}
|
||||
label_mapping = {}
|
||||
out_label_mapping = open(os.path.join(target_dir, 'label_mapping.tsv'), 'w')
|
||||
for mode in modes:
|
||||
outfiles[mode] = open(os.path.join(target_dir, mode + '.tsv'), 'w')
|
||||
input_file = os.path.join(source_dir, naming_map[mode])
|
||||
lines = _read_tsv(input_file)
|
||||
for line in lines:
|
||||
text = line[1]
|
||||
label = line[2]
|
||||
if label == "True":
|
||||
label = line[3]
|
||||
if uncased:
|
||||
text = text.lower()
|
||||
if label not in label_mapping:
|
||||
out_label_mapping.write(f'{label}\t{len(label_mapping)}\n')
|
||||
label_mapping[label] = len(label_mapping)
|
||||
label = label_mapping[label]
|
||||
outfiles[mode].write(f'{text}\t{label}\n')
|
||||
for mode in modes:
|
||||
outfiles[mode].close()
|
||||
out_label_mapping.close()
|
||||
|
||||
|
||||
def process_thucnews(infold, outfold):
|
||||
modes = ['train', 'test']
|
||||
train_size = 0.8
|
||||
if not os.path.exists(infold):
|
||||
link = 'thuctc.thunlp.org/'
|
||||
raise ValueError(f'Data not found at {infold}. ' f'Please download THUCNews from {link}.')
|
||||
|
||||
logging.info(f'Processing THUCNews dataset and store at {outfold}')
|
||||
os.makedirs(outfold, exist_ok=True)
|
||||
|
||||
outfiles = {}
|
||||
for mode in modes:
|
||||
outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'a+', encoding='utf-8')
|
||||
categories = ['体育', '娱乐', '家居', '彩票', '房产', '教育', '时尚', '时政', '星座', '游戏', '社会', '科技', '股票', '财经']
|
||||
for category in categories:
|
||||
label = categories.index(category)
|
||||
category_files = glob.glob(f'{infold}/{category}/*.txt')
|
||||
test_num = int(len(category_files) * (1 - train_size))
|
||||
test_files = category_files[:test_num]
|
||||
train_files = category_files[test_num:]
|
||||
|
||||
for mode in modes:
|
||||
logging.info(f'Processing {mode} data of the category {category}')
|
||||
if mode == 'test':
|
||||
files = test_files
|
||||
else:
|
||||
files = train_files
|
||||
|
||||
if len(files) == 0:
|
||||
logging.info(f'Skipping category {category} for {mode} mode')
|
||||
continue
|
||||
|
||||
for file in tqdm.tqdm(files):
|
||||
with open(file, 'r', encoding='utf-8') as f:
|
||||
news = f.read().strip().replace('\r', '')
|
||||
news = news.replace('\n', '').replace('\t', ' ')
|
||||
outfiles[mode].write(f'{news}\t{label}\n')
|
||||
for mode in modes:
|
||||
outfiles[mode].close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Parse the command-line arguments.
|
||||
parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.")
|
||||
parser.add_argument("--dataset_name", required=True, type=str, choices=['imdb', 'thucnews', 'chemprot'])
|
||||
parser.add_argument(
|
||||
"--source_data_dir", required=True, type=str, help='The path to the folder containing the dataset files.'
|
||||
)
|
||||
parser.add_argument("--target_data_dir", required=True, type=str)
|
||||
parser.add_argument("--do_lower_case", action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
dataset_name = args.dataset_name
|
||||
do_lower_case = args.do_lower_case
|
||||
source_dir = args.source_data_dir
|
||||
target_dir = args.target_data_dir
|
||||
|
||||
if not exists(source_dir):
|
||||
raise FileNotFoundError(f"{source_dir} does not exist.")
|
||||
|
||||
if dataset_name == 'imdb':
|
||||
process_imdb(source_dir, target_dir, do_lower_case)
|
||||
elif dataset_name == 'thucnews':
|
||||
process_thucnews(source_dir, target_dir)
|
||||
elif dataset_name == "chemprot":
|
||||
process_chemprot(source_dir, target_dir, do_lower_case)
|
||||
elif dataset_name == "sst-2":
|
||||
process_sst2(source_dir, target_dir, do_lower_case)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Dataset {dataset_name} is not supported.'
|
||||
+ "Please make sure that you build the preprocessing process for it. "
|
||||
+ "NeMo's format assumes that a data file has a header and each line of the file follows "
|
||||
+ "the format: text [TAB] label. Label is assumed to be an integer."
|
||||
)
|
||||
Reference in New Issue
Block a user