mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-27 08:58:13 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,48 @@
|
||||
# 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 pytest
|
||||
import torch
|
||||
from sklearn.metrics import precision_recall_fscore_support
|
||||
|
||||
from nemo.collections.nlp.metrics.classification_report import ClassificationReport
|
||||
|
||||
|
||||
class ClassificationReportTests:
|
||||
num_classes = 3
|
||||
label_ids = {'a': 0, 'b': 1, 'c': 2}
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_classification_report(self):
|
||||
|
||||
preds = torch.Tensor([0, 1, 1, 1, 2, 2, 0])
|
||||
labels = torch.Tensor([1, 0, 0, 1, 2, 1, 0])
|
||||
|
||||
def __convert_to_tensor(sklearn_metric):
|
||||
return torch.Tensor([round(sklearn_metric * 100)])[0]
|
||||
|
||||
for mode in ['macro', 'micro', 'weighted']:
|
||||
|
||||
classification_report_nemo = ClassificationReport(
|
||||
num_classes=self.num_classes, label_ids=self.label_ids, mode=mode
|
||||
)
|
||||
# pytest.set_trace()
|
||||
precision, recall, f1, _ = classification_report_nemo(preds, labels)
|
||||
tp, fp, fn = classification_report_nemo.tp, classification_report_nemo.fp, classification_report_nemo.fn
|
||||
pr_sklearn, recall_sklearn, f1_sklearn, _ = precision_recall_fscore_support(labels, preds, average=mode)
|
||||
|
||||
self.assertEqual(torch.round(precision), __convert_to_tensor(pr_sklearn), f'wrong precision for {mode}')
|
||||
self.assertEqual(torch.round(recall), __convert_to_tensor(recall_sklearn), f'wrong recall for {mode}')
|
||||
self.assertEqual(torch.round(f1), __convert_to_tensor(f1_sklearn), f'wrong f1 for {mode}')
|
||||
@@ -0,0 +1,83 @@
|
||||
# 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 os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import wget
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models import EntityLinkingModel
|
||||
|
||||
|
||||
def get_cfg():
|
||||
|
||||
language_model = OmegaConf.create(
|
||||
{"pretrained_model_name": "bert-base-uncased", "config_file": None, "config": None, "lm_checkpoint": None}
|
||||
)
|
||||
|
||||
tokenizer = OmegaConf.create(
|
||||
{"tokenizer_name": "bert-base-uncased", "vocab_file": None, "tokenizer_model": None, "do_lower_case": True}
|
||||
)
|
||||
|
||||
model = OmegaConf.create(
|
||||
{
|
||||
"nemo_path": "sap_entity_linking.nemo",
|
||||
"max_seq_length": 128,
|
||||
"language_model": language_model,
|
||||
"tokenizer": tokenizer,
|
||||
"train_ds": None,
|
||||
"validation_ds": None,
|
||||
}
|
||||
)
|
||||
|
||||
cfg = OmegaConf.create({"model": model})
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
class TestEntityLinkingModel:
|
||||
@pytest.mark.unit
|
||||
def test_creation_saving_restoring(self):
|
||||
# Create a new temporary directory
|
||||
with tempfile.TemporaryDirectory() as restore_dir:
|
||||
with tempfile.TemporaryDirectory() as save_dir:
|
||||
model = EntityLinkingModel(cfg=get_cfg().model)
|
||||
assert isinstance(model, EntityLinkingModel)
|
||||
|
||||
save_dir_path = save_dir
|
||||
|
||||
# Where model will be saved
|
||||
model_save_path = os.path.join(save_dir, f"{model.__class__.__name__}.nemo")
|
||||
model.save_to(save_path=model_save_path)
|
||||
|
||||
# Where model will be restored from
|
||||
model_restore_path = os.path.join(restore_dir, f"{model.__class__.__name__}.nemo")
|
||||
shutil.copy(model_save_path, model_restore_path)
|
||||
|
||||
# at this point save_dir should not exist
|
||||
assert save_dir_path is not None and not os.path.exists(save_dir_path)
|
||||
assert not os.path.exists(model_save_path)
|
||||
assert os.path.exists(model_restore_path)
|
||||
|
||||
# attempt to restore
|
||||
model_copy = model.__class__.restore_from(restore_path=model_restore_path)
|
||||
assert model.num_weights == model_copy.num_weights
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestEntityLinkingModel()
|
||||
t.test_creation_saving_restoring()
|
||||
@@ -0,0 +1,86 @@
|
||||
# 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 tempfile
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import nemo.collections.nlp as nemo_nlp
|
||||
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
|
||||
from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
|
||||
|
||||
|
||||
def do_export(model, name: str):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Generate filename in the temporary directory.
|
||||
tmp_file_name = os.path.join(tmpdir, name + '.onnx')
|
||||
# Test export.
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
model.export(tmp_file_name)
|
||||
|
||||
|
||||
class TestHuggingFace:
|
||||
@pytest.mark.unit
|
||||
def test_list_pretrained_models(self):
|
||||
pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list()
|
||||
assert len(pretrained_lm_models) > 0
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_bert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='bert-base-uncased')
|
||||
assert isinstance(model, nemo_nlp.modules.BertEncoder)
|
||||
do_export(model, "bert-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_distilbert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='distilbert-base-uncased')
|
||||
assert isinstance(model, nemo_nlp.modules.DistilBertEncoder)
|
||||
do_export(model, "distilbert-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_roberta_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='roberta-base')
|
||||
assert isinstance(model, nemo_nlp.modules.RobertaEncoder)
|
||||
do_export(model, "roberta-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_albert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='albert-base-v1')
|
||||
assert isinstance(model, nemo_nlp.modules.AlbertEncoder)
|
||||
do_export(model, "albert-base-v1")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_chinese_bert_wwm_model(self):
|
||||
model_name = 'hfl/chinese-bert-wwm'
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
assert isinstance(model, nemo_nlp.modules.BertModule)
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
||||
assert isinstance(tokenizer, AutoTokenizer)
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_arabic_model(self):
|
||||
model_name = 'asafaya/bert-base-arabic'
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
assert isinstance(model, nemo_nlp.modules.BertModule)
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
||||
assert isinstance(tokenizer, AutoTokenizer)
|
||||
@@ -0,0 +1,76 @@
|
||||
# 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.
|
||||
|
||||
try:
|
||||
import apex
|
||||
|
||||
apex_available = True
|
||||
except Exception:
|
||||
apex_available = False
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import onnx
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import nemo.collections.nlp as nemo_nlp
|
||||
from nemo.core.classes import typecheck
|
||||
|
||||
|
||||
def get_pretrained_bert_345m_uncased_model():
|
||||
model_name = "megatron-bert-345m-uncased"
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
return model
|
||||
|
||||
|
||||
class TestMegatron:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_list_pretrained_models(self):
|
||||
pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list()
|
||||
assert len(pretrained_lm_models) > 0
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.skip("Only one Megatron model is allowed")
|
||||
def test_get_model(self):
|
||||
model = get_pretrained_bert_345m_uncased_model()
|
||||
assert isinstance(model, nemo_nlp.modules.MegatronBertEncoder)
|
||||
|
||||
typecheck.set_typecheck_enabled(enabled=False)
|
||||
inp = model.input_example()
|
||||
out = model.forward(*inp)
|
||||
typecheck.set_typecheck_enabled(enabled=True)
|
||||
|
||||
@pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine')
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_onnx_export(self):
|
||||
model = get_pretrained_bert_345m_uncased_model()
|
||||
assert model
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Generate filename in the temporary directory.
|
||||
# Test export.
|
||||
model.export(os.path.join(".", "megatron.onnx"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestMegatron()
|
||||
t.test_onnx_export()
|
||||
@@ -0,0 +1,176 @@
|
||||
# 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 tempfile
|
||||
|
||||
import onnx
|
||||
import pytest
|
||||
import pytorch_lightning as pl
|
||||
import torch
|
||||
import wget
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections import nlp as nemo_nlp
|
||||
from nemo.collections.nlp.models import IntentSlotClassificationModel
|
||||
from nemo.collections.nlp.modules.common import (
|
||||
SequenceClassifier,
|
||||
SequenceRegression,
|
||||
SequenceTokenClassifier,
|
||||
TokenClassifier,
|
||||
)
|
||||
|
||||
|
||||
def classifier_export(obj):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, obj.__class__.__name__ + '.onnx')
|
||||
obj = obj.cuda()
|
||||
obj.export(output=filename)
|
||||
|
||||
|
||||
class TestExportableClassifiers:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_token_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_bert_pretraining_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_token_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(
|
||||
SequenceTokenClassifier(hidden_size=256, num_slots=8, num_intents=8, num_layers=num_layers)
|
||||
)
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(SequenceClassifier(hidden_size=256, num_classes=16, num_layers=num_layers))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_regression_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(SequenceRegression(hidden_size=256, num_layers=num_layers))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_IntentSlotClassificationModel_export_to_onnx(self, dummy_data):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
wget.download(
|
||||
'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/'
|
||||
'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml',
|
||||
tmpdir,
|
||||
)
|
||||
config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml')
|
||||
config = OmegaConf.load(config_file)
|
||||
config = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
|
||||
config.model.data_dir = dummy_data
|
||||
config.trainer.gpus = 1
|
||||
config.trainer.precision = 32
|
||||
config.trainer.accelerator = None
|
||||
trainer = pl.Trainer(**config.trainer)
|
||||
model = IntentSlotClassificationModel(config.model, trainer=trainer)
|
||||
filename = os.path.join(tmpdir, 'isc.onnx')
|
||||
model.export(output=filename)
|
||||
onnx_model = onnx.load(filename)
|
||||
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
|
||||
assert onnx_model.graph.input[0].name == 'input_ids'
|
||||
assert onnx_model.graph.input[1].name == 'attention_mask'
|
||||
assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'intent_logits'
|
||||
assert onnx_model.graph.output[1].name == 'slot_logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.skip('Are we exporting to torchscript?')
|
||||
def test_IntentSlotClassificationModel_export_to_torchscript(self, dummy_data):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
wget.download(
|
||||
'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/'
|
||||
'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml',
|
||||
tmpdir,
|
||||
)
|
||||
config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml')
|
||||
config = OmegaConf.load(config_file)
|
||||
config = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
|
||||
config.model.data_dir = dummy_data
|
||||
config.trainer.gpus = 1
|
||||
config.trainer.precision = 32
|
||||
config.trainer.accelerator = None
|
||||
trainer = pl.Trainer(**config.trainer)
|
||||
model = IntentSlotClassificationModel(config.model, trainer=trainer)
|
||||
filename = os.path.join(tmpdir, 'isc.pt')
|
||||
model.export(output=filename, check_trace=True)
|
||||
torchscript_model = torch.jit.load(filename)
|
||||
inputs = list(torchscript_model.graph.inputs())
|
||||
assert inputs[1].debugName() == 'input_ids.1'
|
||||
assert inputs[2].debugName() == 'input.1'
|
||||
assert inputs[3].debugName() == 'attention_mask.1'
|
||||
modules = list(torchscript_model.modules())
|
||||
assert modules[2].original_name == 'BertEmbeddings'
|
||||
assert modules[46].original_name == 'BertSelfAttention'
|
||||
assert modules[217].original_name == 'SequenceTokenClassifier'
|
||||
assert modules[219].original_name == 'MultiLayerPerceptron'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_TokenClassificationModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.TokenClassificationModel.from_pretrained(model_name="ner_en_bert")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'ner.onnx')
|
||||
model.export(output=filename)
|
||||
onnx_model = onnx.load(filename)
|
||||
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
|
||||
assert onnx_model.graph.input[0].name == 'input_ids'
|
||||
assert onnx_model.graph.output[0].name == 'logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_PunctuationCapitalizationModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name="punctuation_en_distilbert")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'puncap.onnx')
|
||||
model.export(output=filename)
|
||||
onnx_model = onnx.load(filename)
|
||||
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
|
||||
assert onnx_model.graph.input[0].name == 'input_ids'
|
||||
# assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'punct_logits'
|
||||
assert onnx_model.graph.output[1].name == 'capit_logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_QAModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.QAModel.from_pretrained(model_name="qa_squadv2.0_bertbase")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'qa.onnx')
|
||||
model.export(output=filename)
|
||||
onnx_model = onnx.load(filename)
|
||||
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
|
||||
assert onnx_model.graph.input[0].name == 'input_ids'
|
||||
assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'logits'
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def dummy_data(test_data_dir):
|
||||
return os.path.join(test_data_dir, 'nlp', 'dummy_data')
|
||||
@@ -0,0 +1,144 @@
|
||||
# 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 shutil
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models import MTEncDecModel
|
||||
from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig
|
||||
|
||||
|
||||
def export_test(model, suffix, try_script=False):
|
||||
with tempfile.TemporaryDirectory() as restore_folder:
|
||||
enc_filename = os.path.join(restore_folder, 'nmt_e' + suffix)
|
||||
dec_filename = os.path.join(restore_folder, 'nmt_d' + suffix)
|
||||
model.encoder.export(output=enc_filename, try_script=try_script)
|
||||
model.decoder.export(output=dec_filename, try_script=try_script)
|
||||
assert os.path.exists(enc_filename)
|
||||
assert os.path.exists(dec_filename)
|
||||
|
||||
|
||||
def get_cfg():
|
||||
cfg = AAYNBaseConfig()
|
||||
cfg.encoder_tokenizer.tokenizer_name = 'yttm'
|
||||
cfg.encoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model'
|
||||
cfg.decoder_tokenizer.tokenizer_name = 'yttm'
|
||||
cfg.decoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model'
|
||||
cfg.train_ds = None
|
||||
cfg.validation_ds = None
|
||||
cfg.test_ds = None
|
||||
return cfg
|
||||
|
||||
|
||||
class TestMTEncDecModel:
|
||||
@pytest.mark.unit
|
||||
def test_creation_saving_restoring(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
# Create a new temporary directory
|
||||
with tempfile.TemporaryDirectory() as restore_folder:
|
||||
with tempfile.TemporaryDirectory() as save_folder:
|
||||
save_folder_path = save_folder
|
||||
# Where model will be saved
|
||||
model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
|
||||
model.save_to(save_path=model_save_path)
|
||||
# Where model will be restored from
|
||||
model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
|
||||
shutil.copy(model_save_path, model_restore_path)
|
||||
# at this point save_folder should not exist
|
||||
assert save_folder_path is not None and not os.path.exists(save_folder_path)
|
||||
assert not os.path.exists(model_save_path)
|
||||
assert os.path.exists(model_restore_path)
|
||||
# attempt to restore
|
||||
model_copy = model.__class__.restore_from(restore_path=model_restore_path)
|
||||
assert model.num_weights == model_copy.num_weights
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_no_artifact_name_collision(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
with tempfile.TemporaryDirectory() as tmpdir1:
|
||||
model.save_to("nmt_model.nemo")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
model._unpack_nemo_file(path2file="nmt_model.nemo", out_folder=tmpdir)
|
||||
conf = OmegaConf.load(os.path.join(tmpdir, "model_config.yaml"))
|
||||
# Make sure names now differ in saved config
|
||||
assert conf.encoder_tokenizer.tokenizer_model != conf.decoder_tokenizer.tokenizer_model
|
||||
# Make sure names in config start with "nemo:" prefix
|
||||
assert conf.encoder_tokenizer.tokenizer_model.startswith("nemo:")
|
||||
assert conf.decoder_tokenizer.tokenizer_model.startswith("nemo:")
|
||||
# Check if both tokenizers were included
|
||||
assert os.path.exists(os.path.join(tmpdir, conf.encoder_tokenizer.tokenizer_model[5:]))
|
||||
assert os.path.exists(os.path.join(tmpdir, conf.decoder_tokenizer.tokenizer_model[5:]))
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_train_eval_loss(self):
|
||||
cfg = get_cfg()
|
||||
cfg.label_smoothing = 0.5
|
||||
model = MTEncDecModel(cfg=cfg)
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
batch_size = 10
|
||||
time = 32
|
||||
vocab_size = 32000
|
||||
torch.manual_seed(42)
|
||||
tgt_ids = torch.LongTensor(batch_size, time).random_(1, model.decoder_tokenizer.vocab_size)
|
||||
logits = torch.FloatTensor(batch_size, time, vocab_size).random_(-1, 1)
|
||||
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
|
||||
train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
assert not torch.allclose(train_loss, eval_loss) # , (train_loss, eval_loss)
|
||||
|
||||
cfg.label_smoothing = 0
|
||||
model = MTEncDecModel(cfg=cfg)
|
||||
# Train loss == val loss when label smoothing = 0
|
||||
train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
assert torch.allclose(train_loss, eval_loss)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_cpu_export_onnx(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".ts")
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_cpu_export_ts(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".onnx")
|
||||
|
||||
@pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine')
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_gpu_export_ts(self):
|
||||
model = MTEncDecModel(cfg=get_cfg()).cuda()
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".ts")
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_gpu_export_onnx(self):
|
||||
model = MTEncDecModel(cfg=get_cfg()).cuda()
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".onnx")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestMTEncDecModel()
|
||||
# t.test_gpu_export_ts()
|
||||
t.test_train_eval_loss()
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from unittest import TestCase
|
||||
|
||||
import pytest
|
||||
import pytorch_lightning as pl
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import nemo.collections.nlp.models as models
|
||||
|
||||
|
||||
def get_metrics(data_dir, model):
|
||||
trainer = pl.Trainer(gpus=[0])
|
||||
|
||||
model.set_trainer(trainer)
|
||||
model.update_data_dir(data_dir)
|
||||
|
||||
test_ds = OmegaConf.create(
|
||||
{
|
||||
'text_file': 'text_dev.txt',
|
||||
'labels_file': 'labels_dev.txt',
|
||||
'shuffle': False,
|
||||
'num_samples': -1,
|
||||
'batch_size': 8,
|
||||
}
|
||||
)
|
||||
|
||||
model._cfg.dataset.use_cache = False
|
||||
trainer.test(model)
|
||||
model.setup_test_data(test_data_config=test_ds)
|
||||
metrics = trainer.test(model)[0]
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def data_exists(data_dir):
|
||||
return os.path.exists(data_dir)
|
||||
|
||||
|
||||
class TestPretrainedModelPerformance:
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_punct_capit_with_bert(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher'
|
||||
model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_bert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['punct_precision'] - 52.3024) < 0.001
|
||||
assert abs(metrics['punct_recall'] - 58.9220) < 0.001
|
||||
assert abs(metrics['punct_f1'] - 53.2976) < 0.001
|
||||
assert abs(metrics['capit_precision'] - 87.0707) < 0.001
|
||||
assert abs(metrics['capit_recall'] - 87.0707) < 0.001
|
||||
assert abs(metrics['capit_f1'] - 87.0707) < 0.001
|
||||
assert int(model.punct_class_report.total_examples) == 128
|
||||
|
||||
preds_512 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=512)[0]
|
||||
assert preds_512 == 'What can I do for you today?'
|
||||
preds_5 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=5)[0]
|
||||
assert preds_5 == 'What can i'
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_punct_capit_with_distilbert(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher'
|
||||
model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_distilbert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['punct_precision'] - 53.0826) < 0.001
|
||||
assert abs(metrics['punct_recall'] - 56.2905) < 0.001
|
||||
assert abs(metrics['punct_f1'] - 52.4225) < 0.001
|
||||
assert int(model.punct_class_report.total_examples) == 128
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/gmb'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_ner_model(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/gmb'
|
||||
model = models.TokenClassificationModel.from_pretrained("ner_en_bert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['precision'] - 96.0937) < 0.001
|
||||
assert abs(metrics['recall'] - 96.0146) < 0.001
|
||||
assert abs(metrics['f1'] - 95.6076) < 0.001
|
||||
assert int(model.classification_report.total_examples) == 202
|
||||
Reference in New Issue
Block a user