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,119 @@
|
||||
# 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.
|
||||
|
||||
# Config file for text classification with pre-trained BERT models
|
||||
|
||||
trainer:
|
||||
gpus: 1 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1]
|
||||
num_nodes: 1
|
||||
max_epochs: 100
|
||||
max_steps: null # precedence over max_epochs
|
||||
accumulate_grad_batches: 1 # accumulates grads every k batches
|
||||
gradient_clip_val: 0.0
|
||||
amp_level: O0 # O1/O2 for mixed precision
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
accelerator: ddp
|
||||
log_every_n_steps: 1 # Interval of logging.
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
|
||||
checkpoint_callback: False # Provided by exp_manager
|
||||
logger: False # Provided by exp_manager
|
||||
|
||||
model:
|
||||
nemo_path: text_classification_model.nemo # filename to save the model and associated artifacts to .nemo file
|
||||
tokenizer:
|
||||
tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
|
||||
vocab_file: null # path to vocab file
|
||||
tokenizer_model: null # only used if tokenizer is sentencepiece
|
||||
special_tokens: null
|
||||
|
||||
language_model:
|
||||
pretrained_model_name: bert-base-uncased
|
||||
lm_checkpoint: null
|
||||
config_file: null # json file, precedence over config
|
||||
config: null
|
||||
|
||||
classifier_head:
|
||||
num_output_layers: 2
|
||||
fc_dropout: 0.1
|
||||
|
||||
class_labels:
|
||||
class_labels_file : null # optional to specify a file containing the list of the labels
|
||||
|
||||
dataset:
|
||||
num_classes: ??? # The number of classes. 0 < Label <num_classes.
|
||||
do_lower_case: false # true for uncased models, false for cased models, will be set automatically if pre-trained tokenizer model is used
|
||||
max_seq_length: 256 # the maximum length BERT supports is 512
|
||||
class_balancing: null # null or 'weighted_loss'. 'weighted_loss' enables the weighted class balancing of the loss, may be used for handling unbalanced classes
|
||||
use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up
|
||||
|
||||
train_ds:
|
||||
file_path: null
|
||||
batch_size: 64
|
||||
shuffle: true
|
||||
num_samples: -1 # number of samples to be considered, -1 means all the dataset
|
||||
num_workers: 3
|
||||
drop_last: false
|
||||
pin_memory: false
|
||||
|
||||
validation_ds:
|
||||
file_path: null
|
||||
batch_size: 64
|
||||
shuffle: false
|
||||
num_samples: -1 # number of samples to be considered, -1 means all the dataset
|
||||
num_workers: 3
|
||||
drop_last: false
|
||||
pin_memory: false
|
||||
|
||||
test_ds:
|
||||
file_path: null
|
||||
batch_size: 64
|
||||
shuffle: false
|
||||
num_samples: -1 # number of samples to be considered, -1 means all the dataset
|
||||
num_workers: 3
|
||||
drop_last: false
|
||||
pin_memory: false
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 2e-5
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.01
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: WarmupAnnealing
|
||||
# Scheduler params
|
||||
warmup_steps: null
|
||||
warmup_ratio: 0.1
|
||||
last_epoch: -1
|
||||
# pytorch lightning args
|
||||
monitor: val_loss
|
||||
reduce_on_plateau: false
|
||||
|
||||
# List of some sample queries for inference after training is done
|
||||
infer_samples: [
|
||||
'by the end of no such thing the audience , like beatrice , has a watchful affection for the monster .',
|
||||
'director rob marshall went out gunning to make a great one .',
|
||||
'uneasy mishmash of styles and genres .',
|
||||
]
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
|
||||
name: "TextClassification" # The name of your model
|
||||
create_tensorboard_logger: True # Whether you want exp_manger to create a tb logger
|
||||
create_checkpoint_callback: True # Whether you want exp_manager to create a modelcheckpoint callback
|
||||
@@ -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."
|
||||
)
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
This script runs model parallel text classification evaluation.
|
||||
"""
|
||||
import pytorch_lightning as pl
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models.text_classification import TextClassificationModel
|
||||
from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
|
||||
from nemo.core.config import hydra_runner
|
||||
from nemo.utils import logging
|
||||
from nemo.utils.exp_manager import exp_manager
|
||||
|
||||
|
||||
@hydra_runner(config_path="conf", config_name="text_classification_config")
|
||||
def main(cfg: DictConfig) -> None:
|
||||
logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}')
|
||||
trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
|
||||
exp_manager(trainer, cfg.get("exp_manager", None))
|
||||
# TODO: can we drop strict=False
|
||||
model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer, strict=False)
|
||||
model.setup_test_data(test_data_config=cfg.model.test_ds)
|
||||
|
||||
trainer.test(model=model, ckpt_path=None)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,154 @@
|
||||
# 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 contains an example on how to train, evaluate and perform inference with the TextClassificationModel.
|
||||
TextClassificationModel in NeMo supports text classification problems such as sentiment analysis or
|
||||
domain/intent detection for dialogue systems, as long as the data follows the format specified below.
|
||||
|
||||
***Data format***
|
||||
TextClassificationModel requires the data to be stored in TAB separated files (.tsv) with two columns of sentence and
|
||||
label. Each line of the data file contains text sequences, where words are separated with spaces and label separated
|
||||
with [TAB], i.e.:
|
||||
|
||||
[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]
|
||||
|
||||
For example:
|
||||
|
||||
hide new secretions from the parental units[TAB]0
|
||||
that loves its characters and communicates something rather beautiful about human nature[TAB]1
|
||||
...
|
||||
|
||||
If your dataset is stored in another format, you need to convert it to this format to use the TextClassificationModel.
|
||||
|
||||
|
||||
***Setting the configs***
|
||||
The model and the PT trainer are defined in a config file which declares multiple important sections.
|
||||
The most important ones are:
|
||||
model: All arguments that are related to the Model - language model, tokenizer, head classifier, optimizer,
|
||||
schedulers, and datasets/data loaders.
|
||||
trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
|
||||
precision level, etc.
|
||||
|
||||
This script uses the `/examples/nlp/text_classification/conf/text_classification_config.yaml` default config file
|
||||
by default. You may update the config file from the file directly or by using the command line arguments.
|
||||
Other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
|
||||
|
||||
You first need to set the num_classes in the config file which specifies the number of classes in the dataset.
|
||||
Notice that some config lines, including `model.dataset.classes_num`, have `???` as their value, this means that values
|
||||
for these fields are required to be specified by the user. We need to specify and set the `model.train_ds.file_name`,
|
||||
`model.validation_ds.file_name`, and `model.test_ds.file_name` in the config file to the paths of the train, validation,
|
||||
and test files if they exist. We may do it by updating the config file or by setting them from the command line.
|
||||
|
||||
|
||||
***How to run the script?***
|
||||
For example the following would train a model for 50 epochs in 2 GPUs on a classification task with 2 classes:
|
||||
|
||||
# python text_classification_with_bert.py
|
||||
model.dataset.num_classes=2
|
||||
model.train_ds=PATH_TO_TRAIN_FILE
|
||||
model.validation_ds=PATH_TO_VAL_FILE
|
||||
trainer.max_epochs=50
|
||||
trainer.gpus=2
|
||||
|
||||
This script would also reload the last checkpoint after the training is done and does evaluation on the dev set,
|
||||
then performs inference on some sample queries.
|
||||
|
||||
By default, this script uses examples/nlp/text_classification/conf/text_classifciation_config.py config file, and
|
||||
you may update all the params in the config file from the command line. You may also use another config file like this:
|
||||
|
||||
# python text_classification_with_bert.py --config-name==PATH_TO_CONFIG_FILE
|
||||
model.dataset.num_classes=2
|
||||
model.train_ds=PATH_TO_TRAIN_FILE
|
||||
model.validation_ds=PATH_TO_VAL_FILE
|
||||
trainer.max_epochs=50
|
||||
trainer.gpus=2
|
||||
|
||||
***Load a saved model***
|
||||
This script would save the model after training into '.nemo' checkpoint file specified by nemo_path of the model config.
|
||||
You may restore the saved model like this:
|
||||
model = TextClassificationModel.restore_from(restore_path=NEMO_FILE_PATH)
|
||||
|
||||
***Evaluation a saved model on another dataset***
|
||||
# If you wanted to evaluate the saved model on another dataset, you may restore the model and create a new data loader:
|
||||
eval_model = TextClassificationModel.restore_from(restore_path=checkpoint_path)
|
||||
|
||||
# Then, you may create a dataloader config for evaluation:
|
||||
eval_config = OmegaConf.create(
|
||||
{'file_path': cfg.model.test_ds.file_path, 'batch_size': 64, 'shuffle': False, 'num_workers': 3}
|
||||
)
|
||||
eval_model.setup_test_data(test_data_config=eval_config)
|
||||
|
||||
# You need to create a new trainer:
|
||||
eval_trainer = pl.Trainer(gpus=1)
|
||||
eval_model.set_trainer(eval_trainer)
|
||||
eval_trainer.test(model=eval_model, verbose=False)
|
||||
"""
|
||||
import pytorch_lightning as pl
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models.text_classification import TextClassificationModel
|
||||
from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
|
||||
from nemo.core.config import hydra_runner
|
||||
from nemo.utils import logging
|
||||
from nemo.utils.exp_manager import exp_manager
|
||||
|
||||
|
||||
@hydra_runner(config_path="conf", config_name="text_classification_config")
|
||||
def main(cfg: DictConfig) -> None:
|
||||
logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}')
|
||||
trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
|
||||
exp_manager(trainer, cfg.get("exp_manager", None))
|
||||
|
||||
if not cfg.model.train_ds.file_path:
|
||||
raise ValueError("'train_ds.file_path' need to be set for the training!")
|
||||
|
||||
model = TextClassificationModel(cfg.model, trainer=trainer)
|
||||
logging.info("===========================================================================================")
|
||||
logging.info('Starting training...')
|
||||
trainer.fit(model)
|
||||
logging.info('Training finished!')
|
||||
logging.info("===========================================================================================")
|
||||
|
||||
if cfg.model.nemo_path:
|
||||
# '.nemo' file contains the last checkpoint and the params to initialize the model
|
||||
model.save_to(cfg.model.nemo_path)
|
||||
logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}')
|
||||
|
||||
# We evaluate the trained model on the test set if test_ds is set in the config file
|
||||
if cfg.model.test_ds.file_path:
|
||||
logging.info("===========================================================================================")
|
||||
logging.info("Starting the testing of the trained model on test set...")
|
||||
trainer.test(model=model, ckpt_path=None, verbose=False)
|
||||
logging.info("Testing finished!")
|
||||
logging.info("===========================================================================================")
|
||||
|
||||
# perform inference on a list of queries.
|
||||
if "infer_samples" in cfg.model and cfg.model.infer_samples:
|
||||
logging.info("===========================================================================================")
|
||||
logging.info("Starting the inference on some sample queries...")
|
||||
|
||||
# max_seq_length=512 is the maximum length BERT supports.
|
||||
results = model.classifytext(queries=cfg.model.infer_samples, batch_size=16, max_seq_length=512)
|
||||
logging.info('The prediction results of some sample queries with the trained model:')
|
||||
for query, result in zip(cfg.model.infer_samples, results):
|
||||
logging.info(f'Query : {query}')
|
||||
logging.info(f'Predicted label: {result}')
|
||||
|
||||
logging.info("Inference finished!")
|
||||
logging.info("===========================================================================================")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user