NLP refactoring - Stage 2 (#368)

* refactor dataset utils Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * black fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * refactor datasets Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fixes Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * black fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fixes Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * wip black fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * trade example fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * utils refacoting Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * text_cl fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * rm output Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * glue fix, jenkins updated Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * revert jenkins Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Removed CrossEntropy3D. Merged PaddedSmoothedCrossEntropyLossNM with MaskedLanguageModelingLossNM into unified SmoothedCrossEntropyLossNM. Moved SmoothedCrossEntropyLoss into the file for SmoothedCrossEntropyLossNM. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed style. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * wikitext down fixed, jenkins update Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Loss names updated. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bugs. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bugs. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * jenkins Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed targets_id bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * jenkins update Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * jenkins update Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * jenkins update Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * DataProcessor class moved to datasets_utils/ Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Style fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bugs. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bugs. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added boolmasktype. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed style. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed style. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Remove vocab file in jenkins. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added boolmasktype. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Styel fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added optional port to smoothed_cross_entropy_loss Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added MaskType Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Replaced TokenClassifcaitionLoss with crossentropyloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Replaced TokenClassifcaitionLoss with crossentropyloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Replaced TokenClassifcaitionLoss with crossentropyloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Style fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Removed tokenclassificationloss Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added back logits_dim to CrossEntropyLoss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Chnaged crossentropyloss to crossentropylossnm Signed-off-by: VahidooX <vnoroozi@nvidia.com> * refactored and added doc string to data_layers folder Signed-off-by: Yang Zhang <yangzhang@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * added doc string to trainables folder file Signed-off-by: Yang Zhang <yangzhang@nvidia.com> * adding header to squad dataset Signed-off-by: Yang Zhang <yangzhang@nvidia.com> * refactored nm folder utils Signed-off-by: Yang Zhang <yangzhang@nvidia.com> * wkt refactored Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * black Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * docstrings added Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * fix bug Signed-off-by: Yang Zhang <yangzhang@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Bug fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added weights to lossaggregator. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added lossaggregator to jointintentsloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added lossaggregator to jointintentsloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Added lossaggregator to jointintentsloss. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Style fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Style fixed. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * shuffle removed, lgtm fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed lgtm warnings. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * lgtm fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Moved lossaggregator to backedn.common. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * changelog updated Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * bug fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * jenkins fixed to used the same bert in all tests Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed comments. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * jenkins bert default fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed comments. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * deleted unused losses Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * import fix Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed comments. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * bug fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Resolved reviews. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * import fixed Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Updated comments. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * licences updated Signed-off-by: Evelina Bakhturina <ebakhturina@nvidia.com> * Fixed lgtm warnings. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed lgtm warnings. Signed-off-by: VahidooX <vnoroozi@nvidia.com> * Fixed bug. Signed-off-by: VahidooX <vnoroozi@nvidia.com> Co-authored-by: Vahid Noroozi <VahidooX@users.noreply.github.com> Co-authored-by: Yang Zhang <yzhang123@users.noreply.github.com>
NVIDIA · Feb 21, 2020 · 46045f9 · 46045f9
1 parent 2c99584
commit 46045f9
Show file tree

Hide file tree

Showing 94 changed files with 3,750 additions and 3,409 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -85,6 +85,13 @@ files, along with unit tests, examples and tutorials
 ([PR #375](https://github.com/NVIDIA/NeMo/pull/375)) - @titu1994
 
 ### Changed
+- Refactoring of `nemo_nlp` collections: 
+([PR #368](https://github.com/NVIDIA/NeMo/pull/368)) - @VahidooX, @yzhang123, @ekmb
+    - renaming and restructuring of files, folder, and functions in `nemo_nlp`
+    - losses cleaned up. LossAggregatorNM moved to nemo/backends/pytorch/common/losses
+ ([PR #316](https://github.com/NVIDIA/NeMo/pull/316)) - @VahidooX, @yzhang123, @ekmb
+    - renaming and restructuring of files, folder, and functions in `nemo_nlp`
+    - Updated licenses
 - All collections changed to use New Neural Type System.
 ([PR #307](https://github.com/NVIDIA/NeMo/pull/307)) - @okuchaiev
 - Additional Collections Repositories merged into core `nemo_toolkit` package.
@@ -95,10 +102,6 @@ files, along with unit tests, examples and tutorials
 ([PR #286](https://github.com/NVIDIA/NeMo/pull/286)) - @stasbel
 - Major cleanup of Neural Module constructors (init), aiming at increasing the framework robustness: cleanup of NeuralModule initialization logic, refactor of trainer/actions (getting rid of local_params), fixes of several examples and unit tests, extraction and storing of intial parameters (init_params).  
 ([PR #309](https://github.com/NVIDIA/NeMo/pull/309)) - @tkornuta-nvidia
-- Refactoring of `nemo_nlp` collections: 
-([PR #316](https://github.com/NVIDIA/NeMo/pull/316)) - @VahidooX, @yzhang123, @ekmb
-    - renaming of files and restructuring of folder in `nemo_nlp`
-    - Updated licenses
 - Updated nemo's use of the logging library. from nemo import logging is now the reccomended way of using the nemo logger. neural_factory.logger and all other instances of logger are now deprecated and planned for removal in the next version. Please see PR 267 for complete change information.
 ([PR #267](https://github.com/NVIDIA/NeMo/pull/267), [PR #283](https://github.com/NVIDIA/NeMo/pull/283), [PR #305](https://github.com/NVIDIA/NeMo/pull/305), [PR #311](https://github.com/NVIDIA/NeMo/pull/311)) - @blisc
 - Changed Distributed Data Parallel from Apex to Torch

diff --git a/Jenkinsfile b/Jenkinsfile
@@ -91,19 +91,21 @@ pipeline {
       parallel {
         stage ('Text Classification with BERT Test') {
           steps {
-            sh 'cd examples/nlp/text_classification && CUDA_VISIBLE_DEVICES=0 python text_classification_with_bert.py --num_epochs=1 --max_seq_length=50 --dataset_name=jarvis --data_dir=/home/TestData/nlp/retail/ --eval_file_prefix=eval --batch_size=10 --num_train_samples=-1 --do_lower_case --shuffle_data --work_dir=outputs'
+            sh 'cd examples/nlp/text_classification && CUDA_VISIBLE_DEVICES=0 python text_classification_with_bert.py --pretrained_bert_model bert-base-uncased --num_epochs=1 --max_seq_length=50 --dataset_name=jarvis --data_dir=/home/TestData/nlp/retail/ --eval_file_prefix=eval --batch_size=10 --num_train_samples=-1 --do_lower_case --shuffle_data --work_dir=outputs'
             sh 'rm -rf examples/nlp/text_classification/outputs'
           }
         }
         stage ('Dialogue State Tracking - TRADE - Multi-GPUs') {
           steps {
+            sh 'rm -rf /home/TestData/nlp/multiwoz2.1/vocab.pkl'
             sh 'cd examples/nlp/dialogue_state_tracking && CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 dialogue_state_tracking_trade.py --batch_size=10 --eval_batch_size=10 --num_train_samples=-1 --num_eval_samples=-1 --num_epochs=1 --dropout=0.2 --eval_file_prefix=test --shuffle_data --num_gpus=2 --lr=0.001 --grad_norm_clip=10 --work_dir=outputs --data_dir=/home/TestData/nlp/multiwoz2.1'
             sh 'rm -rf examples/nlp/dialogue_state_tracking/outputs'
+            sh 'rm -rf /home/TestData/nlp/multiwoz2.1/vocab.pkl'
           }
         }
         stage ('GLUE Benchmark Test') {
           steps {
-            sh 'cd examples/nlp/glue_benchmark && CUDA_VISIBLE_DEVICES=1 python glue_benchmark_with_bert.py --data_dir /home/TestData/nlp/glue_fake/MRPC --work_dir glue_output --save_step_freq -1 --num_epochs 1 --task_name mrpc --batch_size 2'
+            sh 'cd examples/nlp/glue_benchmark && CUDA_VISIBLE_DEVICES=1 python glue_benchmark_with_bert.py --data_dir /home/TestData/nlp/glue_fake/MRPC --pretrained_bert_model bert-base-uncased --work_dir glue_output --save_step_freq -1 --num_epochs 1 --task_name mrpc --batch_size 2'
             sh 'rm -rf examples/nlp/glue_benchmark/glue_output'
           }
         }
@@ -122,8 +124,8 @@ pipeline {
       parallel {
         stage('Token Classification Training/Inference Test') {
           steps {
-            sh 'cd examples/nlp/token_classification && CUDA_VISIBLE_DEVICES=0 python token_classification.py --data_dir /home/TestData/nlp/token_classification_punctuation/ --batch_size 2 --num_epochs 1 --save_epoch_freq 1 --work_dir token_classification_output --pretrained_bert_model bert-base-cased'
-            sh 'cd examples/nlp/token_classification && DATE_F=$(ls token_classification_output/) && CUDA_VISIBLE_DEVICES=0 python token_classification_infer.py --work_dir token_classification_output/$DATE_F/checkpoints/ --labels_dict /home/TestData/nlp/token_classification_punctuation/label_ids.csv --pretrained_bert_model bert-base-cased'
+            sh 'cd examples/nlp/token_classification && CUDA_VISIBLE_DEVICES=0 python token_classification.py --data_dir /home/TestData/nlp/token_classification_punctuation/ --batch_size 2 --num_epochs 1 --save_epoch_freq 1 --work_dir token_classification_output --pretrained_bert_model bert-base-uncased'
+            sh 'cd examples/nlp/token_classification && DATE_F=$(ls token_classification_output/) && CUDA_VISIBLE_DEVICES=0 python token_classification_infer.py --work_dir token_classification_output/$DATE_F/checkpoints/ --labels_dict /home/TestData/nlp/token_classification_punctuation/label_ids.csv --pretrained_bert_model bert-base-uncased'
             sh 'rm -rf examples/nlp/token_classification/token_classification_output'
           }
         }

diff --git a/docs/sources/source/nlp/bert_pretraining.rst b/docs/sources/source/nlp/bert_pretraining.rst
@@ -191,7 +191,7 @@ For training from raw text use nemo_nlp.BertPretrainingDataLayer, for preprocess
 
             mlm_logits = mlm_classifier(hidden_states=hidden_states)
             mlm_loss = mlm_loss_fn(logits=mlm_logits,
-                                   output_ids=input_data.output_ids,
+                                   labels=input_data.output_ids,
                                    output_mask=input_data.output_mask)
 
             nsp_logits = nsp_classifier(hidden_states=hidden_states)

diff --git a/docs/sources/source/nlp/joint_intent_slot_filling.rst b/docs/sources/source/nlp/joint_intent_slot_filling.rst
@@ -3,9 +3,9 @@ Tutorial
 
 In this tutorial, we are going to implement a joint intent and slot filling system with pretrained BERT model based on
 `BERT for Joint Intent Classification and Slot Filling <https://arxiv.org/abs/1902.10909>`_ :cite:`nlp-slot-chen2019bert`.
-All code used in this tutorial is based on ``examples/nlp/joint_intent_slot_with_bert.py``.
+All code used in this tutorial is based on ``examples/nlp/intent_detection_slot_tagging/joint_intent_slot_with_bert.py``.
 
-There are four pre-trained BERT models that we can select from using the argument `--pretrained_bert_model`. We're currently
+There are a variety pre-trained BERT models that we can select from using the argument `--pretrained_bert_model`. We're currently
 using the script for loading pre-trained models from `pytorch_transformers`. See the list of available pre-trained models
 `here <https://huggingface.co/pytorch-transformers/pretrained_models.html>`__. 
 

diff --git a/examples/nlp/asr_postprocessor/asr_postprocessor.py b/examples/nlp/asr_postprocessor/asr_postprocessor.py
@@ -113,7 +113,7 @@
     args.d_model, num_classes=vocab_size, num_layers=1, log_softmax=True
 )
 
-loss_fn = nemo_nlp.nm.losses.PaddedSmoothedCrossEntropyLossNM(pad_id=tokenizer.pad_id, label_smoothing=0.1)
+loss_fn = nemo_nlp.nm.losses.SmoothedCrossEntropyLoss(pad_id=tokenizer.pad_id, label_smoothing=0.1)
 
 beam_search = nemo_nlp.nm.trainables.BeamSearchTranslatorNM(
     decoder=decoder,
@@ -174,7 +174,7 @@ def create_pipeline(dataset, tokens_in_batch, clean=False, training=True):
         input_ids_tgt=tgt, hidden_states_src=src_hiddens, input_mask_src=src_mask, input_mask_tgt=tgt_mask
     )
     log_softmax = t_log_softmax(hidden_states=tgt_hiddens)
-    loss = loss_fn(logits=log_softmax, target_ids=labels)
+    loss = loss_fn(logits=log_softmax, labels=labels)
     beam_results = None
     if not training:
         beam_results = beam_search(hidden_states_src=src_hiddens, input_mask_src=src_mask)

diff --git a/examples/nlp/dialogue_state_tracking/dialogue_state_tracking_trade.py b/examples/nlp/dialogue_state_tracking/dialogue_state_tracking_trade.py
@@ -25,12 +25,14 @@
 
 import numpy as np
 
+import nemo.backends.pytorch as nemo_backend
+import nemo.backends.pytorch.common.losses
 import nemo.collections.nlp as nemo_nlp
 import nemo.core as nemo_core
 from nemo import logging
 from nemo.backends.pytorch.common import EncoderRNN
 from nemo.collections.nlp.callbacks.state_tracking_trade_callback import eval_epochs_done_callback, eval_iter_callback
-from nemo.collections.nlp.data.datasets.state_tracking_trade_dataset import MultiWOZDataDesc
+from nemo.collections.nlp.data.datasets.multiwoz_dataset import MultiWOZDataDesc
 from nemo.utils.lr_policies import get_lr_policy
 
 parser = argparse.ArgumentParser(description='Dialog state tracking with TRADE model on MultiWOZ dataset')
@@ -97,9 +99,9 @@
     teacher_forcing=args.teacher_forcing,
 )
 
-gate_loss_fn = nemo_nlp.nm.losses.CrossEntropyLoss3D(num_classes=len(data_desc.gating_dict))
-ptr_loss_fn = nemo_nlp.nm.losses.TRADEMaskedCrossEntropy()
-total_loss_fn = nemo_nlp.nm.losses.LossAggregatorNM(num_inputs=2)
+gate_loss_fn = nemo_backend.losses.CrossEntropyLossNM(logits_dim=3)
+ptr_loss_fn = nemo_nlp.nm.losses.MaskedXEntropyLoss()
+total_loss_fn = nemo.backends.pytorch.common.losses.LossAggregatorNM(num_inputs=2)
 
 
 def create_pipeline(num_samples, batch_size, num_gpus, input_dropout, data_prefix, is_training):
@@ -142,7 +144,7 @@ def create_pipeline(num_samples, batch_size, num_gpus, input_dropout, data_prefi
     )
 
     gate_loss = gate_loss_fn(logits=gate_outputs, labels=gate_labels)
-    ptr_loss = ptr_loss_fn(logits=point_outputs, targets=tgt_ids, loss_mask=tgt_lens)
+    ptr_loss = ptr_loss_fn(logits=point_outputs, labels=tgt_ids, length_mask=tgt_lens)
     total_loss = total_loss_fn(loss_1=gate_loss, loss_2=ptr_loss)
 
     if is_training:

diff --git a/examples/nlp/glue_benchmark/glue_benchmark_with_bert.py b/examples/nlp/glue_benchmark/glue_benchmark_with_bert.py
@@ -68,7 +68,7 @@
 import nemo.collections.nlp as nemo_nlp
 import nemo.core as nemo_core
 from nemo import logging
-from nemo.backends.pytorch.common import CrossEntropyLoss, MSELoss
+from nemo.backends.pytorch.common import CrossEntropyLossNM, MSELoss
 from nemo.collections.nlp.callbacks.glue_benchmark_callback import eval_epochs_done_callback, eval_iter_callback
 from nemo.collections.nlp.data import NemoBertTokenizer, SentencePieceTokenizer
 from nemo.collections.nlp.data.datasets.glue_benchmark_dataset import output_modes, processors
@@ -231,7 +231,7 @@
     else:
         model = nemo_nlp.nm.trainables.huggingface.BERT(pretrained_model_name=args.pretrained_bert_model)
     model.restore_from(args.bert_checkpoint)
-    logging.info(f"model resotred from {args.bert_checkpoint}")
+    logging.info(f"model restored from {args.bert_checkpoint}")
 
 hidden_size = model.hidden_size
 
@@ -241,7 +241,7 @@
     glue_loss = MSELoss()
 else:
     pooler = SequenceClassifier(hidden_size=hidden_size, num_classes=num_labels, log_softmax=False)
-    glue_loss = CrossEntropyLoss()
+    glue_loss = CrossEntropyLossNM()
 
 
 def create_pipeline(
@@ -260,8 +260,6 @@ def create_pipeline(
         processor=processor,
         evaluate=evaluate,
         batch_size=batch_size,
-        # num_workers=0,
-        # local_rank=local_rank,
         tokenizer=tokenizer,
         data_dir=args.data_dir,
         max_seq_length=max_seq_length,

diff --git a/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_infer.py b/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_infer.py
@@ -23,7 +23,7 @@
 
 import nemo.collections.nlp.nm.trainables.joint_intent_slot.joint_intent_slot_nm
 from nemo import logging
-from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset import JointIntentSlotDataDesc
+from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset.data_descriptor import JointIntentSlotDataDesc
 
 # Parsing arguments
 parser = argparse.ArgumentParser(description='Joint-intent BERT')

diff --git a/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_infer_b1.py b/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_infer_b1.py
@@ -21,8 +21,8 @@
 
 import nemo.collections.nlp as nemo_nlp
 import nemo.collections.nlp.nm.trainables.joint_intent_slot.joint_intent_slot_nm
-from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset import JointIntentSlotDataDesc
-from nemo.collections.nlp.utils.common_nlp_utils import read_intent_slot_outputs
+from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset.data_descriptor import JointIntentSlotDataDesc
+from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset.inference_utils import read_intent_slot_outputs
 
 # Parsing arguments
 parser = argparse.ArgumentParser(description='Joint-intent BERT')

diff --git a/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_with_bert.py b/examples/nlp/intent_detection_slot_tagging/joint_intent_slot_with_bert.py
@@ -21,16 +21,18 @@
 import numpy as np
 from transformers import BertTokenizer
 
+import nemo
 import nemo.collections.nlp as nemo_nlp
-import nemo.collections.nlp.nm.data_layers.joint_intent_slot_datalayer
-import nemo.collections.nlp.nm.trainables.joint_intent_slot.joint_intent_slot_nm
 from nemo import logging
+from nemo.backends.pytorch.common.losses import CrossEntropyLossNM, LossAggregatorNM
 from nemo.collections.nlp.callbacks.joint_intent_slot_callback import eval_epochs_done_callback, eval_iter_callback
-from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset import JointIntentSlotDataDesc
+from nemo.collections.nlp.data.datasets.joint_intent_slot_dataset.data_descriptor import JointIntentSlotDataDesc
+from nemo.collections.nlp.nm.data_layers import BertJointIntentSlotDataLayer
+from nemo.core import CheckpointCallback, SimpleLossLoggerCallback
 from nemo.utils.lr_policies import get_lr_policy
 
 # Parsing arguments
-parser = argparse.ArgumentParser(description='Joint intent slot filling system with pretrained BERT')
+parser = argparse.ArgumentParser(description='Joint intent detection and slot filling with pre-trained BERT')
 parser.add_argument("--local_rank", default=None, type=int)
 parser.add_argument("--batch_size", default=128, type=int)
 parser.add_argument("--max_seq_length", default=50, type=int)
@@ -87,12 +89,10 @@
 nemo_nlp.huggingface.BERT.list_pretrained_models()
 """
 if args.bert_checkpoint and args.bert_config:
-    pretrained_bert_model = nemo.collections.nlp.nm.trainables.huggingface.BERT(config_filename=args.bert_config)
+    pretrained_bert_model = nemo_nlp.nm.trainables.huggingface.BERT(config_filename=args.bert_config)
     pretrained_bert_model.restore_from(args.bert_checkpoint)
 else:
-    pretrained_bert_model = nemo.collections.nlp.nm.trainables.huggingface.BERT(
-        pretrained_model_name=args.pretrained_bert_model
-    )
+    pretrained_bert_model = nemo_nlp.nm.trainables.huggingface.BERT(pretrained_model_name=args.pretrained_bert_model)
 
 hidden_size = pretrained_bert_model.hidden_size
 
@@ -101,31 +101,29 @@
 )
 
 # Create sentence classification loss on top
-classifier = nemo.collections.nlp.nm.trainables.joint_intent_slot.joint_intent_slot_nm.JointIntentSlotClassifier(
+classifier = nemo_nlp.nm.trainables.JointIntentSlotClassifier(
     hidden_size=hidden_size, num_intents=data_desc.num_intents, num_slots=data_desc.num_slots, dropout=args.fc_dropout
 )
 
 if args.class_balancing == 'weighted_loss':
-    # Using weighted loss will enable weighted loss for both intents and slots
-    # Use the intent_loss_weight hyperparameter to adjust intent loss to
-    # prevent overfitting or underfitting.
-    loss_fn = nemo_nlp.nm.losses.JointIntentSlotLoss(
-        num_slots=data_desc.num_slots,
-        slot_classes_loss_weights=data_desc.slot_weights,
-        intent_classes_loss_weights=data_desc.intent_weights,
-        intent_loss_weight=args.intent_loss_weight,
-    )
+    # To tackle imbalanced classes, you may use weighted loss
+    intent_loss_fn = CrossEntropyLossNM(logits_dim=2, weight=data_desc.intent_weights)
+    slot_loss_fn = CrossEntropyLossNM(logits_dim=3, weight=data_desc.intent_weights)
+
 else:
-    loss_fn = nemo_nlp.nm.losses.JointIntentSlotLoss(num_slots=data_desc.num_slots)
+    intent_loss_fn = CrossEntropyLossNM(logits_dim=2)
+    slot_loss_fn = CrossEntropyLossNM(logits_dim=3)
+
+total_loss_fn = LossAggregatorNM(num_inputs=2, weights=[args.intent_loss_weight, 1.0 - args.intent_loss_weight])
 
 
-def create_pipeline(num_samples=-1, batch_size=32, num_gpus=1, local_rank=0, mode='train'):
+def create_pipeline(num_samples=-1, batch_size=32, num_gpus=1, mode='train'):
     logging.info(f"Loading {mode} data...")
     data_file = f'{data_desc.data_dir}/{mode}.tsv'
     slot_file = f'{data_desc.data_dir}/{mode}_slots.tsv'
     shuffle = args.shuffle_data if mode == 'train' else False
 
-    data_layer = nemo.collections.nlp.nm.data_layers.joint_intent_slot_datalayer.BertJointIntentSlotDataLayer(
+    data_layer = BertJointIntentSlotDataLayer(
         input_file=data_file,
         slot_file=slot_file,
         pad_label=data_desc.pad_label,
@@ -155,35 +153,27 @@ def create_pipeline(num_samples=-1, batch_size=32, num_gpus=1, local_rank=0, mod
 
     intent_logits, slot_logits = classifier(hidden_states=hidden_states)
 
-    loss = loss_fn(
-        intent_logits=intent_logits, slot_logits=slot_logits, loss_mask=loss_mask, intents=intents, slots=slots
-    )
+    intent_loss = intent_loss_fn(logits=intent_logits, labels=intents)
+    slot_loss = slot_loss_fn(logits=slot_logits, labels=slots, loss_mask=loss_mask)
+    total_loss = total_loss_fn(loss_1=intent_loss, loss_2=slot_loss)
 
     if mode == 'train':
-        tensors_to_evaluate = [loss, intent_logits, slot_logits]
+        tensors_to_evaluate = [total_loss, intent_logits, slot_logits]
     else:
         tensors_to_evaluate = [intent_logits, slot_logits, intents, slots, subtokens_mask]
 
-    return tensors_to_evaluate, loss, steps_per_epoch, data_layer
+    return tensors_to_evaluate, total_loss, steps_per_epoch, data_layer
 
 
 train_tensors, train_loss, steps_per_epoch, _ = create_pipeline(
-    args.num_train_samples,
-    batch_size=args.batch_size,
-    num_gpus=args.num_gpus,
-    local_rank=args.local_rank,
-    mode=args.train_file_prefix,
+    args.num_train_samples, batch_size=args.batch_size, num_gpus=args.num_gpus, mode=args.train_file_prefix,
 )
 eval_tensors, _, _, data_layer = create_pipeline(
-    args.num_eval_samples,
-    batch_size=args.batch_size,
-    num_gpus=args.num_gpus,
-    local_rank=args.local_rank,
-    mode=args.eval_file_prefix,
+    args.num_eval_samples, batch_size=args.batch_size, num_gpus=args.num_gpus, mode=args.eval_file_prefix,
 )
 
 # Create callbacks for train and eval modes
-train_callback = nemo.core.SimpleLossLoggerCallback(
+train_callback = SimpleLossLoggerCallback(
     tensors=train_tensors,
     print_func=lambda x: str(np.round(x[0].item(), 3)),
     tb_writer=nf.tb_writer,
@@ -200,7 +190,7 @@ def create_pipeline(num_samples=-1, batch_size=32, num_gpus=1, local_rank=0, mod
 )
 
 # Create callback to save checkpoints
-ckpt_callback = nemo.core.CheckpointCallback(
+ckpt_callback = CheckpointCallback(
     folder=nf.checkpoint_dir, epoch_freq=args.save_epoch_freq, step_freq=args.save_step_freq
 )