mlx-swift-lm/Tests/MLXLMIntegrationTests/IntegrationTestModels.swift at a0bdee5e9ec52d8160ea4c0ee4e08cd7790e3caf · ml-explore/mlx-swift-lm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
// Copyright © 2025 Apple Inc.

import Foundation
import MLXLLM
import MLXLMCommon
import MLXVLM

enum IntegrationTestModelIDs {
    static let llmModelId = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
    static let vlmModelId = "mlx-community/Qwen3-VL-4B-Instruct-4bit"

    static let lfm2ModelId = "mlx-community/LFM2-2.6B-Exp-4bit"
    static let glm4ModelId = "mlx-community/GLM-4-9B-0414-4bit"
    static let mistral3ModelId = "mlx-community/Ministral-3-3B-Instruct-2512-4bit"
    static let nemotronModelId = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit"
    static let qwen35ModelId = "mlx-community/Qwen3.5-2B-4bit"
}

actor IntegrationTestModels {
    static let shared = IntegrationTestModels()

    private init() {}

    private var llmTask: Task<ModelContainer, Error>?
    private var vlmTask: Task<ModelContainer, Error>?

    private var lfm2Task: Task<ModelContainer, Error>?
    private var glm4Task: Task<ModelContainer, Error>?
    private var mistral3Task: Task<ModelContainer, Error>?
    private var nemotronTask: Task<ModelContainer, Error>?
    private var qwen35Task: Task<ModelContainer, Error>?

    func llmContainer() async throws -> ModelContainer {
        if let task = llmTask {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.llmModelId)
            )
        }
        llmTask = task
        return try await task.value
    }

    func vlmContainer() async throws -> ModelContainer {
        if let task = vlmTask {
            return try await task.value
        }

        let task = Task {
            try await VLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.vlmModelId)
            )
        }
        vlmTask = task
        return try await task.value
    }

    func lfm2Container() async throws -> ModelContainer {
        if let task = lfm2Task {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.lfm2ModelId)
            )
        }
        lfm2Task = task
        return try await task.value
    }

    func glm4Container() async throws -> ModelContainer {
        if let task = glm4Task {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.glm4ModelId)
            )
        }
        glm4Task = task
        return try await task.value
    }

    func mistral3Container() async throws -> ModelContainer {
        if let task = mistral3Task {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.mistral3ModelId)
            )
        }
        mistral3Task = task
        return try await task.value
    }

    func nemotronContainer() async throws -> ModelContainer {
        if let task = nemotronTask {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.nemotronModelId)
            )
        }
        nemotronTask = task
        return try await task.value
    }

    func qwen35Container() async throws -> ModelContainer {
        if let task = qwen35Task {
            return try await task.value
        }

        let task = Task {
            try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: IntegrationTestModelIDs.qwen35ModelId)
            )
        }
        qwen35Task = task
        return try await task.value
    }
}