266 lines
8.4 KiB
Swift
266 lines
8.4 KiB
Swift
// Copyright © 2024 Apple Inc.
|
|
|
|
import Foundation
|
|
import MLX
|
|
import MLXNN
|
|
|
|
/// Time/Height/Width struct to represent information about input images.
|
|
public struct THW: Sendable {
|
|
|
|
public let t: Int
|
|
public let h: Int
|
|
public let w: Int
|
|
|
|
public init(_ t: Int, _ h: Int, _ w: Int) {
|
|
self.t = t
|
|
self.h = h
|
|
self.w = w
|
|
}
|
|
|
|
public var values: (Int, Int, Int) {
|
|
(t, h, w)
|
|
}
|
|
|
|
public var product: Int { t * h * w }
|
|
}
|
|
|
|
/// Representation of ``LanguageModel`` input.
|
|
///
|
|
/// This can contain text (tokens), prepared images (`MLXArray`), or other media as
|
|
/// needed. ``LMInput`` is produced by ``UserInputProcessor`` in response
|
|
/// to ``UserInput``.
|
|
///
|
|
/// The ``ModelContext`` holds the ``UserInputProcessor`` associated with a
|
|
/// ``LanguageModel``.
|
|
public struct LMInput {
|
|
public let text: Text
|
|
public let image: ProcessedImage?
|
|
public let video: ProcessedVideo?
|
|
public let audio: ProcessedAudio?
|
|
|
|
/// Representation of tokenized input text.
|
|
public struct Text {
|
|
|
|
/// input token array
|
|
public let tokens: MLXArray
|
|
|
|
/// optional mask array
|
|
public let mask: MLXArray?
|
|
|
|
public init(tokens: MLXArray, mask: MLXArray? = nil) {
|
|
self.tokens = tokens
|
|
self.mask = mask
|
|
}
|
|
|
|
public subscript(
|
|
indices: MLXArrayIndex..., stream stream: StreamOrDevice = .default
|
|
) -> Text {
|
|
Text(tokens: tokens[indices, stream: stream], mask: mask?[indices, stream: stream])
|
|
}
|
|
|
|
public subscript(
|
|
text indices: MLXArrayIndex..., stream stream: StreamOrDevice = .default
|
|
) -> Text {
|
|
Text(tokens: tokens[indices, stream: stream], mask: mask)
|
|
}
|
|
}
|
|
|
|
/// Representation of prepared input image(s).
|
|
public struct ProcessedImage {
|
|
|
|
/// Concatenated pixels from one or more images
|
|
public let pixels: MLXArray
|
|
/// Time, height, and width of the images
|
|
public let frames: [THW]?
|
|
/// Desired number of soft tokens after visual pooling/projecting.
|
|
public let softTokenCount: Int?
|
|
|
|
public init(
|
|
pixels: MLXArray, frames: [THW]? = nil, softTokenCount: Int? = nil
|
|
) {
|
|
self.pixels = pixels
|
|
self.frames = frames
|
|
self.softTokenCount = softTokenCount
|
|
}
|
|
}
|
|
|
|
/// Representation of prepared input video(s).
|
|
/// For now, this is virtually identical to ProcessedImage.
|
|
public struct ProcessedVideo {
|
|
|
|
public let pixels: MLXArray
|
|
public let frames: [THW]?
|
|
|
|
public init(
|
|
pixels: MLXArray, frames: [THW]? = nil
|
|
) {
|
|
self.pixels = pixels
|
|
self.frames = frames
|
|
}
|
|
}
|
|
|
|
/// Representation of prepared input audio.
|
|
public struct ProcessedAudio {
|
|
public let features: MLXArray
|
|
public let invalidMask: MLXArray?
|
|
public let sampleRate: Double
|
|
public let channelCount: Int
|
|
public let duration: TimeInterval
|
|
public let sampleCount: Int
|
|
public let tokenCount: Int
|
|
|
|
public init(
|
|
features: MLXArray,
|
|
invalidMask: MLXArray? = nil,
|
|
sampleRate: Double,
|
|
channelCount: Int,
|
|
duration: TimeInterval,
|
|
sampleCount: Int,
|
|
tokenCount: Int
|
|
) {
|
|
self.features = features
|
|
self.invalidMask = invalidMask
|
|
self.sampleRate = sampleRate
|
|
self.channelCount = channelCount
|
|
self.duration = duration
|
|
self.sampleCount = sampleCount
|
|
self.tokenCount = tokenCount
|
|
}
|
|
}
|
|
|
|
public init(tokens: MLXArray, mask: MLXArray? = nil) {
|
|
self.init(text: .init(tokens: tokens, mask: mask))
|
|
}
|
|
|
|
public init(
|
|
text: LMInput.Text, image: LMInput.ProcessedImage? = nil,
|
|
video: LMInput.ProcessedVideo? = nil,
|
|
audio: LMInput.ProcessedAudio? = nil
|
|
) {
|
|
self.text = text
|
|
self.image = image
|
|
self.video = video
|
|
self.audio = audio
|
|
}
|
|
}
|
|
|
|
/// ``LanguageModel`` step output. This is consumed internally
|
|
/// by the ``TokenIterator``.
|
|
public struct LMOutput {
|
|
|
|
/// logits (one hot vector of probabilities for tokens)
|
|
public let logits: MLXArray
|
|
|
|
/// optional ``State`` to carry forward into the next step
|
|
public let state: State?
|
|
|
|
public struct State {
|
|
public let crossAttentionStates: MLXArray?
|
|
|
|
public init(crossAttentionStates: MLXArray? = nil) {
|
|
self.crossAttentionStates = crossAttentionStates
|
|
}
|
|
}
|
|
|
|
public init(logits: MLXArray, state: LMOutput.State? = nil) {
|
|
self.logits = logits
|
|
self.state = state
|
|
}
|
|
}
|
|
|
|
/// The result of the call to ``LanguageModel/prepare(_:cache:windowSize:)``
|
|
public enum PrepareResult {
|
|
/// tokens to process by the ``TokenIterator``
|
|
case tokens(LMInput.Text)
|
|
|
|
/// logits representing the next token
|
|
case logits(LMOutput)
|
|
}
|
|
|
|
/// Interface for all Language Models (e.g. LLM, VLM).
|
|
///
|
|
/// The language model is typically called by the ``TokenIterator`` and it:
|
|
///
|
|
/// - consumes the ``LMInput``
|
|
/// - calls ``prepare(_:cache:windowSize:)`` to initialize the KVCache and consume the prompt
|
|
/// - calls ``callAsFunction(_:cache:state:)-9kuvf`` for each token, producing an ``LMOutput``
|
|
/// - the ``TokenIterator`` accumulates this information into a ``GenerateResult``
|
|
public protocol LanguageModel: Module {
|
|
|
|
/// Prepare the cache state and consume the ``LMInput``.
|
|
///
|
|
/// This can return:
|
|
/// - ``PrepareResult/tokens(_:)`` if the caller should evaluate the (remaining) tokens normally
|
|
/// - ``PrepareResult/logits(_:)`` to produce the next token from the prompt
|
|
func prepare(_ input: LMInput, cache: [KVCache], windowSize: Int?) throws -> PrepareResult
|
|
|
|
/// Primary entry point to produce a step (single token) from the model
|
|
func callAsFunction(_ input: LMInput.Text, cache: [KVCache]?, state: LMOutput.State?)
|
|
-> LMOutput
|
|
|
|
/// Models may implement this simplified interface if they do not produce any ``LMOutput/State``
|
|
func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray
|
|
|
|
/// create a new array of ``KVCache`` -- automatic implementation if self
|
|
/// implements ``KVCacheDimensionProvider``
|
|
func newCache(parameters: GenerateParameters?) -> [KVCache]
|
|
|
|
/// Optionally preprocess the weights and modify / remove values as needed.
|
|
func sanitize(weights: [String: MLXArray]) -> [String: MLXArray]
|
|
|
|
/// Optionally preprocess the weights with access to safetensor metadata.
|
|
///
|
|
/// The default implementation forwards to ``sanitize(weights:)``.
|
|
/// Models can override this to inspect metadata (e.g. check `metadata["format"] == "mlx"`)
|
|
/// and skip or customize sanitization accordingly.
|
|
func sanitize(weights: [String: MLXArray], metadata: [String: String]) -> [String: MLXArray]
|
|
}
|
|
|
|
extension LanguageModel {
|
|
public func callAsFunction(_ input: LMInput.Text, cache: [KVCache]?, state: LMOutput.State?)
|
|
-> LMOutput
|
|
{
|
|
let logits = callAsFunction(input.tokens, cache: cache)
|
|
return .init(logits: logits)
|
|
}
|
|
|
|
public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray {
|
|
fatalError("callAsFunction(inputs:cache:) not implemented for \(Self.self)")
|
|
}
|
|
}
|
|
|
|
extension LanguageModel {
|
|
public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
|
|
weights
|
|
}
|
|
|
|
public func sanitize(weights: [String: MLXArray], metadata: [String: String]) -> [String:
|
|
MLXArray]
|
|
{
|
|
sanitize(weights: weights)
|
|
}
|
|
}
|
|
|
|
/// Optional protocol that can be implemented by ``LanguageModel`` and will
|
|
/// provide an automatic implementation of ``LanguageModel/newCache(parameters:)``
|
|
public protocol KVCacheDimensionProvider {
|
|
var kvHeads: [Int] { get }
|
|
}
|
|
|
|
extension LanguageModel where Self: KVCacheDimensionProvider {
|
|
public func newCache(parameters: GenerateParameters?) -> [KVCache] {
|
|
// Create one cache per layer (kvHeads.count = number of layers)
|
|
// The number of heads per layer (kvHeads[i]) is not used for cache creation
|
|
let numLayers = kvHeads.count
|
|
|
|
// Follow Python logic: use RotatingKVCache if maxKVSize is provided
|
|
if let maxKVSize = parameters?.maxKVSize {
|
|
return (0 ..< numLayers).map { _ in
|
|
RotatingKVCache(maxSize: maxKVSize, keep: 4)
|
|
}
|
|
} else {
|
|
return (0 ..< numLayers).map { _ in KVCacheSimple() }
|
|
}
|
|
}
|
|
}
|