diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift deleted file mode 100644 index 013b9b6..0000000 --- a/Sources/MarkBase/Embedding/EmbeddingGemmaConfig.swift +++ /dev/null @@ -1,34 +0,0 @@ -import Foundation - -/// EmbeddingGemma model configuration -public struct EmbeddingGemmaConfig: Codable { - public let hiddenSize: Int - public let numHiddenLayers: Int - public let vocabSize: Int - public let numAttentionHeads: Int - public let numKeyValueHeads: Int - public let headDim: Int - public let intermediateSize: Int - public let maxPositionEmbeddings: Int - public let slidingWindow: Int - public let rmsNormEps: Float - public let ropeTheta: Float - public let useBidirectionalAttention: Bool - public let layerTypes: [String] - - enum CodingKeys: String, CodingKey { - case hiddenSize = "hidden_size" - case numHiddenLayers = "num_hidden_layers" - case vocabSize = "vocab_size" - case numAttentionHeads = "num_attention_heads" - case numKeyValueHeads = "num_key_value_heads" - case headDim = "head_dim" - case intermediateSize = "intermediate_size" - case maxPositionEmbeddings = "max_position_embeddings" - case slidingWindow = "sliding_window" - case rmsNormEps = "rms_norm_eps" - case ropeTheta = "rope_theta" - case useBidirectionalAttention = "use_bidirectional_attention" - case layerTypes = "layer_types" - } -} diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift new file mode 100644 index 0000000..3859f42 --- /dev/null +++ b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift @@ -0,0 +1,372 @@ +import Foundation +import Metal +import Accelerate + +/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model +public struct EmbeddingGemmaConfig: Codable { + public let hiddenSize: Int + public let numHiddenLayers: Int + public let vocabSize: Int + public let numAttentionHeads: Int + public let numKeyValueHeads: Int + public let headDim: Int + public let intermediateSize: Int + public let maxPositionEmbeddings: Int + public let slidingWindow: Int + public let rmsNormEps: Float + public let ropeTheta: Float + public let useBidirectionalAttention: Bool + public let layerTypes: [String] + + enum CodingKeys: String, CodingKey { + case hiddenSize = "hidden_size" + case numHiddenLayers = "num_hidden_layers" + case vocabSize = "vocab_size" + case numAttentionHeads = "num_attention_heads" + case numKeyValueHeads = "num_key_value_heads" + case headDim = "head_dim" + case intermediateSize = "intermediate_size" + case maxPositionEmbeddings = "max_position_embeddings" + case slidingWindow = "sliding_window" + case rmsNormEps = "rms_norm_eps" + case ropeTheta = "rope_theta" + case useBidirectionalAttention = "use_bidirectional_attention" + case layerTypes = "layer_types" + } + + public static func load(from modelDir: String) throws -> Self { + let url = URL(fileURLWithPath: modelDir).appendingPathComponent("config.json") + let data = try Data(contentsOf: url) + return try JSONDecoder().decode(Self.self, from: data) + } +} + +/// EmbeddingGemma - Google's 300M parameter embedding model +public final class EmbeddingGemmaModel { + public let config: EmbeddingGemmaConfig + public let engine: MarkBaseEngine + public let tokenizer: Tokenizer + public let reader: SafeTensorsReader + + // GPU Buffers + public var embedTokens: MTLBuffer! + public var finalNorm: MTLBuffer! + public var layerNorms: [[MTLBuffer]] = [] + public var qProjs: [MTLBuffer] = [] + public var kProjs: [MTLBuffer] = [] + public var vProjs: [MTLBuffer] = [] + public var oProjs: [MTLBuffer] = [] + public var qNorms: [MTLBuffer] = [] + public var kNorms: [MTLBuffer] = [] + public var gateProjs: [MTLBuffer] = [] + public var upProjs: [MTLBuffer] = [] + public var downProjs: [MTLBuffer] = [] + + public init(modelDir: String, engine: MarkBaseEngine) throws { + self.engine = engine + self.config = try EmbeddingGemmaConfig.load(from: modelDir) + self.tokenizer = try TokenizerFactory.load(modelDir: modelDir) + self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors") + + try loadWeights() + print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))") + } + + private func loadWeights() throws { + let hs = config.hiddenSize + let intermedi = config.intermediateSize + let nKV = config.numKeyValueHeads + let hDim = config.headDim + + // Embedding table [vocab, hidden] + let embedData = try readTensor("embed_tokens.weight") + embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)! + + for i in 0.. [Float] { + var tokens = tokenizer.encode(text: text) + if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) } + guard !tokens.isEmpty else { return [] } + + let seqLen = tokens.count, hs = config.hiddenSize + + // Embedding lookup + let inputBuf = try lookupEmbeddings(tokens: tokens) + + // Forward through layers + var hidden = inputBuf + for idx in 0..