Building a Private Classifier with Apple's On-Device Model and NLContextualEmbedding

#ai #swift #macos #ios

Apple's on-device model is a small language model that ships with macOS and iOS. It runs offline, costs nothing per call, and never sends your data anywhere. This article covers what it is good at, where it falls short, and how to build a classification engine from two pieces: NLContextualEmbedding and the Foundation Models framework.

How good is it?

Size and context. The model has roughly 3 billion parameters. Its context window was about 4,096 tokens at launch, and that budget covers instructions, prompt and response together. Check the context size on your target OS version, because it can change between releases. A long document will not fit, so plan to chunk it.

Speed. It runs on the Neural Engine and GPU, and generation is fast enough to feel interactive on Apple silicon. Time to first token is short, so short classification calls come back quickly. Throughput varies a lot with the chip and with the length of the output, so measure on your own hardware before quoting numbers.

What it is good for:

Classifying and tagging text.
Extracting structured fields from short text.
Summarising or rewriting short passages.
Anything where privacy or offline use is a requirement.

What it is not good for:

Broad world knowledge.
Heavy reasoning or maths.
Long documents without chunking.

Structured output is the key feature

With @Generable you get typed Swift values back instead of free text. For classification this means the model can only answer with the shape you define:

import FoundationModels

@Generable
struct Verdict {
    @Guide(description: "Exactly one of the allowed category names")
    var category: String
    @Guide(description: "Confidence from 0 to 1")
    var confidence: Double
}

let session = LanguageModelSession(instructions: """
Classify the text into one of: Invoice, Contract, Resume, Other.
""")
let result = try await session.respond(to: text, generating: Verdict.self)

Add embeddings: a fast, cheap first pass

Asking the language model about every item is the slow path. NLContextualEmbedding is a BERT-style model built into the NaturalLanguage framework. It turns text into vectors, so you can score items by similarity to each category without any generation:

import NaturalLanguage

let model = NLContextualEmbedding(language: .english)!
try await model.requestAssets()   // one-time asset download
try model.load()

func embed(_ text: String) throws -> [Float] {
    let r = try model.embeddingResult(for: text, language: .english)
    var sum = [Float](repeating: 0, count: model.dimension)
    var n: Float = 0
    r.enumerateTokenVectors(in: text.startIndex..<text.endIndex) { v, _ in
        for i in 0..<sum.count { sum[i] += Float(v[i]) }
        n += 1
        return true
    }
    return sum.map { $0 / max(n, 1) }   // mean pooling
}

Mean-pool the token vectors, then compare with cosine similarity.

The classification engine

Define categories with a name and a short description. Add 3 to 5 example phrases each if you can.
Embed each category once and cache the vectors.
Embed the input and compute cosine similarity against every category.
If the best score is clearly ahead of the second, accept it. That path is fast and uses no generation.
If the top two are close, ask the language model to choose between just those two candidates, using @Generable. Narrowing to two keeps the prompt short and the answer reliable.

func classify(_ text: String) async throws -> String {
    let v = try embed(text)
    let ranked = categories
        .map { ($0.name, cosine(v, $0.vector)) }
        .sorted { $0.1 > $1.1 }
    if ranked[0].1 - ranked[1].1 > 0.08 { return ranked[0].0 }
    let pair = "\(ranked[0].0) or \(ranked[1].0)"
    let s = LanguageModelSession(instructions: "Choose between: \(pair).")
    return try await s.respond(to: text, generating: Verdict.self).content.category
}

The 0.08 margin is a starting point. Tune it against a labelled set of your own, and keep that set as a regression test.

Practical tips

Check SystemLanguageModel.default.availability first. The device may not support Apple Intelligence, or it may be turned off.
Chunk long inputs and classify each chunk, or classify the first chunk and a sample from the rest.
Keep instructions short, because they count against the context window.
Build a golden set early. It is the only way to know whether a prompt change helped.

Wrap-up

Embeddings give you speed and a stable baseline. The language model gives you judgement on the hard cases. Used together they make a private, offline classifier with no API keys and no per-call cost.

DEV Community