I am struggling to understand how to perform inference with a pre-trained HuggingFace model loaded as a TensorFlow Keras model. Context In my case, I am trying to fine-tune a pre-trained DistilBert classifier. I have something as follows to preprocess my data and load/train my model: from transformers import TFAutoModelForSequenceClassification tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased") model = TFAutoModelForSequenceClassification.from_pretrained( "distilbert-base-uncased", num_labels=2, id2label=id2label, label2id=label2id ) # add another layer tf_train = model.prepare_tf_dataset(question_train_test_split['train'], batch_size=16, shuffle=True, tokenizer=tokenizer) model.compile(optimizer=tf.keras.optimizers.Adam(2e-5)) # freeze the first transformer layer of model model.layers[0].trainable=False print('Model Architecture:') print(model.summary()) model.fit(tf_train, epochs=3) Where question_train_test_split is a HuggingFace Dataset object instance. This snippet of code works perfectly as expected, it loads the HuggingFace model as a tf.keras layer. This even trains properly with the .fit method. However, I am having issues when I want to perform predictions. I understand that I need to tokenize my string input, however, I would like to load the tokenizer as a tf.keras layer. I've looked everywhere for a way to do this and I was not able to find any way. Ideally, I would like something like this: user_input = 'When were the Beatles formed?' model_input = tokenizer(user_input) # THIS HF TOKENIZER SHOULD BE A tf.keras LAYER model = model(model_input) This is so that I can save the entire model (with the tokenizer and the transformer layers + the classifier layers) into a TensorFlow SavedModel. If there's any pointers out there to convert a HuggingFace tokenizer into a TensorFlow Keras layer, I would appreciate the pointer. Continue reading...
