Auto-generate attention mask

ModernBERT-base-4bit.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a92e9ca48ff235d69963aae1243cc46012dfafac748c1a28e548c5a43cb032c
-size 278858
+oid sha256:9a9d1c57cce5d2cdb960bbeb3cfd97443a0752e0542ab958eb636c6a63a87774
+size 279063

ModernBERT-base-4bit.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "9BD93E54-BA9D-470F-B2B6-0C4E6EF68B69": {
+        "E15EBF2F-BF39-4526-8AB8-51E73AFEC7FD": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Specification",
             "name": "model.mlmodel",
             "path": "com.apple.CoreML/model.mlmodel"
         },
-        "9EE76E38-2157-40A9-B3E0-1974022135F2": {
+        "FD5E8367-55E6-413C-8437-713B00D64144": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Weights",
             "name": "weights",
             "path": "com.apple.CoreML/weights"
         }
     },
-    "rootModelIdentifier": "9BD93E54-BA9D-470F-B2B6-0C4E6EF68B69"
+    "rootModelIdentifier": "E15EBF2F-BF39-4526-8AB8-51E73AFEC7FD"
 }

ModernBERT-base.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:769355987eb544d7ce0e25260e07dcc407b42a7044aa08835e3d434019b6629d
-size 266877
+oid sha256:67cc80bf6060a2d08059fe98c081fc9ebc34b074252143ca4dfa03d3e4f277df
+size 267080

ModernBERT-base.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "5DE8F5DD-BC7D-4275-90A0-85F8F3B565E9": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Weights",
-            "name": "weights",
-            "path": "com.apple.CoreML/weights"
-        },
-        "6B85C700-26D7-47DD-9768-DA523F859464": {
+        "A373F851-2CE5-43A4-8511-ECFD46CE8A03": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Specification",
             "name": "model.mlmodel",
             "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "C8A72ADC-A8C6-48A6-8B70-5F91F4815174": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
         }
     },
-    "rootModelIdentifier": "6B85C700-26D7-47DD-9768-DA523F859464"
+    "rootModelIdentifier": "A373F851-2CE5-43A4-8511-ECFD46CE8A03"
 }

ModernBERT-large-4bit.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc436ae7fd0bc2824bd12f0c0049db5d15234510d9ec73b78fa441d440b455ac
-size 350781
+oid sha256:dbe210fd61ab3ced4e7041c483472e46a3383d4e44849a6d5cdba2c4eea3614b
+size 350974

ModernBERT-large-4bit.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "69217F75-A051-4A83-B7B8-3F3A04468AD6": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Specification",
-            "name": "model.mlmodel",
-            "path": "com.apple.CoreML/model.mlmodel"
-        },
-        "D932257A-B814-45A2-B0A2-9CD3B727AF75": {
+        "5C57D828-673E-436E-ABD8-F8FAF5D3A45E": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Weights",
             "name": "weights",
             "path": "com.apple.CoreML/weights"
+        },
+        "BDC69D2B-C6B6-4E19-A9C2-B2333CA42723": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
         }
     },
-    "rootModelIdentifier": "69217F75-A051-4A83-B7B8-3F3A04468AD6"
+    "rootModelIdentifier": "BDC69D2B-C6B6-4E19-A9C2-B2333CA42723"
 }

ModernBERT-large.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ab69e9172a16cb1a6ed2efd5557c957446df8b6db7bc7cbfd83be23635ba71f6
-size 335647
+oid sha256:420ef83f761a8cc1c6baa280486ea89dd2892bce55a32d471af9a28a27b28c38
+size 335838

ModernBERT-large.mlpackage/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "D1999926-212D-4488-9717-A26B486AE1EE": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Specification",
-            "name": "model.mlmodel",
-            "path": "com.apple.CoreML/model.mlmodel"
-        },
-        "E3431F6D-5367-47DA-8179-3E46922823A0": {
+        "20D16BF3-9CD7-4631-8373-8CA41044B1AE": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Weights",
             "name": "weights",
             "path": "com.apple.CoreML/weights"
+        },
+        "652E02D8-970C-42B2-806B-A5B2911B653D": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
         }
     },
-    "rootModelIdentifier": "D1999926-212D-4488-9717-A26B486AE1EE"
+    "rootModelIdentifier": "652E02D8-970C-42B2-806B-A5B2911B653D"
 }

README.md

@@ -29,24 +29,21 @@ 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 attentionMask = MLShapedArray<Int32>(scalars: tokens.map { _ in 1 }, shape: [1, tokens.count])
-let input = ModernBERT_baseInput(
-    input_ids: inputIDs,
-    attention_mask: attentionMask
-)
+let input = ModernBERT_baseInput(input_ids: inputIDs)
 
 print("Predicting…")
 let output = try await model.prediction(input: input)
 let logits = output.logitsShapedArray
-let maskPosition = tokens.firstIndex(
-    of: tokenizer.encode(text: "[MASK]", addSpecialTokens: false)[0]
-)!
-let maskLogits = logits[0, maskPosition]
-let maskToken = Int(await MLTensor(maskLogits).argmax().shapedArray(of: Int32.self).scalar!)
-let maskText = tokenizer.decode(tokens: [maskToken])
 
+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.
-print(text.replacingOccurrences(of: "[MASK]", with: maskText.trimmingCharacters(in: .whitespaces)))
 ```