ModernBERT-CoreML

This repo contains ModernBERT-base and ModernBERT-large converted to CoreML.

Example Usage

import CoreML
import Tokenizers

let text = "The capital of Ireland is [MASK]."

print("Loading…")
let model = try await ModernBERT_base.load()
let tokenizer = try await AutoTokenizer.from(pretrained: "answerdotai/ModernBERT-base")

print("Tokenizing…")
let tokens = tokenizer(text)
let inputIDs = MLShapedArray(scalars: tokens.map(Int32.init), shape: [1, tokens.count])
let input = ModernBERT_baseInput(input_ids: inputIDs)

print("Predicting…")
let output = try await model.prediction(input: input)
let logits = output.logitsShapedArray

print("Decoding…")
let maskPosition = tokens.firstIndex(of: tokenizer.convertTokenToId("[MASK]")!)!
let predictedTokenID = await MLTensor(logits[0, maskPosition]).argmax().shapedArray(of: Int32.self).scalar!
let predictedTokenText = tokenizer.decode(tokens: [Int(predictedTokenID)])

print("Result:")
print(text.replacingOccurrences(of: "[MASK]", with: predictedTokenText.trimmingCharacters(in: .whitespaces)))
// The capital of Ireland is Dublin.

Conversion

uv run https://hf.co/finnvoorhees/ModernBERT-CoreML/raw/main/convert.py

usage: convert.py [-h] [--model MODEL] [--quantize]

Convert ModernBERT to CoreML

options:
  -h, --help     show this help message and exit
  --model MODEL  Model name
  --quantize     Linear quantize model