Save DistilBert Model and Convert

[ffalkenberg]

Hello,

I want to save a transformer model to TensorFlow 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?

[patil-suraj]

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

[ffalkenberg]

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?

[gante]

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), 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 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 :)

Working script

```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"]) ```

[ffalkenberg]

This works nicely, thank you alot @gante !