shirochenkov90/embeddinggemma-300m-coreml

EmbeddingGemma-300m — Core ML

Core ML conversion of google/embeddinggemma-300m for native macOS / iOS apps (Swift + Core ML). Weights are unchanged; only the format is.

The package contains the complete sentence-transformers pipeline, so the output is the final embedding — nothing to post-process on the Swift side:

Gemma3 transformer → mean pooling over attention_mask → Dense(768→3072) → Dense(3072→768) → L2 normalization

Files

Core ML interface

• —Enumerated shapes: pass exactly (1, 64) or (1, 256); both inputs must have the same shape in one call. Use 64 for short queries (faster), 256 for passages. Inputs longer than 256 tokens must be truncated.
• —Padding goes on the RIGHT: real tokens first, then pad id 0 with attention_mask = 0. Positions are computed as 0…S-1 inside the model, so left padding would change the result.
• —Batch size is 1.
• —The tokenizer prepends `<bos>` (id 2) and appends `<eos>` (id 1). Reproduce exactly that when tokenizing in Swift. Example: проверка → [2, 7877, 144813, 1].
• —Output is already L2-normalized: cosine similarity = dot product.
• —Embedding dimension: 768. Matryoshka truncation to 512/256/128 works as in the original model: take the first N values and re-normalize.
• —Minimum deployment target: macOS 15 / iOS 18 (required for enumerated shapes on two inputs). Converted with compute_units = ALL.
• —On the Neural Engine the runtime may log E5RT … tensor_buffer has known strides while the model has FlexibleShapeInfo — a warning caused by the enumerated shapes; predictions are still correct (see verification below).

Task prefixes (prompts)

Prepend the prefix to the raw text, exactly as in the original sentence-transformers config (note the trailing space):

For search:

query    = "task: search result | query: " + text
document = "title: none | text: " + text        // or "title: <title> | text: " + text

Swift usage sketch

import CoreML

let config = MLModelConfiguration()
config.computeUnits = .all
let model = try embeddinggemma_300m(configuration: config)   // class generated by Xcode from the .mlpackage

/// tokenIds must already contain <bos> (2) at the start and <eos> (1) at the end.
func embed(tokenIds: [Int32]) throws -> [Float] {
    let length = tokenIds.count <= 64 ? 64 : 256
    precondition(tokenIds.count <= length)
    let ids = try MLMultiArray(shape: [1, NSNumber(value: length)], dataType: .int32)
    let mask = try MLMultiArray(shape: [1, NSNumber(value: length)], dataType: .int32)
    for i in 0..<length {
        ids[i] = NSNumber(value: i < tokenIds.count ? tokenIds[i] : 0)   // pad id 0
        mask[i] = NSNumber(value: i < tokenIds.count ? 1 : 0)
    }
    let out = try model.prediction(input_ids: ids, attention_mask: mask)
    let e = out.embedding
    return (0..<e.count).map { Float(truncating: e[$0]) }
}

Verification (this exact package, Apple Silicon, compute units ALL)

Cosine between the PyTorch pipeline (SentenceTransformer.encode, float32) and Core ML:

Semantic check (Core ML): query task: search result | query: про деньги vs document A title: none | text: обсудили бюджет на следующий квартал → cos 0.2469; vs document B title: none | text: починили баг в плеере → cos 0.1923.

Cross-script check: скан vs Scan → cos 0.7522 in Core ML (0.7524 in PyTorch).

How it was converted

torch.jit.trace of a single nn.Module wrapping the whole sentence-transformers pipeline (transformer with an explicit bidirectional 4-D attention mask and position ids, mean pooling, both Dense layers, L2 norm), then coremltools.convert(..., convert_to="mlprogram", compute_precision=FLOAT16, minimum_deployment_target=macOS15, inputs with EnumeratedShapes [(1, 64), (1, 256)]). Before conversion the wrapper was checked against SentenceTransformer.encode (cos = 1.0000).

Versions: torch 2.7.0, transformers 5.17.0, sentence-transformers 6.0.1, coremltools 9.0, numpy 2.3 (numpy ≥ 2.4 breaks coremltools 9.0 — apple/coremltools#2633).

License

Gemma Terms of Use — see LICENSE. This is a format conversion of Google's model with unchanged weights; use is subject to the Gemma Prohibited Use Policy.