Export ThutmoseTagger ITN Torchscript module
State: closed | Pull request: no
Labels: name: stale

Hello, I'm trying to export a ThutmoseTaggerModel using TorchScript. I was able to export the forward method, but not other methods. I was able to trace the forward function and run it in C++ by running this code:

```
from helpers import ITN_MODEL, instantiate_model_and_trainer
from omegaconf import DictConfig, OmegaConf
import torch
import omegaconf
from nemo.collections.nlp.data.text_normalization_as_tagging import bert_example

cfg_path = "/home/msancinetti/NeMo/examples/nlp/text_normalization_as_tagging/conf/thutmose_tagger_itn_config.yaml"
cfg = omegaconf.OmegaConf.load(cfg_path)
cfg["pretrained_model"] = "itn_en_thutmose_bert"
_, model = instantiate_model_and_trainer(cfg, ITN_MODEL, False)
model.eval()
model.to('cpu')
input_example = [torch.zeros([64,128], dtype=torch.int).to("cpu")] * 3
model.set_input_names(["input_1"])
model.export('itn_model.ts', verbose=True, input_example=input_example)
```
For this to work I had to:

- Add `~/NeMo/examples/nlp/text_normalization_as_tagging/` to my PYTHONPATH for the first import
- Add the following code to the ThutmoseTaggerModel class:
```
    @property
    def input_names(self):
        return self._input_names

    def set_input_names(self, new_value):
        self._input_names = new_value
```
(because the input_names property is necessary to run the export method)


However, the model.export method only traces the forward function. I would like to export other methods (that may not be traceable, so I think I have to use torch.jit.script). The problem is that if I try to script any method that uses "self", I get an error.

As a minimal example, I added this method to the ThutmoseTaggerModel class:

```
    def dummy(self):
        return self.num_labels
```

When I try to script that function, I get this error:

```
In [1]: torch.jit.script(model.dummy)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 1
----> 1 torch.jit.script(model.dummy)

File ~/miniconda3/envs/nemo_debug/lib/python3.8/site-packages/torch/jit/_script.py:1381, in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1379 if _rcb is None:
   1380     _rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
-> 1381 fn = torch._C._jit_script_compile(
   1382     qualified_name, ast, _rcb, get_default_args(obj)
   1383 )
   1384 # Forward docstrings
   1385 fn.__doc__ = obj.__doc__

RuntimeError: 
'Tensor (inferred)' object has no attribute or method 'num_labels'.:
  File "/home/msancinetti/miniconda3/envs/nemo_debug/lib/python3.8/site-packages/nemo/collections/nlp/models/text_normalization_as_tagging/thutmose_tagger.py", line 503
    def dummy(self):
        return self.num_labels
               ~~~~~~~~~~~~~~~ <--- HERE
```

I also get an error if I try to script the whole model with `torch.jit.script(model)`

_TL;DR:_ 
How do I export methods from ThutmoseTagger ITN model other than forward to run them in C++?

Any help will be appreciated!



hi @MarceloSancinetti , this was my model, but I do not know about torchscript so can't help with that. Basically, ThutmoseTagger is just a token classification bert-based model. If you have questions about meaning of its elements, you can ask

Hello @bene-ges, thank you for your response. I've been using your model for a while and it works very well, I'm just trying to figure out how to deploy it. I'll let you know if I have any questions about the model. 

@MarceloSancinetti, in case if you don't find other ways, you can try the following

you can rewrite this class in pure pytorch, inheriting from nn.Module
`class MyThutmoseModel(nn.Module):`
then include only `init` and `forward` methods.

If you unpack the nemo checkpoint
`tar xvf itn_en_thutmose_bert.nemo`
you will see a number of files packed with the model,
among them `model_weights.ckpt`
then you can initialize your class and load pretrained parameters into it:
```
    model = MyThutmoseModel()
    state_dict = torch.load("model_weights.ckpt", map_location=device)
    model.load_state_dict(state_dict, strict=False)
```
and so you will get a pure pytorch model that you can use for inference.     
`model_weights.ckpt` - is just a dictionary - look at its keys to better understand which class members it will affect. 

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

This issue was closed because it has been inactive for 7 days since being marked as stale.

@MarceloSancinetti Hello! I've faced similar need to deploy TT model
I wonder if you have implemented this solution?
Can I find your implementation somewhere?

Hello @Dan1aR, I ended up exporting only the forward function and re implementing the pre and post processing part of the inference function in C++ myself.

@MarceloSancinetti, Thank's a lot for quick reply!
Can you please share your solution? That would save me a lot of trouble
Or point me to some resources, where i could find one?