Upgrade to mlx-swift-lm 3.x with Hugging Face integration - #1
Upgrade to mlx-swift-lm 3.x with Hugging Face integration#1ronaldmannak wants to merge 3 commits into
Conversation
Bump mlx-swift-lm from 2.30.3 to 3.31.4 and migrate ChromaEmbedder to the
3.x MLXEmbedders API. The 2->3 rewrite decoupled the model downloader and
tokenizer out of MLXEmbedders, so the embedder now loads via
EmbedderModelFactory with the Hugging Face integration packages wired in at
the call site.
Package.swift:
- mlx-swift-lm: from 2.30.3 -> from 3.31.4
- Add swift-huggingface (from 0.9.0) and swift-transformers (from 1.3.0),
the integration packages the 3.x embedders require (versions per the
mlx-swift-lm 3.31.4 docs)
- Chroma target now pulls MLXEmbedders, MLXLMCommon, MLXHuggingFace,
HuggingFace, and Tokenizers
ChromaEmbedder.swift:
- Replace `import Hub` with MLXLMCommon / MLXHuggingFace / HuggingFace /
Tokenizers
- Model configs move from ModelConfiguration.* to EmbedderRegistry.*
- ModelContainer -> EmbedderModelContainer; drop HubApi
- loadModelContainer(hub:configuration:) ->
EmbedderModelFactory.shared.loadContainer(from:using:configuration:) using
the #hubDownloader() / #huggingFaceTokenizerLoader() macros
- Encode closures use the current `perform { context in }` form (the 3-arg
overload is deprecated in 3.x)
Package.resolved is removed because it pinned the old graph and cannot be
regenerated in this Linux environment (no Swift toolchain / MLX needs Metal);
SwiftPM regenerates it on the first macOS resolve/build.
Docs (README, api-notes) updated to describe the 3.x integration-package
model; scripts/update_dependencies.swift annotated as stale.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019ZY2q1wQyxRpAQGSRrAtjg
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the mlx-swift-lm dependency to version 3.x, integrating swift-huggingface and swift-transformers for model downloading and tokenization. It refactors ChromaEmbedder to utilize the new EmbedderModelContainer and EmbedderModelFactory APIs. The review feedback highlights opportunities to prevent redundant model loading, ensure Swift 6 concurrency safety by converting ChromaEmbedder to an actor, and optimize batch embedding performance by reducing GPU-to-CPU memory transfers.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| public func loadModel() async throws { | ||
| do { |
There was a problem hiding this comment.
Redundant Model Loading
Calling loadModel() multiple times will trigger redundant downloads and model initialization. Since loading an embedding model is an expensive operation, we should guard against reloading if the model is already initialized.
| public func loadModel() async throws { | |
| do { | |
| public func loadModel() async throws { | |
| guard !isInitialized else { return } | |
| do { |
There was a problem hiding this comment.
Fixed in 0b4e214. Added guard !isInitialized else { return } at the top of loadModel() so repeat calls are no-ops and don't re-download/re-initialize the model. Thanks!
Generated by Claude Code
| // ModelContainer actor manages the model loading and inference | ||
| private var modelContainer: ModelContainer? | ||
| // EmbedderModelContainer actor manages the model loading and inference | ||
| private var modelContainer: EmbedderModelContainer? |
There was a problem hiding this comment.
Thread Safety and Sendability under Swift 6
ChromaEmbedder is currently a standard class with mutable state (modelContainer and isInitialized) and async methods. Under Swift 6's strict concurrency rules (which are enabled in Package.swift via .swiftLanguageMode(.v6)), this class is not thread-safe and does not conform to Sendable.
To prevent data races and ensure full Swift 6 compatibility, consider converting ChromaEmbedder into an actor or marking it as @MainActor. Since all of its public methods are already async, converting it to an actor would be a highly compatible and safe change.
There was a problem hiding this comment.
Deferring this one for now, not because it's wrong but because it doesn't belong in this PR. Reasoning:
- Not a regression from this change.
ChromaEmbedderwas already aclasswith the same mutable state under.swiftLanguageMode(.v6)before this PR; the dependency bump didn't introduce a new concurrency error. The compiler doesn't forceSendable/actor isolation here because instances aren't sent across isolation boundaries (thecontainer.perform { context in … }closures don't captureself). - It's an API-affecting, architectural change. Converting to
actormakes the synchronous extension API actor-isolated — e.g.createCollection(name:) throwsinChromaEmbedderExtensions.swiftand synchronous property reads — so existing call sites would need newawaits. That's a source-compatibility change I'd rather not bundle into a dependency upgrade. - I can't compile-verify it here. This is a Metal/MLX package that only builds on macOS/Apple silicon; I don't want to land an unverifiable refactor of the concurrency model in this PR.
I think it's a reasonable follow-up on its own (actor + nonisolated on the passthrough members), just not mixed into the dep bump. Happy to open a separate PR for it if you'd like — @ronaldmannak your call.
Generated by Claude Code
| let modelOutput = model(input, positionIds: nil, tokenTypeIds: tokenTypes, attentionMask: attentionMask) | ||
| let pooled = pooling(modelOutput, mask: attentionMask, normalize: true) | ||
| let modelOutput = context.model(input, positionIds: nil, tokenTypeIds: tokenTypes, attentionMask: attentionMask) | ||
| let pooled = context.pooling(modelOutput, mask: attentionMask, normalize: true) |
There was a problem hiding this comment.
Batch Embedding Performance Optimization
In the subsequent loop (lines 181-183), calling pooled[i].asArray(Float.self) triggers texts.count separate synchronous GPU-to-CPU memory transfers. This introduces significant overhead, especially for larger batch sizes.
To optimize performance, you can call pooled.asArray(Float.self) once to copy the entire batch to the CPU as a flat array, and then slice it in Swift. This reduces the transfer overhead to a single operation.
Here is how you can implement this optimization:
let flatArray = pooled.asArray(Float.self)
let dim = pooled.shape[1]
var embeddings = [[Float]]()
embeddings.reserveCapacity(texts.count)
for i in 0..<texts.count {
let start = i * dim
let end = start + dim
embeddings.append(Array(flatArray[start..<end]))
}
return embeddingsThere was a problem hiding this comment.
Fixed in 0b4e214. encodeBatch now copies the whole pooled batch to the CPU with a single pooled.asArray(Float.self) and slices per row using dim = pooled.shape[1], replacing the previous per-row pooled[i].asArray(...) (which did one GPU→CPU transfer per embedding).
Generated by Claude Code
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e9c75fb0cf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let modelOutput = model(input, positionIds: nil, tokenTypeIds: tokenTypes, attentionMask: attentionMask) | ||
| let pooled = pooling(modelOutput, mask: attentionMask, normalize: true) | ||
| let modelOutput = context.model(input, positionIds: nil, tokenTypeIds: tokenTypes, attentionMask: attentionMask) | ||
| let pooled = context.pooling(modelOutput, mask: attentionMask, normalize: true) |
There was a problem hiding this comment.
With the 3.31.4 upgrade, MLXEmbedders.BertModel.callAsFunction truncates inputs and the model-side attention mask when a sequence exceeds maxPositionEmbeddings; for a batch containing a long document, modelOutput is shortened while the attentionMask passed here still has the original maxLen. Mean/last/max pooling for common SentenceTransformer-style embedders then sees incompatible sequence lengths (or out-of-range last-token indices), so embed(texts:) can fail for >512-token inputs. Rebuild or slice the pooling mask to match the returned hidden-state length before passing it to context.pooling.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in 0b4e214. Confirmed against the 3.31.4 source: BertModel.callAsFunction (Bert.swift) truncates inp/mask/typeIds/posIds to maxPositionEmbeddings and returns hidden states of that shortened length, while encodeBatch was still handing the original full-length attentionMask to context.pooling — a shape mismatch (mean/max) or out-of-range last-token index for >context-window batches. encodeBatch now slices the pooling mask to the returned hidden-state length:
let seqLen = modelOutput.hiddenStates?.dim(1) ?? maxLen
let pooledMask = attentionMask[0..., ..<seqLen]
let pooled = context.pooling(modelOutput, mask: pooledMask, normalize: true)Since truncation keeps the leading tokens, slicing the mask to the first seqLen columns stays aligned. The single-text path passes no mask (pooling defaults to ones matching the hidden states), so it was already safe.
Generated by Claude Code
| import MLXLMCommon | ||
| import MLXHuggingFace | ||
| import HuggingFace | ||
| import Tokenizers |
There was a problem hiding this comment.
When developers follow the documented use_local_framework.sh / use_release_framework.sh workflow, those scripts still rewrite Package.swift to depend only on mlx-swift-examples / MLXEmbedders, so these newly imported 3.x integration modules are no longer linked and the package fails to compile after switching frameworks. Update the generated manifests in those scripts alongside this dependency upgrade or stop documenting that workflow.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 0b4e214. Both scripts/use_local_framework.sh and scripts/use_release_framework.sh now regenerate a Package.swift that matches the new graph: mlx-swift-lm from: "3.31.4" plus swift-huggingface and swift-transformers, all five products (MLXEmbedders, MLXLMCommon, MLXHuggingFace, HuggingFace, Tokenizers), swift-tools-version: 6.2, iOS 17, Swift 6 language mode, and the test target — differing only in the FFI binaryTarget (local path vs. remote url/checksum). They previously emitted mlx-swift-examples exact 2.25.6 with only MLXEmbedders, which would have failed to compile the 3.x imports after a framework switch (and were already stale vs. the committed manifest).
Generated by Claude Code
- Guard loadModel() with `guard !isInitialized else { return }` to avoid
redundant model downloads/initialization on repeat calls (Gemini).
- encodeBatch: slice the pooling mask to the model's returned hidden-state
length. Bert-style models truncate inputs to maxPositionEmbeddings, so for
>context-window batches the full-length mask no longer matched the hidden
states and pooling would fail (Codex).
- encodeBatch: copy the pooled batch to the CPU in a single asArray transfer
and slice per row, instead of one GPU->CPU transfer per embedding (Gemini).
- use_local_framework.sh / use_release_framework.sh: regenerate the current
dependency graph (mlx-swift-lm 3.31.4 + swift-huggingface + swift-transformers,
five products, tools 6.2, iOS 17, Swift 6 language mode, test target) instead
of the old mlx-swift-examples / MLXEmbedders-only manifest, so the documented
framework-switch workflow still compiles (Codex).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019ZY2q1wQyxRpAQGSRrAtjg
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9da7ee6485
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // #expect(info.name == name) | ||
| // #expect(info.collectionId == createdId) | ||
| // #expect(info.numDocuments == 0) | ||
| let info = try Chroma.getCollection(collectionName: name) |
There was a problem hiding this comment.
This PR still uses the same published 1.0.2 XCFramework/checksum in Package.swift and does not change the generated FFI wrapper, so the getCollection call that was documented here as hanging for both valid and invalid names is still the one CI will execute. Re-enabling it makes swift test hang in the default package configuration; keep these assertions skipped until the binary framework is updated.
Useful? React with 👍 / 👎.
2 participants
Summary
Upgrade ChromaEmbedder to use mlx-swift-lm 3.x with the new Hugging Face integration packages for model downloading and tokenization. This modernizes the embedding pipeline to use the latest factory-based API.
Key Changes
mlx-swift-lmfrom 2.30.3 to 3.31.4, addedswift-huggingface(0.9.0+) andswift-transformers(1.3.0+)Hubimport withMLXLMCommon,MLXHuggingFace,HuggingFace, andTokenizersMLXEmbedders.loadModelContainer(hub:configuration:)toEmbedderModelFactory.shared.loadContainer(from:using:configuration:)with#hubDownloader()and#huggingFaceTokenizerLoader()macrosModelContainerreferences toEmbedderModelContainerand refactored theperformclosure to use a context object withmodel,tokenizer, andpoolingproperties.bge_micro) toEmbedderRegistrystatic properties (e.g.,EmbedderRegistry.bge_micro)hubApiproperty in favor of the factory-based approachImplementation Details
The refactoring maintains backward compatibility at the public API level while modernizing the internal embedding pipeline. The factory pattern with macro-based dependency injection (
#hubDownloader(),#huggingFaceTokenizerLoader()) provides cleaner separation of concerns and better testability compared to the previous HubApi instance approach.https://claude.ai/code/session_019ZY2q1wQyxRpAQGSRrAtjg