feat: migration to mlx-swift-lm v3
This commit is contained in:
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
)
|
||||
|
||||
@@ -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 }
|
||||
)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user