Exporting PunctuationCapitalizationModel model to onnx
State: closed | Pull request: no

I am trying to export the nemo_nlp PunctuationCapitalizationModel to onnx format. 
I was able to do so using this block of code:
```
model = PunctuationCapitalizationModel.restore_from(restore_path="punctuation_en_distilbert.nemo")
model.export("punctuation_en_distilbert.onnx")
```

This works fine and I get the onnx file and I am able to load it using onnx.load and check it using onnx.checker.check_model.

The issues are that exported onnx model has only 2 inputs: 
1.  

```
[name: "input_ids"
type {
  tensor_type {
    elem_type: 7
    shape {
      dim {
        dim_param: "input_ids_dynamic_axes_1"
      }
      dim {
        dim_param: "input_ids_dynamic_axes_2"
      }
    }
  }
}
, name: "attention_mask"
type {
  tensor_type {
    elem_type: 7
    shape {
      dim {
        dim_param: "attention_mask_dynamic_axes_1"
      }
      dim {
        dim_param: "attention_mask_dynamic_axes_2"
      }
    }
  }
}
]
```

and 

2. if I call it like:

```
options = onnxruntime.SessionOptions()
options.intra_op_num_threads = 1
sess_punct = onnxruntime.InferenceSession("punctuation_en_distilbert.onnx", options, providers=["CPUExecutionProvider"])

modelPunct = PunctuationCapitalizationModel.restore_from("punctuation_en_distilbert.nemo")
data_loader = create_infer_dataloader(modelPunct.tokenizer, queries)

for batch in data_loader:
    input_ids, input_type_ids, input_mask, subtokens_mask = batch
    ort_inputs = {"input_ids": to_numpy(input_ids),"attention_mask":to_numpy(input_mask)}
    punct_logits, capit_logits = sess_punct.run(None, input_feed = ort_inputs)
```

where `create_infer_dataloader()` is 

```
def create_infer_dataloader(tokenizer, queries):
    batch_size = len(queries)

    dataset = BertPunctuationCapitalizationInferDataset(
        tokenizer=tokenizer, queries=queries, max_seq_length=512
    )
    return torch.utils.data.DataLoader(
        dataset=dataset,
        collate_fn=dataset.collate_fn,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=False,
    )
```

it gives me an error: `onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float)) , expected: (tensor(int64))`

So the question is what is wrong with the exported model? 
Why the first question is why there are only 2 input parameters instead of 3, that we can see in different tutorials or source code of the PunctuationCapitalizationModel model  and why data types are messed

And if I change `punctuation_en_distilbert.nemo` to `punctuation_en_bert.nemo` than .onnx model has 3 inputs, but the issue with 
`onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float)) , expected: (tensor(int64))` for `"attention_mask"` still exist

In forward DistilBERT takes only input_ids and attention_mask (no token_type_ids as in BERT).
@borisfom could you help with the ONNX part?

Well, finally I ended up with the next code:

Export .onnx and config yaml:

```
model = PunctuationCapitalizationModel.restore_from(restore_path="punctuation_en_distilbert.nemo")
model.export("punctuation_en_distilbert.onnx")

model = PunctuationCapitalizationModel.restore_from(restore_path="punctuation_en_distilbert.nemo", return_config = True)
textfile = open("punctuation_en_distilbert.yaml", "w")
textfile.write(str(OmegaConf.to_yaml(model)))
textfile.close()
```

Export tokenizer:
```
from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('distilbert-base-uncased')
tokenizer.save_pretrained("./my_model_directory/")
```
and rename `tokenizer_config.json` to `distilbert-base-uncased.json`
Also need to download https://huggingface.co/distilbert-base-uncased/blob/main/tokenizer.json to the my_model_directory folder and change `distilbert-base-uncased.json` content a bit, set `"tokenizer_file"` to `"tokenizer.json"`
Now make processing with onnx:

```
import numpy as np
import torch
import yaml
from nemo.collections.common.tokenizers import AutoTokenizer
from nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset import \
    BertPunctuationCapitalizationInferDataset
import onnxruntime

import logging
logging.getLogger('nemo_logger').setLevel(logging.ERROR)

def create_infer_dataloader(tokenizer, queries):
    batch_size = len(queries)

    dataset = BertPunctuationCapitalizationInferDataset(
        tokenizer=tokenizer, queries=queries, max_seq_length=512
    )
    return torch.utils.data.DataLoader(
        dataset=dataset,
        collate_fn=dataset.collate_fn,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=False,
    )

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

queries = ["by the end of no such thing the audience like beatrice has a watchful affection for the monster", \
           "let me guess you're the kind of guy that ignores the rules cause it makes you feel in control am i right \
           you're not wrong you think that's cute do you think it's cute"]

tokenizer = AutoTokenizer(pretrained_model_name="./my_model_directory/distilbert-base-uncased.json", \
                          vocab_file="./my_model_directory/vocab.txt")

data_loader = create_infer_dataloader(tokenizer, queries)

with open("punctuation_en_distilbert.yaml") as f:
    params = yaml.safe_load(f)

all_punct_preds = []
all_capit_preds = []

options = onnxruntime.SessionOptions()
options.intra_op_num_threads = 1

sess_punct = onnxruntime.InferenceSession("punctuation_en_distilbert.onnx", options, providers=["CPUExecutionProvider"])  

# most of processing was taken from 
# https://github.com/NVIDIA/NeMo/blob/5839aee402f314aa413b28e9042b1e1cac10a114/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py#L403
for batch in data_loader:
    input_ids, input_type_ids, input_mask, subtokens_mask = batch
    input_mask = input_mask.type(torch.int64)  # why?
    ort_inputs = {"attention_mask": to_numpy(input_mask), "input_ids": to_numpy(input_ids)}

    punct_logits, capit_logits = sess_punct.run(None, input_feed = ort_inputs)

    subtokens_mask = subtokens_mask > 0.5

    punct_preds = [
        list(p_l[subtokens_mask[i]]) for i, p_l in enumerate(np.argmax(punct_logits, axis=-1))
    ]
    capit_preds = [
        list(c_l[subtokens_mask[i]]) for i, c_l in enumerate(np.argmax(capit_logits, axis=-1))
    ]

    all_punct_preds.extend(punct_preds)
    all_capit_preds.extend(capit_preds)

punct_ids_to_labels = {v: k for k, v in params['punct_label_ids'].items()}
capit_ids_to_labels = {v: k for k, v in params['capit_label_ids'].items()}

queries = [q.strip().split() for q in queries]

result = []

for i, query in enumerate(queries):
    punct_preds = all_punct_preds[i]
    capit_preds = all_capit_preds[i]
    if len(query) != len(punct_preds):
        # removing the end of phrase punctuation of the truncated segment
        punct_preds[-1] = 0
        max_len = len(punct_preds)
        query = query[:max_len]

    query_with_punct_and_capit = ''
    for j, word in enumerate(query):
        punct_label = punct_ids_to_labels[punct_preds[j]]
        capit_label = capit_ids_to_labels[capit_preds[j]]

        if capit_label != params['dataset']['pad_label']:
            word = word.capitalize()
        query_with_punct_and_capit += word
        if punct_label != params['dataset']['pad_label']:
            query_with_punct_and_capit += punct_label
        query_with_punct_and_capit += ' '

    result.append(query_with_punct_and_capit.strip())
```

not sure if I did everything correct with tokenizer export and import, but it works 

I am still facing the same issue after following the steps. ValueError: Model requires 3 inputs. Input Feed contains 2 
Did I miss something?