MLX Cleanup #187

jkrukowski · 2024-07-21T19:25:30Z

In this PR:

added new params to WhisperKit:
- mlxModel if specified, it will be used to specify which mlx model to download
- mlxModelRepo if specified, it will be used to specify the repo where the mlx model is stored
- mlxModelFolder if specified, it will point to the local folder where the mlx model is stored
added new options to CLIArguments:
- mlxModel, mlxModelPath see the explanation above
- featureExtractorType, audioEncoderType, textDecoderType to specify which type (coreml or or mlx) should be used for given pipeline step
updated readme to reflect the changes above
updated makefile to reflect the changes above
simplified CI unit test

To do:

update MLX to the latest version, right now blocked by .strides property is misleading -- PR #96 changes affected our implementation ml-explore/mlx-swift#117, edit: fixed
update the way we build brew, we should use xcodebuild instead of swift build

- updated readme - updated makefile

jkrukowski · 2024-07-21T19:26:02Z

.github/workflows/unit-tests.yml

-          xcodebuild clean build-for-testing -scheme ${{ matrix.run-config['scheme'] }} -destination '${{ matrix.run-config['clean-destination'] }}' -skipPackagePluginValidation | xcpretty
-          xcodebuild test -only-testing WhisperKitTests/UnitTests -scheme ${{ matrix.run-config['scheme'] }} -destination '${{ matrix.run-config['test-destination'] }}' -skipPackagePluginValidation
+          xcodebuild clean build-for-testing test \
+            ${{ matrix.run-config['test-cases'] }} \
+            -scheme ${{ matrix.run-config['scheme'] }} \
+            -destination '${{ matrix.run-config['test-destination'] }}' \
+            -skipPackagePluginValidation


this is merged into single command xcodebuild clean build-for-testing test

jkrukowski · 2024-07-21T19:26:33Z

Makefile

+	@xcodebuild CLANG_ENABLE_CODE_COVERAGE=NO VALID_ARCHS=arm64 clean build \
+		-configuration Release \
+		-scheme whisperkit-Package \
+		-destination generic/platform=macOS \
+		-derivedDataPath .build/.xcodebuild/ \
+		-clonedSourcePackagesDirPath .build/ \
+		-skipPackagePluginValidation


we have to use xcodebuild because of mlx

jkrukowski · 2024-07-21T19:26:53Z

Makefile

-	@swift build -c release --product whisperkit-cli
+	@xcodebuild CLANG_ENABLE_CODE_COVERAGE=NO VALID_ARCHS=arm64 clean build \
+		-configuration Release \
+		-scheme whisperkit-cli \
+		-destination generic/platform=macOS \
+		-derivedDataPath .build/.xcodebuild/ \
+		-clonedSourcePackagesDirPath .build/ \
+		-skipPackagePluginValidation


 test:
 	@echo "Running tests..."
-	@swift test -v
+	@xcodebuild clean build-for-testing test \
+		-scheme whisperkit-Package \
+		-only-testing WhisperKitMLXTests/MLXUnitTests \
+		-only-testing WhisperKitTests/UnitTests \
+		-destination 'platform=macOS,arch=arm64' \
+		-skipPackagePluginValidation


we have to use xcodebuild because of mlx

jkrukowski · 2024-07-21T19:28:10Z

Package.swift

-            .package(url: "https://github.com/ml-explore/mlx-swift", branch: "main"),
+            .package(url: "https://github.com/ml-explore/mlx-swift", revision: "d6d9472da5bf7ec2654e8914bd1d15622f45b6a9"),


the latest version of mlx 0.16.0 breaks the implementation, let's pin it to the last commit before the breaking one, I'll have to talk to mlx team about the issue

fixed, solution provided here ml-explore/mlx-swift#117

jkrukowski · 2024-07-21T19:29:32Z

README.md

-Task {
-   let pipe = try? await WhisperKit()
-   let transcription = try? await pipe!.transcribe(audioPath: "path/to/your/audio.{wav,mp3,m4a,flac}")?.text
-    print(transcription)
-}
+let pipe = try await WhisperKit()
+// Transcribe the audio file
+let transcription = try await pipe.transcribe(audioPath: "path/to/your/audio.{wav,mp3,m4a,flac}")?.text
+// Print the transcription
+print(transcription)


removed not needed Task, decided to remove ? from try ?. When the potential user does copy paste from the readme I think it's better when the error is not ignored

ZachNagengast

Looking great, in the final stages of completing the MLX support project 🎉. Left a few comments and questions here to consider.

Makefile

ZachNagengast · 2024-08-03T09:18:33Z

README.md

 ```

 ### Model Selection

-WhisperKit automatically downloads the recommended model for the device if not specified. You can also select a specific model by passing in the model name:
+You have to specify the model by passing the model name:


Why this change?

Right now user has to pass either model and / or mlxModel, depending whether mlx is used or not

README.md

ZachNagengast · 2024-08-03T09:23:55Z

Sources/WhisperKit/Core/WhisperKit.swift

    public func prewarmModels() async throws {
        try await loadModels(prewarmMode: true)
    }

    public func loadModels(
        prewarmMode: Bool = false
    ) async throws {
-        modelState = prewarmMode ? .prewarming : .loading
+        assert(modelFolder != nil || mlxModelFolder != nil, "Please specify `modelFolder` or `mlxModelFolder`")


Will this crash on prod builds? If so it should throw instead of crash.

it's not going to crash on a prod build

ZachNagengast · 2024-08-03T09:25:46Z

Sources/WhisperKit/Core/WhisperKit.swift

@@ -211,114 +228,80 @@ open class WhisperKit {
        }
    }

-    /// Sets up the model folder either from a local path or by downloading from a repository.
-    public func setupModels(


Is this deleted for a specific reason? Would be preferred to keep with with some adjustments since it was previously a public function and may be getting called elsewhere.

previously setupModels would set modelFolder property. Right now we have to set either modelFolder or / and mlxModelFolder depending on the pipeline setup. In this context this function doesn't make sense any more

README.md

ZachNagengast

Approved pending the setupModels change (also to run the tests)

ZachNagengast · 2024-08-10T22:41:42Z

Test passing ✅

jkrukowski added 2 commits July 18, 2024 09:57

update mlx-swift

2f3780f

- reverted mlx version

fd52493

- updated readme - updated makefile

jkrukowski commented Jul 21, 2024

View reviewed changes

jkrukowski added 6 commits July 21, 2024 21:42

reversed

c486bd8

fixed tests

799fca7

updated mlx-swift

c258a41

updated makefile

dd65efe

remove device change, fft can run on gpu now

60cf3be

updated readme, added tests

6e6d45d

jkrukowski marked this pull request as ready for review July 25, 2024 15:21

updated readme

65514fb

ZachNagengast requested changes Aug 3, 2024

View reviewed changes

jkrukowski added 4 commits August 10, 2024 16:20

review changes

15d9b80

review changes

c728996

CI model cache path fix

2e56036

tests failed

801e04a

jkrukowski requested a review from ZachNagengast August 10, 2024 15:04

ZachNagengast approved these changes Aug 10, 2024

View reviewed changes

ZachNagengast added 6 commits August 10, 2024 13:08

Update package.swift

f4cedc5

Keep setupModels with adjustments

a6b984e

Test CI skip cache

d0bf45f

Test CI package name change

95c2a1f

Test CI optional CLI settings

3808eb8

Use correct logits size for MLX

aa1bbfa

ZachNagengast merged commit 2ea8426 into argmaxinc:mlx-support Aug 10, 2024
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MLX Cleanup #187

MLX Cleanup #187

jkrukowski commented Jul 21, 2024 •

edited

Loading

jkrukowski Jul 21, 2024

jkrukowski Jul 21, 2024

jkrukowski Jul 21, 2024

jkrukowski Jul 21, 2024

jkrukowski Jul 25, 2024

jkrukowski Jul 21, 2024

ZachNagengast left a comment

ZachNagengast Aug 3, 2024

jkrukowski Aug 10, 2024 •

edited

Loading

ZachNagengast Aug 3, 2024

jkrukowski Aug 10, 2024

ZachNagengast Aug 3, 2024

jkrukowski Aug 10, 2024

ZachNagengast left a comment

ZachNagengast commented Aug 10, 2024

		.package(url: "https://github.com/ml-explore/mlx-swift", branch: "main"),
		.package(url: "https://github.com/ml-explore/mlx-swift", revision: "d6d9472da5bf7ec2654e8914bd1d15622f45b6a9"),

MLX Cleanup #187

MLX Cleanup #187

Conversation

jkrukowski commented Jul 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ZachNagengast left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkrukowski Aug 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ZachNagengast left a comment

Choose a reason for hiding this comment

ZachNagengast commented Aug 10, 2024

jkrukowski commented Jul 21, 2024 •

edited

Loading

jkrukowski Aug 10, 2024 •

edited

Loading