Ptuning Vocab error
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

I am trying to do Ptuning with Nemo with my custom Megtraon gpt2 and tokenizers that were trained using huggingface tokenizers. But when I provide all the configs and run, I get a vocabulary error.

```
root@ptuning-7d8497bf49-r8spb:/workspace/sean/ptuning# PYTHONPATH=NeMo-feature_ptune python ptune_misogyny.py
NOTE! Installing ujson may make loading annotations faster.
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py:107: LightningDeprecationWarning: Argument `num_nodes` in `DDPPlugin` is deprecated in v1.4, and will be removed in v1.6. Notice that it will be overriden by the trainer setting.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py:113: LightningDeprecationWarning: Argument `sync_batchnorm` in `DDPPlugin` is deprecated in v1.4, and will be removed in v1.6. Notice that it will be overriden by the trainer setting.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:324: LightningDeprecationWarning: Passing <nemo.collections.nlp.parts.nlp_overrides.NLPDDPPlugin object at 0x7fc6c7dcd610> `strategy` to the `plugins` flag in Trainer has been deprecated in v1.5 and will be removed in v1.7. Use `Trainer(strategy=<nemo.collections.nlp.parts.nlp_overrides.NLPDDPPlugin object at 0x7fc6c7dcd610>)` instead.
      rank_zero_deprecation(
    
Using 16bit native Automatic Mixed Precision (AMP)
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:48: LightningDeprecationWarning: Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7. Use `max_steps = -1` instead.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/callback_connector.py:147: LightningDeprecationWarning: Setting `Trainer(checkpoint_callback=False)` is deprecated in v1.5 and will be removed in v1.7. Please consider using `Trainer(enable_checkpointing=False)`.
      rank_zero_deprecation(
    
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
[NeMo I 2022-01-28 06:50:58 exp_manager:283] Experiments will be logged at /workspace/sean/ptuning/nemo_experiments/PTuneTextClassification/2022-01-28_06-50-58
[NeMo I 2022-01-28 06:50:58 exp_manager:648] TensorboardLogger has been set up
[NeMo W 2022-01-28 06:50:58 exp_manager:901] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to -1. Please ensure that max_steps will run for at least 1 epochs to ensure that checkpointing will not error out.
[NeMo W 2022-01-28 06:50:58 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:243: LightningDeprecationWarning: `ModelCheckpoint(every_n_val_epochs)` is deprecated in v1.4 and will be removed in v1.6. Please use `every_n_epochs` instead.
      rank_zero_deprecation(
    
[NeMo I 2022-01-28 06:50:58 ptune_text_classification_model:375] Dataloader config or file_path for the validation is missing, so no data loader for test is created!
[NeMo I 2022-01-28 06:50:58 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /workspace/sean/ptuning/tamil_gpt_tokenizer/vocab.json
[NeMo I 2022-01-28 06:50:58 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /workspace/sean/ptuning/tamil_gpt_tokenizer/vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:51:08 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /tmp/tmpyb_4gkds/749983edcb2a4816bacb8628b0e6beeb_vocab.json
[NeMo I 2022-01-28 06:51:08 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /tmp/tmpyb_4gkds/749983edcb2a4816bacb8628b0e6beeb_vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:51:08 megatron_gpt_model:754] Padded vocab_size: 32000, original vocab_size: 32000, dummy tokens: 0.
[NeMo I 2022-01-28 06:51:09 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /tmp/tmpyb_4gkds/749983edcb2a4816bacb8628b0e6beeb_vocab.json
[NeMo I 2022-01-28 06:51:09 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /tmp/tmpyb_4gkds/749983edcb2a4816bacb8628b0e6beeb_vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:51:09 megatron_gpt_model:754] Padded vocab_size: 32000, original vocab_size: 32000, dummy tokens: 0.
[NeMo I 2022-01-28 06:51:10 save_restore_connector:154] Model MegatronGPTModel was successfully restored from /workspace/sean/ptuning/tamil_gpt_344m.nemo.
[NeMo E 2022-01-28 06:51:10 ptune_text_classification_model:105] class NOT is not part of the vocabulary. Please add it to your vocab
[NeMo E 2022-01-28 06:51:10 ptune_text_classification_model:105] class OFF is not part of the vocabulary. Please add it to your vocab
Traceback (most recent call last):
  File "ptune_misogyny.py", line 41, in <module>
    model_ptune = PTuneTextClassificationModel(cfg=config.model, trainer=trainer)
  File "/workspace/sean/ptuning/NeMo-feature_ptune/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py", line 106, in __init__
    self.allowed_vocab_ids = set(self.vocab[token_wrapper(k)] for k in cfg.dataset.classes)
  File "/workspace/sean/ptuning/NeMo-feature_ptune/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py", line 106, in <genexpr>
    self.allowed_vocab_ids = set(self.vocab[token_wrapper(k)] for k in cfg.dataset.classes)
KeyError: 'ĠNOT'
root@ptuning-7d8497bf49-r8spb:/workspace/sean/ptuning# PYTHONPATH=NeMo-feature_ptune python ptune_misogyny.py
NOTE! Installing ujson may make loading annotations faster.
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py:107: LightningDeprecationWarning: Argument `num_nodes` in `DDPPlugin` is deprecated in v1.4, and will be removed in v1.6. Notice that it will be overriden by the trainer setting.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py:113: LightningDeprecationWarning: Argument `sync_batchnorm` in `DDPPlugin` is deprecated in v1.4, and will be removed in v1.6. Notice that it will be overriden by the trainer setting.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:324: LightningDeprecationWarning: Passing <nemo.collections.nlp.parts.nlp_overrides.NLPDDPPlugin object at 0x7fe7e13cd640> `strategy` to the `plugins` flag in Trainer has been deprecated in v1.5 and will be removed in v1.7. Use `Trainer(strategy=<nemo.collections.nlp.parts.nlp_overrides.NLPDDPPlugin object at 0x7fe7e13cd640>)` instead.
      rank_zero_deprecation(
    
Using 16bit native Automatic Mixed Precision (AMP)
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:48: LightningDeprecationWarning: Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7. Use `max_steps = -1` instead.
      rank_zero_deprecation(
    
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/callback_connector.py:147: LightningDeprecationWarning: Setting `Trainer(checkpoint_callback=False)` is deprecated in v1.5 and will be removed in v1.7. Please consider using `Trainer(enable_checkpointing=False)`.
      rank_zero_deprecation(
    
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
[NeMo I 2022-01-28 06:52:33 exp_manager:283] Experiments will be logged at /workspace/sean/ptuning/nemo_experiments/PTuneTextClassification/2022-01-28_06-52-33
[NeMo I 2022-01-28 06:52:33 exp_manager:648] TensorboardLogger has been set up
[NeMo W 2022-01-28 06:52:33 exp_manager:901] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to -1. Please ensure that max_steps will run for at least 1 epochs to ensure that checkpointing will not error out.
[NeMo W 2022-01-28 06:52:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:243: LightningDeprecationWarning: `ModelCheckpoint(every_n_val_epochs)` is deprecated in v1.4 and will be removed in v1.6. Please use `every_n_epochs` instead.
      rank_zero_deprecation(
    
[NeMo I 2022-01-28 06:52:33 ptune_text_classification_model:375] Dataloader config or file_path for the validation is missing, so no data loader for test is created!
[NeMo I 2022-01-28 06:52:34 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /workspace/sean/ptuning/tamil_gpt_tokenizer/vocab.json
[NeMo I 2022-01-28 06:52:34 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /workspace/sean/ptuning/tamil_gpt_tokenizer/vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:52:44 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /tmp/tmpfxcnfv61/749983edcb2a4816bacb8628b0e6beeb_vocab.json
[NeMo I 2022-01-28 06:52:44 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /tmp/tmpfxcnfv61/749983edcb2a4816bacb8628b0e6beeb_vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:52:44 megatron_gpt_model:754] Padded vocab_size: 32000, original vocab_size: 32000, dummy tokens: 0.
[NeMo I 2022-01-28 06:52:45 tokenizer_utils:193] Getting Megatron tokenizer for pretrained model name: huggingface_model_tamil and custom vocab file: /tmp/tmpfxcnfv61/749983edcb2a4816bacb8628b0e6beeb_vocab.json
[NeMo I 2022-01-28 06:52:45 tokenizer_utils:125] Getting HuggingFace AutoTokenizer with pretrained_model_name: huggingface_model_tamil, vocab_file: /tmp/tmpfxcnfv61/749983edcb2a4816bacb8628b0e6beeb_vocab.json, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2022-01-28 06:52:45 megatron_gpt_model:754] Padded vocab_size: 32000, original vocab_size: 32000, dummy tokens: 0.
[NeMo I 2022-01-28 06:52:45 save_restore_connector:154] Model MegatronGPTModel was successfully restored from /workspace/sean/ptuning/tamil_gpt_344m.nemo.
[NeMo E 2022-01-28 06:52:46 ptune_text_classification_model:105] class தாக்குதல்-அல்ல is not part of the vocabulary. Please add it to your vocab
[NeMo E 2022-01-28 06:52:46 ptune_text_classification_model:105] class தாக்குதல் is not part of the vocabulary. Please add it to your vocab
Traceback (most recent call last):
  File "ptune_misogyny.py", line 41, in <module>
    model_ptune = PTuneTextClassificationModel(cfg=config.model, trainer=trainer)
  File "/workspace/sean/ptuning/NeMo-feature_ptune/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py", line 106, in __init__
    self.allowed_vocab_ids = set(self.vocab[token_wrapper(k)] for k in cfg.dataset.classes)
  File "/workspace/sean/ptuning/NeMo-feature_ptune/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py", line 106, in <genexpr>
    self.allowed_vocab_ids = set(self.vocab[token_wrapper(k)] for k in cfg.dataset.classes)
KeyError: 'Ġதாக்குதல்-அல்ல'
```

From the error trace, I can infer that at first, the model is loading the custom vocab file, but on the second time it is loading the default vocab file(I assume that is English), so I tried replacing the classes to English but still I am getting same error


**Steps/Code to reproduce bug**

My configurations
```


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
  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:
  tensor_model_parallel_size: 1 # tensor model parallel size used in the LM model
  seed: 1234
  nemo_path: null # filename to save the model and associated artifacts to .nemo file
  use_lm_finetune: False  # whether fine tune the language model
  pseudo_token: '[PROMPT]' # pseudo prompt tokens

  tokenizer:
    library: 'megatron'
    type: 'GPT2BPETokenizer'
    model: "huggingface_model_tamil"
    vocab_file: "tamil_gpt_tokenizer/vocab.json"
    merge_file: "tamil_gpt_tokenizer/merges.txt" 

  language_model:
    nemo_file: tamil_gpt_344m.nemo 

  prompt_encoder:
    template: [3, 3, 0]
    dropout: 0.0
    num_layers: 2

  dataset:
    classes: ['தாக்குதல்-அல்ல','தாக்குதல்']                    #['positive', 'neutral', 'negative']
#['எதுவும்-இல்லை','தவறான-நடத்தை','எதிர்-பேச்சு','பெண் -ெறுப்பு','இனவெறி','நம்பிக்கை-பேச்சு','ஓரினச்சேர்க்கை-வெறுப்பு','டிரான்ஸ்ஃபோபிக்']
  train_ds:
    file_path: data/offensive/train.txt
    batch_size: 8
    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: 8
    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: data/offensive/test.txt
    batch_size: 8
    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: 1e-5
    # optimizer arguments
    betas: [0.9, 0.999]
    weight_decay: 0.0005

    # 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: [
    'பிராமின்ஸ்லையும் மனசாட்சி உள்ள ஒரு மனுஷன் இருக்கான் யா சூப்பர் தல.',
    'ஆமா அவன் மலையாளி பஞ்சம் பிழைக்க வந்த வந்தேரி மலையாள கிறுத்துவ தேவிடியா மவன் ',
    'மதிப்பிற்குறிய ஆனந்த் சீனிவாசன் போன்ற முற்போக்கான எண்ணம் உள்ள பிராமண சமூதாயத்தில் தங்களை வாழ்த்துவதில் பெருமகிழ்சி கொள்கிறேன்.; அப்பிராமண பள்ளியில் பாதிக்கப்பட்டவர்கள் பிராமண குழந்தைகளும் இருக்கும் போது இந்த புகார்கள் எப்படி பிராமண எதிர்பாகும் எனபதை பெரியமூளை உள்ள  மதுவாந்திக்கு புரியவைக்க; யாராலும் முடியாது; அப்படியா பட்ட மூளை....',
  ]

exp_manager:
  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
  name: "PTuneTextClassification"  # 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

``````



**Expected behavior**

Ptuning should start

**Environment overview (please complete the following information)**

 - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
 - Method of NeMo install: [pip install or from source]. Please specify exact commands you used to install.
 - If method of install is [Docker], provide `docker pull` & `docker run` commands used

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

Add any other context about the problem here.
Example: GPU model