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

File	What it is
embeddinggemma-300m.mlpackage/	Core ML model (ML Program, float16 weights and compute)
tokenizer.json, tokenizer_config.json, special_tokens_map.json	Tokenizer files from the source model, unchanged
config.json	Transformer config from the source model, unchanged
LICENSE	Gemma Terms of Use + conversion notice

Core ML interface

Name	Direction	dtype	Shape
input_ids	input	int32	(1, 64) or (1, 256)
attention_mask	input	int32	(1, 64) or (1, 256)
embedding	output	float32	(1, 768)

• 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):

Prompt name	Prefix
query	task: search result | query:
document	title: none | text:
Retrieval-query	task: search result | query:
Retrieval-document	title: none | text:
Reranking	task: search result | query:
BitextMining	task: search result | query:
STS	task: sentence similarity | query:
PairClassification	task: sentence similarity | query:
Classification	task: classification | query:
MultilabelClassification	task: classification | query:
Clustering	task: clustering | query:
Summarization	task: summarization | query:
InstructionRetrieval	task: code retrieval | query:

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:

Text	Tokens	Shape	cos(PyTorch, Core ML)
Russian passage (document prefix)	153	(1, 256)	0.999963
Russian short query (query prefix)	13	(1, 64)	0.999966

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.