Save DistilBert Model and Convert
Author: ffalkenberg
Pull request: no

Hello,

I want to save a transformer model to [TensorFlow Lite](https://www.tensorflow.org/lite), in order to put in on a mobile device.

```
from transformers import AutoModelForTokenClassification, AutoTokenizer
model = AutoModelForTokenClassification.from_pretrained("elastic/distilbert-base-uncased-finetuned-conll03-english")
type(model)
```
transformers.models.distilbert.modeling_distilbert.DistilBertForTokenClassification

Then i tried to save the model in tensorflow format:

```tf.saved_model.save(model, saved_model_dir)```
@jplu  https://github.com/huggingface/transformers/issues/6864#issuecomment-690086932

Here, I get the following Error:

ValueError: Expected an object of type `Trackable`, such as `tf.Module` or a subclass of the `Trackable` class, for export. Got DistilBertForTokenClassification(
  (distilbert): DistilBertModel ...



As a second step I was planning to convert it tflite:

```
converter = tf.lite.TFLiteConverter.from_saved_model(model_path) 
tflite_model = converter.convert()
```

What am I doing wrong and is there an easier way to do that?



Comments:
Hi, you are using a pytorch model here. `AutoModelForTokenClassification` will load the pt version of `DistilBert`. For TF one should use `TFAutoModelForTokenClassification`.

> 
thank you alot @patil-suraj 👍  , i am using the following now, which works fine:

```
model = TFDistilBertForTokenClassification.from_pretrained(model_path, from_pt=True)
inputs_1 = tokenizer("Hugging Face Inc. is a company based in New York City.", return_tensors="tf")
outputs = model(inputs_1["input_ids"])
tf.saved_model.save(model, "/path/to/file")
```
The import works as in the following, but ...

```
loaded = tf.saved_model.load("/path/to/file")
inference_func = loaded.signatures["serving_default"]
inputs_2 = tokenizer("Hugging Face Inc", return_tensors="tf")

outputs_1 = inference_func(input_ids=inputs_1["input_ids"])
outputs_2 = inference_func(input_ids=inputs_2["input_ids"])
```
i can only do infer for inputs_2 not for inputs_1
I get the following error:

```
InvalidArgumentError:  Incompatible shapes: [1,15,768] vs. [1,5,768]
	 [[node tf_distil_bert_for_token_classification_3/distilbert/embeddings/add
 (defined at /tmp/ipykernel_377675/508456377.py:1)
]] [Op:__inference_signature_wrapper_541974]

Errors may have originated from an input operation.
Input Source operations connected to node tf_distil_bert_for_token_classification_3/distilbert/embeddings/add:
In[0] tf_distil_bert_for_token_classification_3/distilbert/embeddings/Identity:	
In[1] tf_distil_bert_for_token_classification_3/distilbert/embeddings/Identity_1:
```
how can this be fixed and what is the reason for this?

Hi @Schnittchenkraus! The issue is that, in essence, `loaded` (the TF loaded object) is a constrained version of `model` (an instance of a Hugging Face class). `loaded` stores the computation graph, which is generated with the first call, while `model` is a bunch of TF operations that can accept variable-length tensors.

Now, to the solution. Typically, the correct solution revolves around defining signatures (see [this doc](https://www.tensorflow.org/api_docs/python/tf/saved_model/save)), where you can define inputs with variable length. However, our code does not fit the requirements for this, so it is not doable. 

The less correct solution is to make sure the computation graph and the inputs are constrained to a fixed length (for instance, the maximum sequence length, 512). To do that, you have to pad the inputs and override the first input to the model. Without getting into too many details, if you run the script below while changing [this line](https://github.com/huggingface/transformers/blob/master/src/transformers/file_utils.py#L312) to `DUMMY_INPUTS = [[1]*512]`, you should be able to get what you want.

This is overly convoluted, but we are working on making it easier :)

<details><summary>Working script</summary>

```python
from functools import partial
import tensorflow as tf
from transformers import TFDistilBertForTokenClassification, AutoTokenizer

model_path = "elastic/distilbert-base-uncased-finetuned-conll03-english"
model = TFDistilBertForTokenClassification.from_pretrained(model_path, from_pt=True)
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenize_fn = partial(tokenizer, return_tensors="tf", padding="max_length", truncation=True)
inputs_1 = tokenize_fn("Hugging Face Inc. is a company based in New York City.")
inputs_2 = tokenize_fn("Hugging Face Inc.")

outputs_2 = model(inputs_2)
outputs_1 = model(inputs_1)

tf.saved_model.save(model, "~/test_model")
loaded = tf.saved_model.load("~/test_model")
inference_func = loaded.signatures["serving_default"]

outputs_1 = inference_func(input_ids=inputs_1["input_ids"])
outputs_2 = inference_func(input_ids=inputs_2["input_ids"])
```

</details>

This works nicely, thank you alot @gante !