Cannot run another HuggingFace Model

[ierezell]

I saved an official distilbert-base-multilingual-cased model in the tensorflow savedModel format with the code given :

import os
import tensorflow as tf
from transformers import (TFDistilBertModel, DistilBertConfig, DistilBertTokenizer, DistilBertModel)

model_name = "distilbert-base-multilingual-cased"
if not os.path.exists("./models/" + model_name):
    os.makedirs("./models/" + model_name)

tokenizer = DistilBertTokenizer.from_pretrained(model_name)
tokenizer.save_pretrained("./models/" + model_name)

pt_model = DistilBertModel.from_pretrained(model_name)
pt_model.save_pretrained("./models/" + model_name)
config = DistilBertConfig.from_json_file("./models/" + model_name +
                                         "/config.json")
model = TFDistilBertModel.from_pretrained("./models/" + model_name,
                                          from_pt=True,
                                          config=config)

concrete_function = tf.function(model.call).get_concrete_function([
    tf.TensorSpec([None, 384], tf.int32, name="input_ids"),
    tf.TensorSpec([None, 384], tf.int32, name="attention_mask")
])
tf.saved_model.save(
    model,
    "./models/" + model_name,
    signatures=concrete_function)
)

I then load the model and the tokenizer perfectly with

const model = await tf.node.loadSavedModel(path_to_distilbert);
const tokenizer = await BertWordPieceTokenizer.fromOptions({
    vocabFile: path.join(chemin, "vocab.txt"),
    lowercase: true,
  });

But when I want to predict to get simply the embeddings of a sentence with :

const result = tf.tidy(() => {
    const inputTensor = tf.tensor(tokens.ids, undefined, "int32");
    const maskTensor = tf.tensor(tokens.attentionMask, undefined, "int32");
    return model.predict({
      ["input_ids"]: inputTensor,
      ["attention_mask"]: maskTensor,
    }) as tf.NamedTensorMap;
  });
  console.log(result);

I got Incompatibles shapes :

Error: Session fail to run with error: [_Derived_]{{function_node __inference_call_4832}} {{function_node __inference_call_4832}} Incompatible shapes: [9,768] vs. [1,384,768]
         [[{{node distilbert/embeddings/add}}]]
         [[StatefulPartitionedCall/StatefulPartitionedCall]]
    at NodeJSKernelBackend.runSavedModel (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-node/dist/nodejs_kernel_backend.js:1580:43)
    at TFSavedModel.predict (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-node/dist/saved_model.js:347:52)
    at /mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/dist/test.js:31:22
    at /mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:3134:22
    at Engine.scopedRun (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:3144:23)
    at Engine.tidy (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:3133:21)
    at Object.tidy (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:7777:19)
    at test (/mnt/Documents/Projets/BotPress/R_D/R_D_MLjs_pipeline/sdk/dist/test.js:27:23)

I tried to expandDims, to remove the 384 in the concreteFunction, to same without concrete function, no way to get inference...

I tried with node 13.13.0

Thanks in advance, I know the problem is s abit borderline...

EDIT : By putting tf.TensorSpec([None, None], ...) I was able to load the model, I'm letting this for those who try to load a HuggingFace model in Nodejs.

Still it's not the way to do it, it works in this special case but I know it could be done in a cleaner way