feat: migration to mlx-swift-lm v3

This commit is contained in:
2026-04-30 09:18:37 +02:00
parent 4ad46ec1ea
commit 3502266ff9
13 changed files with 211 additions and 228 deletions

View File

@@ -3,30 +3,27 @@ import XCTest
@testable import MLX_Server
final class LocalModelResolverTests: XCTestCase {
func testDiscoverModelsInfersTextOnlyMetadataAndDirectorySize() throws {
let base = try makeTempModelsRoot()
let repoDirectory = try makeRepoDirectory(base: base, owner: "example", repo: "text-only")
let configURL = repoDirectory.appendingPathComponent("config.json")
let modelURL = repoDirectory.appendingPathComponent("model.safetensors")
let tokenizerURL = repoDirectory.appendingPathComponent("tokenizer.json")
func testDiscoverSystemHFModelsInfersTextOnlyMetadata() throws {
let base = try makeTempHFCache()
let snapshotDir = try makeHFSnapshot(base: base, repoId: "example/text-only")
try writeJSON(
[
"architectures": ["LlamaForCausalLM"],
"max_position_embeddings": 32768,
],
to: configURL
to: snapshotDir.appendingPathComponent("config.json")
)
try Data(repeating: 0x11, count: 64).write(to: modelURL)
try Data(repeating: 0x22, count: 19).write(to: tokenizerURL)
try Data(repeating: 0x11, count: 64).write(to: snapshotDir.appendingPathComponent("model.safetensors"))
try Data(repeating: 0x22, count: 19).write(to: snapshotDir.appendingPathComponent("tokenizer.json"))
let expectedSize = Int64(
try Data(contentsOf: configURL).count
+ Data(contentsOf: modelURL).count
+ Data(contentsOf: tokenizerURL).count
try Data(contentsOf: snapshotDir.appendingPathComponent("config.json")).count
+ Data(contentsOf: snapshotDir.appendingPathComponent("model.safetensors")).count
+ Data(contentsOf: snapshotDir.appendingPathComponent("tokenizer.json")).count
)
let discovered = LocalModelResolver.discoverModels(in: base)
let discovered = LocalModelResolver.discoverSystemHFModels(in: base)
let model = try XCTUnwrap(discovered.first)
XCTAssertEqual(model.repoId, "example/text-only")
@@ -36,21 +33,25 @@ final class LocalModelResolverTests: XCTestCase {
XCTAssertEqual(model.sizeBytes, expectedSize)
}
func testDiscoverModelsInfersVisionMetadataFromProcessorFiles() throws {
let base = try makeTempModelsRoot()
let repoDirectory = try makeRepoDirectory(base: base, owner: "example", repo: "vision-model")
func testDiscoverSystemHFModelsInfersVisionMetadata() throws {
let base = try makeTempHFCache()
let snapshotDir = try makeHFSnapshot(base: base, repoId: "example/vision-model")
try writeJSON(
[
"text_config": ["max_position_embeddings": 262144],
"vision_config": ["hidden_size": 768],
],
to: repoDirectory.appendingPathComponent("config.json")
to: snapshotDir.appendingPathComponent("config.json")
)
try writeJSON(["processor_class": "Qwen3VLProcessor"], to: repoDirectory.appendingPathComponent("tokenizer_config.json"))
try Data(repeating: 0x33, count: 12).write(to: repoDirectory.appendingPathComponent("processor_config.json"))
try Data(repeating: 0x44, count: 8).write(to: repoDirectory.appendingPathComponent("model.safetensors.index.json"))
try writeJSON(
["processor_class": "Qwen3VLProcessor"],
to: snapshotDir.appendingPathComponent("tokenizer_config.json")
)
try Data(repeating: 0x33, count: 12).write(to: snapshotDir.appendingPathComponent("processor_config.json"))
try Data(repeating: 0x44, count: 8).write(to: snapshotDir.appendingPathComponent("model.safetensors.index.json"))
let discovered = LocalModelResolver.discoverModels(in: base)
let discovered = LocalModelResolver.discoverSystemHFModels(in: base)
let model = try XCTUnwrap(discovered.first)
XCTAssertEqual(model.repoId, "example/vision-model")
@@ -155,7 +156,7 @@ final class LocalModelResolverTests: XCTestCase {
XCTAssertTrue(config.supportsTools)
}
private func makeTempModelsRoot() throws -> URL {
private func makeTempHFCache() throws -> URL {
let root = FileManager.default.temporaryDirectory
.appendingPathComponent(UUID().uuidString, isDirectory: true)
try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)
@@ -165,16 +166,18 @@ final class LocalModelResolverTests: XCTestCase {
return root
}
private func makeRepoDirectory(base: URL, owner: String, repo: String) throws -> URL {
let directory = base
.appendingPathComponent(owner, isDirectory: true)
.appendingPathComponent(repo, isDirectory: true)
try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
return directory
private func makeHFSnapshot(base: URL, repoId: String, hash: String = "abc123") throws -> URL {
let slug = repoId.replacingOccurrences(of: "/", with: "--")
let snapshotDir = base
.appendingPathComponent("models--\(slug)", isDirectory: true)
.appendingPathComponent("snapshots", isDirectory: true)
.appendingPathComponent(hash, isDirectory: true)
try FileManager.default.createDirectory(at: snapshotDir, withIntermediateDirectories: true)
return snapshotDir
}
private func writeJSON(_ object: Any, to url: URL) throws {
let data = try JSONSerialization.data(withJSONObject: object, options: [.prettyPrinted, .sortedKeys])
try data.write(to: url)
}
}
}

View File

@@ -1,7 +1,9 @@
import Foundation
import Hub
import HuggingFace
import MLXHuggingFace
import MLXLMCommon
import MLXVLM
import Tokenizers
import XCTest
@testable import MLX_Server
@@ -671,10 +673,9 @@ private actor LocalGemmaFixture {
}
let loadTask = Task<ModelContainer, Error> {
let cachesDir = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first
let hub = HubApi(downloadBase: cachesDir, cache: nil)
return try await VLMModelFactory.shared.loadContainer(
hub: hub,
from: #hubDownloader(HubClient.default),
using: #huggingFaceTokenizerLoader(),
configuration: ModelConfiguration(directory: localDir),
progressHandler: { _ in }
)

View File

@@ -1,8 +1,10 @@
import Foundation
import Hub
import HuggingFace
import MLX
import MLXHuggingFace
import MLXLMCommon
import MLXVLM
import Tokenizers
import XCTest
@testable import MLX_Server
@@ -230,10 +232,9 @@ private actor LocalGemmaFixture {
}
let loadTask = Task<ModelContainer, Error> {
let cachesDir = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first
let hub = HubApi(downloadBase: cachesDir, cache: nil)
return try await VLMModelFactory.shared.loadContainer(
hub: hub,
from: #hubDownloader(HubClient.default),
using: #huggingFaceTokenizerLoader(),
configuration: ModelConfiguration(directory: localDir),
progressHandler: { _ in }
)

View File

@@ -249,4 +249,11 @@ private final class NonStandardCache: KVCache {
) -> MLXFast.ScaledDotProductAttentionMaskMode {
.none
}
func copy() -> any KVCache {
let c = NonStandardCache(tokenCount: 0, headDim: 0)
c.state = state
c.offset = offset
return c
}
}

View File

@@ -388,4 +388,10 @@ private final class TestTrimRecordingCache: KVCache {
) -> MLXFast.ScaledDotProductAttentionMaskMode {
.none
}
func copy() -> any KVCache {
let c = TestTrimRecordingCache(offset: offset, trimmable: trimmable)
c.state = state
return c
}
}