60 lines
2.0 KiB
Swift
60 lines
2.0 KiB
Swift
// Copyright © 2024 Apple Inc.
|
|
|
|
import Foundation
|
|
import MLX
|
|
import MLXNN
|
|
|
|
/// Load model weights.
|
|
///
|
|
/// This is typically called via ``ModelFactory/load(from:configuration:progressHandler:)``.
|
|
/// This function loads all `safetensor` files in the given `modelDirectory`,
|
|
/// calls ``LanguageModel/sanitize(weights:metadata:)`` to allow per-model preprocessing,
|
|
/// applies optional quantization, and
|
|
/// updates the model with the weights.
|
|
public func loadWeights(
|
|
modelDirectory: URL, model: LanguageModel,
|
|
quantization: BaseConfiguration.Quantization? = nil,
|
|
perLayerQuantization: BaseConfiguration.PerLayerQuantization? = nil
|
|
) throws {
|
|
// load the weights and collect metadata from the first safetensor file
|
|
var weights = [String: MLXArray]()
|
|
var metadata = [String: String]()
|
|
let enumerator = FileManager.default.enumerator(
|
|
at: modelDirectory, includingPropertiesForKeys: nil)!
|
|
for case let url as URL in enumerator {
|
|
if url.pathExtension == "safetensors" {
|
|
let (w, m) = try loadArraysAndMetadata(url: url)
|
|
for (key, value) in w {
|
|
weights[key] = value
|
|
}
|
|
if metadata.isEmpty {
|
|
metadata = m
|
|
}
|
|
}
|
|
}
|
|
|
|
// per-model cleanup (models can inspect metadata to customize behavior)
|
|
weights = model.sanitize(weights: weights, metadata: metadata)
|
|
|
|
// quantize if needed
|
|
if quantization != nil || perLayerQuantization != nil {
|
|
quantize(model: model) { path, module in
|
|
if weights["\(path).scales"] != nil {
|
|
if let perLayerQuantization {
|
|
return perLayerQuantization.quantization(layer: path)?.asTuple
|
|
} else {
|
|
return quantization?.asTuple
|
|
}
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
|
|
// apply the loaded weights
|
|
let parameters = ModuleParameters.unflattened(weights)
|
|
try model.update(parameters: parameters, verify: [.all])
|
|
|
|
eval(model)
|
|
}
|