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 it 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)

Tested Python 3.10 and torch 2.3.1 with on Ubuntu 20.04 (does not work on Windows 10).