[Python] Exporting a Bert-based PyTorch model to CoreML. How can I make the CoreML model work...

I use the code below to export a Bert-based PyTorch model to CoreML.

Since I used

dummy_input = tokenizer("A French fan", return_tensors="pt")


the CoreML model only works with that input when tested on macOS. How can I make the CoreML model work for any input (i.e., any text)?


Export script:

# -*- coding: utf-8 -*-
"""Core ML Export
pip install transformers torch coremltools nltk
"""
import os
from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
import torch.nn as nn
import nltk
import coremltools as ct

nltk.download('punkt')

# Load the model and tokenizer
model_path = os.path.join('model')
model = AutoModelForTokenClassification.from_pretrained(model_path, local_files_only=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)

# Modify the model's forward method to return a tuple
class ModifiedModel(nn.Module):
def __init__(self, model):
super(ModifiedModel, self).__init__()
self.model = model
self.device = model.device # Add the device attribute

def forward(self, input_ids, attention_mask, token_type_ids=None):
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
return outputs.logits


modified_model = ModifiedModel(model)

# Export to Core ML
def convert_to_coreml(model, tokenizer):
# Define a dummy input for tracing
dummy_input = tokenizer("A French fan", return_tensors="pt")
dummy_input = {k: v.to(model.device) for k, v in dummy_input.items()}

# Trace the model with the dummy input
traced_model = torch.jit.trace(model, (
dummy_input['input_ids'], dummy_input['attention_mask'], dummy_input.get('token_type_ids')))

# Convert to Core ML
inputs = [
ct.TensorType(name="input_ids", shape=dummy_input['input_ids'].shape),
ct.TensorType(name="attention_mask", shape=dummy_input['attention_mask'].shape)
]
if 'token_type_ids' in dummy_input:
inputs.append(ct.TensorType(name="token_type_ids", shape=dummy_input['token_type_ids'].shape))

mlmodel = ct.convert(traced_model, inputs=inputs)

# Save the Core ML model
mlmodel.save("model.mlmodel")
print("Model exported to Core ML successfully")

convert_to_coreml(modified_model, tokenizer)


To use the exported model:

import os
from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
import torch.nn as nn
import nltk
import coremltools as ct
from coremltools.models import MLModel
import numpy as np
from transformers import AutoTokenizer
import nltk

nltk.download('punkt')

# Load the Core ML model
model = MLModel('model.mlmodel')

# Load the tokenizer
model_path = 'model'
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)

def prepare_input(text, tokenizer):
tokens = nltk.tokenize.word_tokenize(text)
tokenized_inputs = tokenizer(tokens, is_split_into_words=True, return_tensors="np")
input_ids = tokenized_inputs['input_ids'].astype(np.int32)
attention_mask = tokenized_inputs['attention_mask'].astype(np.int32)

input_data = {
'input_ids': input_ids,
'attention_mask': attention_mask
}

if 'token_type_ids' in tokenized_inputs:
input_data['token_type_ids'] = tokenized_inputs['token_type_ids'].astype(np.int32)

return input_data, tokens

def predict(text):
# Prepare the input
input_data, tokens = prepare_input(text, tokenizer)

# Make the prediction
prediction = model.predict(input_data)

# Extract the predicted labels
logits = prediction['output'] # Adjust this key according to your model's output
predicted_label = np.argmax(logits, axis=-1)[0]

# Display the results
for word, label in zip(tokens, predicted_label):
print(f"{word}: {model.model_description.outputDescriptions[0].dictionaryType.int64KeyType.stringDictionary[label]}")

# Test the model with a sentence
predict("A French fan")


The script only works with the example 'A French Fan'. When I tried another example predict("A footbal fan is standing in the stadium."), it triggers an error:

NSLocalizedDescription = "MultiArray shape (1 x 12) does not match the shape (1 x 5) specified in the model description";


Environments:

• Export script: Tested Python 3.10 and torch 2.3.1 on Ubuntu 20.04 (does not work on Windows 10).
• Prediction script: must be run on macOS 10.13+, as CoreML model only supports predictions on macOS 10.13+.

Continue reading...