Repository: NVIDIA/NeMo

Issue:
PunctuationCapitalizationModel fails on a sentence with specific length
**Describe the bug**

`nemo.collections.nlp.models.PunctuationCapitalizationModel.add_punctuation_capitalization` fails when inferencing on a sentence with specific length

**Steps/Code to reproduce bug**

```
import nemo.collections.nlp
punct = nemo.collections.nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name="punctuation_en_bert").eval().cuda()
punct.add_punctuation_capitalization(["here we are waiting for the meat to fry how long it takes for it to rummage for about a minute as it browns a big drop on the other side"])
```

Results in:
```
[NeMo I 2021-10-13 12:05:07 punctuation_capitalization_model:575] Using batch size 1 for inference
[NeMo I 2021-10-13 12:05:07 punctuation_capitalization_dataset:533] Max length: 34
[NeMo I 2021-10-13 12:05:07 data_preprocessing:358] Some stats of the lengths of the sequences:
[NeMo I 2021-10-13 12:05:07 data_preprocessing:360] Min: 32 |                  Max: 32 |                  Mean: 32.0 |                  Median: 32.0
[NeMo I 2021-10-13 12:05:07 data_preprocessing:366] 75 percentile: 32.00
[NeMo I 2021-10-13 12:05:07 data_preprocessing:367] 99 percentile: 32.00
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-762f40629049> in <module>
----> 1 punct.add_punctuation_capitalization(["here we are waiting for the meat to fry how long it takes for it to rummage for about a minute as it browns a big drop on the other side"])

~/miniconda3/envs/v2v-en/lib/python3.9/site-packages/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py in add_punctuation_capitalization(self, queries, batch_size, max_seq_length, step, margin, return_labels)
    578         try:
    579             self.eval()
--> 580             infer_datalayer = self._setup_infer_dataloader(queries, batch_size, max_seq_length, step, margin)
    581             # Predicted labels for queries. List of labels for every query
    582             all_punct_preds: List[List[int]] = [[] for _ in queries]

~/miniconda3/envs/v2v-en/lib/python3.9/site-packages/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py in _setup_infer_dataloader(self, queries, batch_size, max_seq_length, step, margin)
    361             margin = self._cfg.dataset.margin
    362
--> 363         dataset = BertPunctuationCapitalizationInferDataset(
    364             tokenizer=self.tokenizer, queries=queries, max_seq_length=max_seq_length, step=step, margin=margin
    365         )

~/miniconda3/envs/v2v-en/lib/python3.9/site-packages/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py in __init__(self, queries, tokenizer, max_seq_length, step, margin)
    638         self, queries: List[str], tokenizer: TokenizerSpec, max_seq_length: int = 128, step: int = 32, margin: int = 16
    639     ):
--> 640         features = get_features_infer(
    641             queries=queries, max_seq_length=max_seq_length, tokenizer=tokenizer, step=step, margin=margin
    642         )

~/miniconda3/envs/v2v-en/lib/python3.9/site-packages/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py in get_features_infer(queries, tokenizer, max_seq_length, step, margin)
    542     for q_i, query_st in enumerate(st):
    543         q_inp_ids, q_segment_ids, q_subtokens_mask, q_inp_mask, q_quantities_of_preceding_words = [], [], [], [], []
--> 544         for i in range(0, max(len(query_st), length) - length + step, step):
    545             subtokens = [tokenizer.cls_token] + query_st[i : i + length] + [tokenizer.sep_token]
    546             q_inp_ids.append(tokenizer.tokens_to_ids(subtokens))

ValueError: range() arg 3 must not be zero
```


**Expected behavior**

Method doesn't fail and return punctuated and capitalized result

**Environment overview (please complete the following information)**

Local machine
nemo installed with `pip install nemo_toolkit['all']`

```
$ pip freeze | grep nemo
nemo-toolkit==1.4.0
```

**Environment details**

Ubuntu 18.04.5 LTS
PyTorch 1.9.1
Python 3.9.7

**Additional context**

It seems like this line of code can result in step being zero: https://github.com/NVIDIA/NeMo/blame/main/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py#L536



Discussion:
@PeganovAnton could you please take a look into this?

Resolving patch:
diff --git a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
--- a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
+++ b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
@@ -529,11 +529,18 @@ def get_features_infer(
         st.append(subtokens)
         stm.append(subtokens_mask)
     _check_max_seq_length_and_margin_and_step(max_seq_length, margin, step)
-    max_seq_length = min(max_seq_length, max(sent_lengths) + 2)
+    if max_seq_length > max(sent_lengths) + 2:
+        max_seq_length = max(sent_lengths) + 2
+        # If `max_seq_length` is greater than maximum length of input query, parameters ``margin`` and ``step`` are
+        # not used will not be used.
+        step = 1
+        # Maximum number of word subtokens in segment. The first and the last tokens in segment are CLS and EOS
+        length = max_seq_length - 2
+    else:
+        # Maximum number of word subtokens in segment. The first and the last tokens in segment are CLS and EOS
+        length = max_seq_length - 2
+        step = min(length - margin * 2, step)
     logging.info(f'Max length: {max_seq_length}')
-    # Maximum number of word subtokens in segment. The first and the last tokens in segment are CLS and EOS
-    length = max_seq_length - 2
-    step = min(length - margin * 2, step)
     get_stats(sent_lengths)
     all_input_ids, all_segment_ids, all_subtokens_mask, all_input_mask, all_input_mask = [], [], [], [], []
     all_quantities_of_preceding_words, all_query_ids, all_is_first, all_is_last = [], [], [], []


Test patch: