From ac4d9175b6d77f03acfebad5250801b2ffee863d Mon Sep 17 00:00:00 2001 From: shavit Date: Sat, 30 Mar 2024 15:09:04 -0400 Subject: [PATCH 1/6] Create embedding module --- Package.swift | 6 +- Sources/Embedding/Embedding.swift | 142 ++++++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 2 deletions(-) create mode 100644 Sources/Embedding/Embedding.swift diff --git a/Package.swift b/Package.swift index de98aa6..d0fdebd 100644 --- a/Package.swift +++ b/Package.swift @@ -7,7 +7,7 @@ let package = Package( name: "swift-transformers", platforms: [.iOS(.v16), .macOS(.v13)], products: [ - .library(name: "Transformers", targets: ["Tokenizers", "Generation", "Models"]), + .library(name: "Transformers", targets: ["Tokenizers", "Generation", "Models", "Embedding"]), .executable(name: "transformers", targets: ["TransformersCLI"]), .executable(name: "hub-cli", targets: ["HubCLI"]), ], @@ -26,11 +26,13 @@ let package = Package( .target(name: "TensorUtils"), .target(name: "Generation", dependencies: ["Tokenizers", "TensorUtils"]), .target(name: "Models", dependencies: ["Tokenizers", "Generation", "TensorUtils"]), + .target(name: "Embedding", dependencies: ["Hub", "Tokenizers"]), .testTarget(name: "TokenizersTests", dependencies: ["Tokenizers", "Models", "Hub"], resources: [.process("Resources"), .process("Vocabs")]), .testTarget(name: "HubTests", dependencies: ["Hub"]), .testTarget(name: "PreTokenizerTests", dependencies: ["Tokenizers", "Hub"]), .testTarget(name: "TensorUtilsTests", dependencies: ["TensorUtils"]), .testTarget(name: "NormalizerTests", dependencies: ["Tokenizers", "Hub"]), - .testTarget(name: "PostProcessorTests", dependencies: ["Tokenizers", "Hub"]) + .testTarget(name: "PostProcessorTests", dependencies: ["Tokenizers", "Hub"]), + .testTarget(name: "EmbeddingTests", dependencies: ["Embedding", "Tokenizers", "Hub"]) ] ) diff --git a/Sources/Embedding/Embedding.swift b/Sources/Embedding/Embedding.swift new file mode 100644 index 0000000..c46e1a3 --- /dev/null +++ b/Sources/Embedding/Embedding.swift @@ -0,0 +1,142 @@ +import Hub +import Tokenizers +import CoreML +import Accelerate + + +public protocol Embedding {} + +public struct AutoEmbedding {} // Otherwise AutoModel + +extension AutoEmbedding { + public static func from(pretrained model: String, hubApi: HubApi = .shared) async throws -> Embedding { + return try await BGEM3Model(repoName: model, hubApi: hubApi) + } +} + +class BERTEmbedding: Embedding { // Otherwise BERTModel + private let wordEmbedding: BNNS.EmbeddingLayer + private let positionEmbedding: BNNS.EmbeddingLayer + private let tokenTypeEmbedding: BNNS.EmbeddingLayer + private let normalization: BNNS.NormalizationLayer + private let dropout: BNNS.DropoutLayer + + private let positionEmbeddingType = "absolute" + + init(repoName: String) { fatalError() } + + public func callAsFunction(inputIds: MLMultiArray? = nil, + tokenTypeIDs: MLMultiArray? = nil, + positionIDs: MLMultiArray? = nil, + inputEmbeds: MLMultiArray? = nil, + pastKeyValuesLength: Int = 0) -> MLMultiArray { + fatalError() + } +} + +class BGEM3Model: Embedding { + + struct Output { + let lastHidddenState: MLMultiArray // batchSize, sequenceLength, hiddenSize + let hiddenStates: MLMultiArray? + let attentions: MLMultiArray? + + let loss: MLMultiArray? + let scores: MLMultiArray? + let pReps: MLMultiArray? + let qReps: MLMultiArray? + } + + let withSparse = false + let withDense = true + let withColbert = false + + let shouldNormalize = false +// let poolingMethod = "cls" +// let negativesCrossDevice = false +// let temperature = 1.0 +// let enableSubBatch = true +// let unifiedFinetuning = true +// let useSelfDistill = false +// let colbertDim: Int? = nil +// let selfDistillStartStep: Int? = nil + + private let tokenizer: Tokenizer + private let denseLayer: BNNS.FullyConnectedLayer + private let sparseLayer: BNNS.FullyConnectedLayer + private let colbertLayer: BNNS.FullyConnectedLayer + + init(repoName: String, hubApi: HubApi) async throws { + let config = LanguageModelConfigurationFromHub(modelName: repoName) + self.tokenizer = try await AutoTokenizer.from(pretrained: repoName, hubApi: hubApi) + + let hiddenSize = try await config.modelConfig.hiddenSize?.intValue ?? 384 + let colbertDim: Int? = nil + let denseInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + let denseOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(colbertDim ?? hiddenSize, stride: 2)) + let denseWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + self.denseLayer = BNNS.FullyConnectedLayer(input: denseInput, output: denseOutput, weights: denseWeights, bias: nil, activation: .identity)! + + let sparseInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + let sparseOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(1, stride: 2)) + let sparseWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + self.sparseLayer = BNNS.FullyConnectedLayer(input: sparseInput, output: sparseOutput, weights: sparseWeights, bias: nil, activation: .identity)! + + let colbertInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + let colbertOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(1, stride: 2)) + let colbertWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) + self.colbertLayer = BNNS.FullyConnectedLayer(input: colbertInput, output: colbertOutput, weights: colbertWeights, bias: nil, activation: .identity)! + } + + public func callAsFunction(_ textInput: (indices: MLMultiArray, attentionMask: MLMultiArray)) -> Output { + fatalError() + } + + private func forward(textInput: (indices: MLMultiArray, attentionMask: MLMultiArray)) -> [String: MLMultiArray] { + let lastHiddenState = self(textInput).lastHidddenState + + var output = [String: MLMultiArray]() + if withDense { + output["dense"] = self.dense(hiddenState: lastHiddenState, mask: textInput.attentionMask) + } + if withSparse { + output["sparse"] = self.sparse(hiddenState: lastHiddenState, mask: textInput.attentionMask) + } + if withColbert { + output["colbert"] = self.colbert(hiddenState: lastHiddenState, mask: textInput.attentionMask) + } + + if shouldNormalize { + if withDense { + // TODO: Normalize output["dense"] = + fatalError() + } + if withColbert { + // TODO: Normalize output["colbert"] = + fatalError() + } + } + + return output + } + + private func dense(hiddenState: MLMultiArray, mask: MLMultiArray) -> MLMultiArray { + assert(hiddenState.shape.count == 2) + var data = [Float]() + data.reserveCapacity(hiddenState.count) + + for index in 0.. MLMultiArray { + fatalError() + } + + private func colbert(hiddenState: MLMultiArray, mask: MLMultiArray) -> MLMultiArray { + fatalError() + } +} From 67d6c3d72131b2abc6bbb979e8ae6a733d3a010e Mon Sep 17 00:00:00 2001 From: shavit Date: Tue, 9 Jul 2024 14:22:36 -0400 Subject: [PATCH 2/6] BERT Embedding model Create layers with weights --- Package.swift | 2 +- Sources/Embedding/Embedding.swift | 260 ++++++++++++++++-------------- 2 files changed, 137 insertions(+), 125 deletions(-) diff --git a/Package.swift b/Package.swift index d0fdebd..8ef7eff 100644 --- a/Package.swift +++ b/Package.swift @@ -33,6 +33,6 @@ let package = Package( .testTarget(name: "TensorUtilsTests", dependencies: ["TensorUtils"]), .testTarget(name: "NormalizerTests", dependencies: ["Tokenizers", "Hub"]), .testTarget(name: "PostProcessorTests", dependencies: ["Tokenizers", "Hub"]), - .testTarget(name: "EmbeddingTests", dependencies: ["Embedding", "Tokenizers", "Hub"]) + .testTarget(name: "EmbeddingTests", dependencies: ["Embedding", "Tokenizers", "Hub"], resources: [.process("Resources"), .process("Vocabs")]) ] ) diff --git a/Sources/Embedding/Embedding.swift b/Sources/Embedding/Embedding.swift index c46e1a3..e4ef3e0 100644 --- a/Sources/Embedding/Embedding.swift +++ b/Sources/Embedding/Embedding.swift @@ -4,139 +4,151 @@ import CoreML import Accelerate -public protocol Embedding {} +class BERTEmbedding { -public struct AutoEmbedding {} // Otherwise AutoModel - -extension AutoEmbedding { - public static func from(pretrained model: String, hubApi: HubApi = .shared) async throws -> Embedding { - return try await BGEM3Model(repoName: model, hubApi: hubApi) - } -} - -class BERTEmbedding: Embedding { // Otherwise BERTModel - private let wordEmbedding: BNNS.EmbeddingLayer - private let positionEmbedding: BNNS.EmbeddingLayer - private let tokenTypeEmbedding: BNNS.EmbeddingLayer - private let normalization: BNNS.NormalizationLayer - private let dropout: BNNS.DropoutLayer - - private let positionEmbeddingType = "absolute" - - init(repoName: String) { fatalError() } - - public func callAsFunction(inputIds: MLMultiArray? = nil, - tokenTypeIDs: MLMultiArray? = nil, - positionIDs: MLMultiArray? = nil, - inputEmbeds: MLMultiArray? = nil, - pastKeyValuesLength: Int = 0) -> MLMultiArray { - fatalError() - } -} - -class BGEM3Model: Embedding { - - struct Output { - let lastHidddenState: MLMultiArray // batchSize, sequenceLength, hiddenSize - let hiddenStates: MLMultiArray? - let attentions: MLMultiArray? - - let loss: MLMultiArray? - let scores: MLMultiArray? - let pReps: MLMultiArray? - let qReps: MLMultiArray? - } - - let withSparse = false - let withDense = true - let withColbert = false - - let shouldNormalize = false -// let poolingMethod = "cls" -// let negativesCrossDevice = false -// let temperature = 1.0 -// let enableSubBatch = true -// let unifiedFinetuning = true -// let useSelfDistill = false -// let colbertDim: Int? = nil -// let selfDistillStartStep: Int? = nil - - private let tokenizer: Tokenizer - private let denseLayer: BNNS.FullyConnectedLayer - private let sparseLayer: BNNS.FullyConnectedLayer - private let colbertLayer: BNNS.FullyConnectedLayer - - init(repoName: String, hubApi: HubApi) async throws { - let config = LanguageModelConfigurationFromHub(modelName: repoName) - self.tokenizer = try await AutoTokenizer.from(pretrained: repoName, hubApi: hubApi) - - let hiddenSize = try await config.modelConfig.hiddenSize?.intValue ?? 384 - let colbertDim: Int? = nil - let denseInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - let denseOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(colbertDim ?? hiddenSize, stride: 2)) - let denseWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - self.denseLayer = BNNS.FullyConnectedLayer(input: denseInput, output: denseOutput, weights: denseWeights, bias: nil, activation: .identity)! + typealias Weights = [String: MLMultiArray] + + var shape: [NSNumber] {[ + NSNumber(value: maxPositionEmbeddings), + NSNumber(value: hiddenSize), + ]} + + private let weights: Weights + + private let positionEmbeddingType: String + private let hiddenSize: Int + private let vocabSize: Int + private let maxPositionEmbeddings: Int + private let typeVocabSize: Int + private let padTokenID: Int + private let normalizationEpsilon: Float + private let dropoutRate: Float = 1e-1 + private let hiddenActivation: BNNS.ActivationFunction = .geluApproximation2(alpha: 1e-1, beta: 1e-1) + + private var allocations: [BNNSNDArrayDescriptor] = [] + + private lazy var wordEmbedding: BNNS.EmbeddingLayer = { + let input = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Int64.self, shape: .vector(maxPositionEmbeddings)) + allocations.append(input) + let dictData: [Float32] = weights["bert.embeddings.word_embeddings.weight"]!.toArray() + let dict = BNNSNDArrayDescriptor.allocate(initializingFrom: dictData, shape: .matrixColumnMajor(hiddenSize, vocabSize)) + allocations.append(dict) + let output = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(output) - let sparseInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - let sparseOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(1, stride: 2)) - let sparseWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - self.sparseLayer = BNNS.FullyConnectedLayer(input: sparseInput, output: sparseOutput, weights: sparseWeights, bias: nil, activation: .identity)! + return BNNS.EmbeddingLayer(input: input, output: output, dictionary: dict, paddingIndex: 0, maximumNorm: 0, normType: .l2, scalesGradientByFrequency: false)! + }() + + private lazy var positionEmbedding: BNNS.EmbeddingLayer = { + let input = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Int64.self, shape: .vector(maxPositionEmbeddings)) + allocations.append(input) + let dictData: [Float32] = weights["bert.embeddings.position_embeddings.weight"]!.toArray() + let dict = BNNSNDArrayDescriptor.allocate(initializingFrom: dictData, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(dict) + let output = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(output) + + return BNNS.EmbeddingLayer(input: input, output: output, dictionary: dict, paddingIndex: -1, maximumNorm: 0, normType: .l2, scalesGradientByFrequency: true)! + }() + + private lazy var tokenTypeEmbedding: BNNS.EmbeddingLayer = { + let input = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Int64.self, shape: .vector(maxPositionEmbeddings)) + allocations.append(input) + let dictData: [Float32] = weights["bert.embeddings.token_type_embeddings.weight"]!.toArray() + let dict = BNNSNDArrayDescriptor.allocate(initializingFrom: dictData, shape: .matrixColumnMajor(hiddenSize, typeVocabSize)) + allocations.append(dict) + let output = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(output) - let colbertInput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - let colbertOutput = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(1, stride: 2)) - let colbertWeights = BNNSNDArrayDescriptor(dataType: .float16, shape: .vector(hiddenSize, stride: 2)) - self.colbertLayer = BNNS.FullyConnectedLayer(input: colbertInput, output: colbertOutput, weights: colbertWeights, bias: nil, activation: .identity)! - } - - public func callAsFunction(_ textInput: (indices: MLMultiArray, attentionMask: MLMultiArray)) -> Output { - fatalError() + return BNNS.EmbeddingLayer(input: input, output: output, dictionary: dict, paddingIndex: -1, maximumNorm: 0, normType: .l2, scalesGradientByFrequency: true)! + }() + + private lazy var normalization: BNNS.NormalizationLayer = { + let input = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixRowMajor(maxPositionEmbeddings, hiddenSize)) + allocations.append(input) + let output = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixRowMajor(maxPositionEmbeddings, hiddenSize)) + allocations.append(output) + + let betaWA: MLMultiArray! = weights["bert.embeddings.LayerNorm.beta"] ?? weights["bert.embeddings.LayerNorm.bias"] + let beta = BNNSNDArrayDescriptor.allocate(initializingFrom: betaWA.toArray() as [Float32], shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(beta) + + let gammaWA: MLMultiArray! = weights["bert.embeddings.LayerNorm.gamma"] ?? weights["bert.embeddings.LayerNorm.weight"] + let gamma = BNNSNDArrayDescriptor.allocate(initializingFrom: gammaWA.toArray() as [Float32], shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(gamma) + + return BNNS.NormalizationLayer(type: .batch(movingMean: nil, movingVariance: nil), input: input, output: output, beta: beta, gamma: gamma, epsilon: normalizationEpsilon, activation: hiddenActivation)! + }() + + private lazy var dropout: BNNS.DropoutLayer = { + let input = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(input) + let output = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + allocations.append(output) + + return BNNS.DropoutLayer(input: input, output: output, rate: dropoutRate, seed: 0, control: 0)! + }() + + deinit { + allocations.forEach({ $0.deallocate() }) } - private func forward(textInput: (indices: MLMultiArray, attentionMask: MLMultiArray)) -> [String: MLMultiArray] { - let lastHiddenState = self(textInput).lastHidddenState - - var output = [String: MLMultiArray]() - if withDense { - output["dense"] = self.dense(hiddenState: lastHiddenState, mask: textInput.attentionMask) - } - if withSparse { - output["sparse"] = self.sparse(hiddenState: lastHiddenState, mask: textInput.attentionMask) + init(config: Config, weights: Weights = [:]) { + assert(config.model_type!.stringValue == "bert") + for key in [ + "bert.embeddings.word_embeddings.weight", + "bert.embeddings.position_embeddings.weight", + "bert.embeddings.token_type_embeddings.weight", + ] { assert(weights.keys.contains(where: { $0 == key })) } + assert(weights.keys.contains(where: { $0 == "bert.embeddings.LayerNorm.beta" || $0 == "bert.embeddings.LayerNorm.bias" })) + assert(weights.keys.contains(where: { $0 == "bert.embeddings.LayerNorm.gamma" || $0 == "bert.embeddings.LayerNorm.weight" })) + assert(config.hidden_act!.stringValue == "gelu") + assert("absolute" == config.position_embedding_type!.stringValue!) + self.positionEmbeddingType = config.position_embedding_type!.stringValue! + self.hiddenSize = config.hidden_size!.intValue! + self.vocabSize = config.vocab_size!.intValue! + self.maxPositionEmbeddings = config.max_position_embeddings!.intValue! + self.typeVocabSize = config.type_vocab_size!.intValue! + self.padTokenID = config.pad_token_id!.intValue! + self.normalizationEpsilon = Float(config.layer_norm_eps!.doubleValue!) + self.weights = weights + } + + public func callAsFunction(inputIDs: [Int64], + tokenTypeIDs: [Int64]? = nil, + positionIDs: [Int64]? = nil) -> MLMultiArray { + let inputLength = inputIDs.count + let inputIDs: [Int64] = inputIDs.padded(length: maxPositionEmbeddings) + let wordInput = BNNSNDArrayDescriptor.allocate(initializingFrom: inputIDs, shape: .vector(inputIDs.count)) + let wordOutput = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, inputIDs.count)) + defer { + wordInput.deallocate() + wordOutput.deallocate() } - if withColbert { - output["colbert"] = self.colbert(hiddenState: lastHiddenState, mask: textInput.attentionMask) + try! wordEmbedding.apply(batchSize: 1, input: wordInput, output: wordOutput) + + let positionIDs = positionIDs ?? Array(stride(from: 0, through: Int64(inputLength - 1), by: 1)) + let positionInput = BNNSNDArrayDescriptor.allocate(initializingFrom: positionIDs.padded(length: maxPositionEmbeddings), shape: .vector(maxPositionEmbeddings)) + let positionOutput = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + defer { + positionInput.deallocate() + positionOutput.deallocate() } - - if shouldNormalize { - if withDense { - // TODO: Normalize output["dense"] = - fatalError() - } - if withColbert { - // TODO: Normalize output["colbert"] = - fatalError() - } + try! self.positionEmbedding.apply(batchSize: 1, input: positionInput, output: positionOutput) + + let tokenTypeIDs: [Int64] = tokenTypeIDs ?? Array(repeating: 0, count: maxPositionEmbeddings) + let typeInput = BNNSNDArrayDescriptor.allocate(initializingFrom: tokenTypeIDs, shape: .vector(maxPositionEmbeddings)) + let typeOutput = BNNSNDArrayDescriptor.allocateUninitialized(scalarType: Float32.self, shape: .matrixColumnMajor(hiddenSize, maxPositionEmbeddings)) + defer { + typeInput.deallocate() + typeOutput.deallocate() } + try! self.tokenTypeEmbedding.apply(batchSize: 1, input: typeInput, output: typeOutput) - return output - } - - private func dense(hiddenState: MLMultiArray, mask: MLMultiArray) -> MLMultiArray { - assert(hiddenState.shape.count == 2) - var data = [Float]() - data.reserveCapacity(hiddenState.count) - - for index in 0.. MLMultiArray { - fatalError() - } + let multiWord = try! wordOutput.makeMultiArray(of: Float32.self, shape: shape) + let multiPosition = try! positionOutput.makeMultiArray(of: Float32.self, shape: shape) + let multiType = try! typeOutput.makeMultiArray(of: Float32.self, shape: shape) - private func colbert(hiddenState: MLMultiArray, mask: MLMultiArray) -> MLMultiArray { - fatalError() + return multiWord + multiPosition + multiType } } From 126845fa06edd59d9f4f46341ff22eca63d1b203 Mon Sep 17 00:00:00 2001 From: shavit Date: Tue, 9 Jul 2024 14:24:02 -0400 Subject: [PATCH 3/6] Array padding --- Package.swift | 2 +- Sources/TensorUtils/Array+Utils.swift | 8 +++++++ Tests/EmbeddingTests/EmbeddingTests.swift | 7 ++++++ Tests/TensorUtilsTests/ArrayUtilsTests.swift | 23 ++++++++++++++++++++ 4 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 Sources/TensorUtils/Array+Utils.swift create mode 100644 Tests/EmbeddingTests/EmbeddingTests.swift create mode 100644 Tests/TensorUtilsTests/ArrayUtilsTests.swift diff --git a/Package.swift b/Package.swift index 8ef7eff..5050e92 100644 --- a/Package.swift +++ b/Package.swift @@ -33,6 +33,6 @@ let package = Package( .testTarget(name: "TensorUtilsTests", dependencies: ["TensorUtils"]), .testTarget(name: "NormalizerTests", dependencies: ["Tokenizers", "Hub"]), .testTarget(name: "PostProcessorTests", dependencies: ["Tokenizers", "Hub"]), - .testTarget(name: "EmbeddingTests", dependencies: ["Embedding", "Tokenizers", "Hub"], resources: [.process("Resources"), .process("Vocabs")]) + .testTarget(name: "EmbeddingTests", dependencies: ["Embedding", "Tokenizers", "Hub", "TensorUtils"], resources: [.process("Resources"), .process("Vocabs")]) ] ) diff --git a/Sources/TensorUtils/Array+Utils.swift b/Sources/TensorUtils/Array+Utils.swift new file mode 100644 index 0000000..5353afb --- /dev/null +++ b/Sources/TensorUtils/Array+Utils.swift @@ -0,0 +1,8 @@ +import Foundation + + +extension Array where Element: Numeric { + func padded(length maxLength: Int) -> Array { + self + Array(repeating: 0, count: Swift.max(maxLength - count, 0)) + } +} diff --git a/Tests/EmbeddingTests/EmbeddingTests.swift b/Tests/EmbeddingTests/EmbeddingTests.swift new file mode 100644 index 0000000..1e68b37 --- /dev/null +++ b/Tests/EmbeddingTests/EmbeddingTests.swift @@ -0,0 +1,7 @@ +import XCTest +@testable import Tokenizers +@testable import Hub +@testable import TensorUtils +@testable import Embedding + +class EmbeddingTests: XCTestCase { } diff --git a/Tests/TensorUtilsTests/ArrayUtilsTests.swift b/Tests/TensorUtilsTests/ArrayUtilsTests.swift new file mode 100644 index 0000000..73d55c6 --- /dev/null +++ b/Tests/TensorUtilsTests/ArrayUtilsTests.swift @@ -0,0 +1,23 @@ +import XCTest +@testable import TensorUtils + +class ArrayUtilsTests: XCTestCase { + + func testPaddedArrayWhenNeedPadding() { + let array = [1, 2, 3, 4] + let paddedArray = array.padded(length: 7) + XCTAssertEqual(paddedArray, [1, 2, 3, 4, 0, 0, 0]) + } + + func testNoPaddingForTheSamePaddingLength() { + let array = [1, 2, 3, 4] + let paddedArray = array.padded(length: 4) + XCTAssertEqual(paddedArray, [1, 2, 3, 4]) + } + + func testNoPaddingForShorterPaddingLength() { + let array = [1, 2, 3, 4] + let paddedArray = array.padded(length: 2) + XCTAssertEqual(paddedArray, [1, 2, 3, 4]) + } +} From e0259ba290eddb7b4b7b79d6a4ae066d1420d3d5 Mon Sep 17 00:00:00 2001 From: shavit Date: Tue, 9 Jul 2024 14:24:26 -0400 Subject: [PATCH 4/6] MLMultiArray toArray --- Sources/TensorUtils/MLMultiArray+Utils.swift | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/Sources/TensorUtils/MLMultiArray+Utils.swift b/Sources/TensorUtils/MLMultiArray+Utils.swift index ddb2760..2e29a7c 100644 --- a/Sources/TensorUtils/MLMultiArray+Utils.swift +++ b/Sources/TensorUtils/MLMultiArray+Utils.swift @@ -198,3 +198,15 @@ extension MLMultiArray { return s + "]" } } + +extension MLMultiArray { + func toArray() -> Array { + let stride = MemoryLayout.stride + let allocated = UnsafeMutableRawBufferPointer.allocate(byteCount: self.count * stride, alignment: MemoryLayout.alignment) + return self.withUnsafeBytes { ptr in + memcpy(allocated.baseAddress!, ptr.baseAddress!, self.count * stride) + let start = allocated.bindMemory(to: T.self).baseAddress! + return Array(UnsafeBufferPointer(start: start, count: self.count)) + } + } +} From 4c5e6ac0ba22aa69f136a44f0f33154f2837ec7d Mon Sep 17 00:00:00 2001 From: shavit Date: Tue, 9 Jul 2024 14:30:02 -0400 Subject: [PATCH 5/6] Add MLMultiArray helpers and tests * Test addition of MLMultiArrays and broadcasting * Create generic number array from MLMultiArray * Add double type to config --- Sources/Hub/Hub.swift | 1 + Sources/TensorUtils/Array+Utils.swift | 2 +- Sources/TensorUtils/BNNS+Utils.swift | 21 ++++ Sources/TensorUtils/MLMultiArray+Utils.swift | 42 ++++++- .../MLMultiArrayUtilsTests.swift | 106 ++++++++++++++++++ 5 files changed, 170 insertions(+), 2 deletions(-) create mode 100644 Sources/TensorUtils/BNNS+Utils.swift create mode 100644 Tests/TensorUtilsTests/MLMultiArrayUtilsTests.swift diff --git a/Sources/Hub/Hub.swift b/Sources/Hub/Hub.swift index 8deeee8..d5aa649 100644 --- a/Sources/Hub/Hub.swift +++ b/Sources/Hub/Hub.swift @@ -90,6 +90,7 @@ public struct Config { } public var intValue: Int? { value as? Int } + public var doubleValue: Double? { value as? Double } public var boolValue: Bool? { value as? Bool } public var stringValue: String? { value as? String } diff --git a/Sources/TensorUtils/Array+Utils.swift b/Sources/TensorUtils/Array+Utils.swift index 5353afb..38c32a9 100644 --- a/Sources/TensorUtils/Array+Utils.swift +++ b/Sources/TensorUtils/Array+Utils.swift @@ -1,7 +1,7 @@ import Foundation -extension Array where Element: Numeric { +public extension Array where Element: Numeric { func padded(length maxLength: Int) -> Array { self + Array(repeating: 0, count: Swift.max(maxLength - count, 0)) } diff --git a/Sources/TensorUtils/BNNS+Utils.swift b/Sources/TensorUtils/BNNS+Utils.swift new file mode 100644 index 0000000..662445d --- /dev/null +++ b/Sources/TensorUtils/BNNS+Utils.swift @@ -0,0 +1,21 @@ +import Accelerate +import CoreML.MLMultiArray + + +public extension BNNSNDArrayDescriptor { + func makeMultiArray(of: T.Type, shape: [NSNumber]) throws -> MLMultiArray { + let dataType: MLMultiArrayDataType + switch of { + case is Int32.Type: dataType = .int32 + case is Float32.Type: dataType = .float32 + case is Double.Type: dataType = .double + default: fatalError("type not supported") + } + + let strides = shape.dropFirst().reversed().reduce(into: [1]) { acc, a in + acc.insert(acc[0].intValue * a.intValue as NSNumber, at: 0) + } + + return try MLMultiArray(dataPointer: self.data!, shape: shape, dataType: dataType, strides: strides) + } +} diff --git a/Sources/TensorUtils/MLMultiArray+Utils.swift b/Sources/TensorUtils/MLMultiArray+Utils.swift index 2e29a7c..623bc87 100644 --- a/Sources/TensorUtils/MLMultiArray+Utils.swift +++ b/Sources/TensorUtils/MLMultiArray+Utils.swift @@ -8,6 +8,7 @@ import Foundation import CoreML +import Accelerate public extension MLMultiArray { /// All values will be stored in the last dimension of the MLMultiArray (default is dims=1) @@ -199,7 +200,7 @@ extension MLMultiArray { } } -extension MLMultiArray { +public extension MLMultiArray { func toArray() -> Array { let stride = MemoryLayout.stride let allocated = UnsafeMutableRawBufferPointer.allocate(byteCount: self.count * stride, alignment: MemoryLayout.alignment) @@ -210,3 +211,42 @@ extension MLMultiArray { } } } + +public extension MLMultiArray { + static func +(lhs: MLMultiArray, rhs: MLMultiArray) -> MLMultiArray { + assert(lhs.dataType == rhs.dataType && lhs.dataType == .float32) + assert(lhs.shape.count == rhs.shape.count && lhs.shape[1].intValue == rhs.shape[1].intValue) + + let outShape: [NSNumber] + let outLength: Int + var ptr0: UnsafeMutablePointer + var ptr1: UnsafeMutablePointer + if lhs.shape[0].intValue >= rhs.shape[0].intValue { + assert(rhs.shape[0].intValue == 1 || lhs.shape == rhs.shape) // A[m, n], B[1, n] || B[m, n] + outShape = lhs.shape + outLength = lhs.count + ptr0 = UnsafeMutablePointer(OpaquePointer(lhs.withUnsafeMutableBytes({ ptr, _ in ptr.baseAddress! }))) + ptr1 = UnsafeMutablePointer(OpaquePointer(rhs.withUnsafeMutableBytes({ ptr, _ in ptr.baseAddress! }))) + } else { + assert(lhs.shape[0].intValue == 1) // Swap when A[1, n], B[m, n] + outShape = rhs.shape + outLength = rhs.count + ptr0 = UnsafeMutablePointer(OpaquePointer(rhs.withUnsafeMutableBytes({ ptr, _ in ptr.baseAddress! }))) + ptr1 = UnsafeMutablePointer(OpaquePointer(lhs.withUnsafeMutableBytes({ ptr, _ in ptr.baseAddress! }))) + } + + let output = try! MLMultiArray(shape: outShape, dataType: .float32) + var ptrOutput = UnsafeMutablePointer(OpaquePointer(output.withUnsafeMutableBytes({ ptr, _ in ptr.baseAddress! }))) + vDSP_vadd(ptr0, 1, ptr1, 1, ptrOutput, 1, vDSP_Length(outLength)) + + if lhs.shape[0].intValue != rhs.shape[0].intValue { + for _ in 1...stride + let allocated = UnsafeMutableRawBufferPointer.allocate(byteCount: array.count * stride, alignment: MemoryLayout.alignment) + defer { allocated.deallocate() } + _ = array.withUnsafeBufferPointer { ptr in + memcpy(allocated.baseAddress!, ptr.baseAddress!, array.count * stride) + } + let multiArray = try MLMultiArray(dataPointer: allocated.baseAddress!, shape: [4, 10], dataType: .float32, strides: [10, 1]) + let output = multiArray + multiArray + XCTAssertEqual(output.count, array.count) + XCTAssertEqual(output.count, multiArray.count) + + for index in 0...stride + let allocA = UnsafeMutableRawBufferPointer.allocate(byteCount: array.count * stride, alignment: MemoryLayout.alignment) + defer { allocA.deallocate() } + let allocB = UnsafeMutableRawBufferPointer.allocate(byteCount: 10 * stride, alignment: MemoryLayout.alignment) + defer { allocB.deallocate() } + + _ = array.withUnsafeBufferPointer { ptr in + memcpy(allocA.baseAddress!, ptr.baseAddress!, array.count * stride) + } + _ = Array(repeating: 10, count: 10).withUnsafeBufferPointer { ptr in + memcpy(allocB.baseAddress!, ptr.baseAddress!, 10 * stride) + } + + let A = try MLMultiArray(dataPointer: allocA.baseAddress!, shape: [4, 10], dataType: .float32, strides: [10, 1]) + XCTAssertEqual(A.count, 40) + let B = try MLMultiArray(dataPointer: allocB.baseAddress!, shape: [1, 10], dataType: .float32, strides: [10, 1]) + XCTAssertEqual(B.count, 10) + + _ = A + B + _ = A + B + B + _ = A + B + B + let expectedArray: [Float32] = [ + 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, + ] + let output = A + B + XCTAssertEqual(output.floats, expectedArray) + } + + func testAdditionRowReverseOrder() throws { + let array: [Float32] = [ + 01, 02, 03, 04, 05, 06, 07, 08, 09, 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, + ] + let stride = MemoryLayout.stride + let allocA = UnsafeMutableRawBufferPointer.allocate(byteCount: array.count * stride, alignment: MemoryLayout.alignment) + defer { allocA.deallocate() } + let allocB = UnsafeMutableRawBufferPointer.allocate(byteCount: 10 * stride, alignment: MemoryLayout.alignment) + defer { allocB.deallocate() } + + _ = array.withUnsafeBufferPointer { ptr in + memcpy(allocA.baseAddress!, ptr.baseAddress!, array.count * stride) + } + _ = Array(repeating: 10, count: 10).withUnsafeBufferPointer { ptr in + memcpy(allocB.baseAddress!, ptr.baseAddress!, 10 * stride) + } + + let A = try MLMultiArray(dataPointer: allocA.baseAddress!, shape: [4, 10], dataType: .float32, strides: [10, 1]) + XCTAssertEqual(A.count, 40) + let B = try MLMultiArray(dataPointer: allocB.baseAddress!, shape: [1, 10], dataType: .float32, strides: [10, 1]) + XCTAssertEqual(B.count, 10) + XCTAssertEqual(B + A, A + B) + _ = A + B + _ = A + B + B + _ = A + B + B + let expectedArray: [Float32] = [ + 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, + ] + let output = B + A + XCTAssertEqual(output.floats, expectedArray) + } +} From f8d57e178333ee89f54ac2385a478c5ba1296eb7 Mon Sep 17 00:00:00 2001 From: shavit Date: Tue, 9 Jul 2024 14:30:14 -0400 Subject: [PATCH 6/6] Support only float32 to make multi arrays --- Sources/TensorUtils/BNNS+Utils.swift | 13 +++-------- Tests/TensorUtilsTests/BNNSUtilTests.swift | 26 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 10 deletions(-) create mode 100644 Tests/TensorUtilsTests/BNNSUtilTests.swift diff --git a/Sources/TensorUtils/BNNS+Utils.swift b/Sources/TensorUtils/BNNS+Utils.swift index 662445d..f5b70ae 100644 --- a/Sources/TensorUtils/BNNS+Utils.swift +++ b/Sources/TensorUtils/BNNS+Utils.swift @@ -3,19 +3,12 @@ import CoreML.MLMultiArray public extension BNNSNDArrayDescriptor { - func makeMultiArray(of: T.Type, shape: [NSNumber]) throws -> MLMultiArray { - let dataType: MLMultiArrayDataType - switch of { - case is Int32.Type: dataType = .int32 - case is Float32.Type: dataType = .float32 - case is Double.Type: dataType = .double - default: fatalError("type not supported") - } - + func makeMultiArray(of numericType: T.Type, shape: [NSNumber]) throws -> MLMultiArray { + assert(numericType == Float32.self) let strides = shape.dropFirst().reversed().reduce(into: [1]) { acc, a in acc.insert(acc[0].intValue * a.intValue as NSNumber, at: 0) } - return try MLMultiArray(dataPointer: self.data!, shape: shape, dataType: dataType, strides: strides) + return try MLMultiArray(dataPointer: self.data!, shape: shape, dataType: .float32, strides: strides) } } diff --git a/Tests/TensorUtilsTests/BNNSUtilTests.swift b/Tests/TensorUtilsTests/BNNSUtilTests.swift new file mode 100644 index 0000000..86ea958 --- /dev/null +++ b/Tests/TensorUtilsTests/BNNSUtilTests.swift @@ -0,0 +1,26 @@ +import XCTest +import Accelerate +@testable import TensorUtils + +class BNNSUtilsTests: XCTestCase { + + func testMakeMultiArrayFromDescriptor() throws { + let rowCount = 4 + let dimSize = 6 + let dictData: [Float32] = [ + 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, + 19, 20, 21, 22, 23, 24, + ] + let dict = BNNSNDArrayDescriptor.allocate(initializingFrom: dictData, shape: .matrixColumnMajor(dimSize, rowCount)) + let shape: [NSNumber] = [ + NSNumber(value: rowCount), + NSNumber(value: dimSize), + ] + let multiArray = try dict.makeMultiArray(of: Float32.self, shape: shape) + XCTAssertEqual(multiArray.toArray(), dictData) + XCTAssertEqual(multiArray.floats!, dictData) + } + +}