v2: Initial clean branch with unit tests + CI/CD pipeline
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,182 @@
|
||||
import Foundation
|
||||
import Accelerate
|
||||
|
||||
public final class PCA: Codable, @unchecked Sendable {
|
||||
public let inputDimension: Int
|
||||
public let outputDimension: Int
|
||||
public let mean: [Float]
|
||||
public let components: [[Float]]
|
||||
public let explainedVariance: [Float]
|
||||
public let whiteningEnabled: Bool
|
||||
public let sampleCount: Int
|
||||
private let trainingData: [[Float]]
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case inputDimension, outputDimension, mean, components, explainedVariance, whiteningEnabled, sampleCount, trainingData
|
||||
}
|
||||
|
||||
public init(inputDimension: Int, outputDimension: Int, mean: [Float], components: [[Float]], explainedVariance: [Float], whiteningEnabled: Bool = false, sampleCount: Int = 0, trainingData: [[Float]] = []) {
|
||||
self.inputDimension = inputDimension
|
||||
self.outputDimension = outputDimension
|
||||
self.mean = mean
|
||||
self.components = components
|
||||
self.explainedVariance = explainedVariance
|
||||
self.whiteningEnabled = whiteningEnabled
|
||||
self.sampleCount = sampleCount
|
||||
self.trainingData = trainingData
|
||||
}
|
||||
|
||||
public func transform(_ input: [Float]) throws -> [Float] {
|
||||
guard input.count == inputDimension else {
|
||||
throw PCAError.dimensionMismatch
|
||||
}
|
||||
var centered = [Float](repeating: 0, count: inputDimension)
|
||||
for i in 0..<inputDimension {
|
||||
centered[i] = input[i] - mean[i]
|
||||
}
|
||||
var result = [Float](repeating: 0, count: outputDimension)
|
||||
for j in 0..<outputDimension {
|
||||
var dot: Float = 0
|
||||
for i in 0..<inputDimension {
|
||||
dot += centered[i] * components[j][i]
|
||||
}
|
||||
result[j] = dot
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
public func transformWhitened(_ input: [Float]) throws -> [Float] {
|
||||
var result = try transform(input)
|
||||
for j in 0..<outputDimension {
|
||||
let denom = sqrt(max(explainedVariance[j], 1e-10))
|
||||
result[j] /= denom
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
public func explainedVarianceRatio() -> [Float] {
|
||||
let total = explainedVariance.reduce(0, +)
|
||||
guard total > 0 else { return explainedVariance.map { _ in 0 } }
|
||||
return explainedVariance.map { $0 / total }
|
||||
}
|
||||
|
||||
public func cumulativeExplainedVarianceRatio() -> [Float] {
|
||||
let ratios = explainedVarianceRatio()
|
||||
var cum: [Float] = []
|
||||
var sum: Float = 0
|
||||
for r in ratios {
|
||||
sum += r
|
||||
cum.append(sum)
|
||||
}
|
||||
return cum
|
||||
}
|
||||
|
||||
public func save(to url: URL) throws {
|
||||
let encoder = JSONEncoder()
|
||||
let data = try encoder.encode(self)
|
||||
try data.write(to: url)
|
||||
}
|
||||
|
||||
public static func load(from url: URL) throws -> PCA {
|
||||
let data = try Data(contentsOf: url)
|
||||
let decoder = JSONDecoder()
|
||||
return try decoder.decode(PCA.self, from: data)
|
||||
}
|
||||
|
||||
public static func train(data: [[Float]], outputDimension: Int, whitening: Bool = false) throws -> PCA {
|
||||
guard let first = data.first else { throw PCAError.noData }
|
||||
let n = data.count
|
||||
let d = first.count
|
||||
let k = min(outputDimension, d, n)
|
||||
|
||||
guard k > 0 else { throw PCAError.invalidDimension }
|
||||
|
||||
var mean = [Float](repeating: 0, count: d)
|
||||
for i in 0..<n {
|
||||
for j in 0..<d {
|
||||
mean[j] += data[i][j]
|
||||
}
|
||||
}
|
||||
for j in 0..<d {
|
||||
mean[j] /= Float(n)
|
||||
}
|
||||
|
||||
var A = [Float](repeating: 0, count: n * d)
|
||||
for j in 0..<d {
|
||||
for i in 0..<n {
|
||||
A[i + j * n] = data[i][j] - mean[j]
|
||||
}
|
||||
}
|
||||
|
||||
let m = n
|
||||
var m32 = Int32(m)
|
||||
var n32 = Int32(d)
|
||||
var lda = Int32(m)
|
||||
var ldu = Int32(1)
|
||||
var ldvt = Int32(d)
|
||||
var s = [Float](repeating: 0, count: min(m, d))
|
||||
var u = [Float](repeating: 0, count: 1)
|
||||
var vt = [Float](repeating: 0, count: d * d)
|
||||
var lwork = Int32(-1)
|
||||
var work: [Float] = [0]
|
||||
var info = Int32(0)
|
||||
var jobU = Int8(78)
|
||||
var jobVT = Int8(65)
|
||||
|
||||
sgesvd_(&jobU, &jobVT, &m32, &n32, &A, &lda, &s, &u, &ldu, &vt, &ldvt, &work, &lwork, &info)
|
||||
guard info == 0 else { throw PCAError.svdFailed }
|
||||
|
||||
lwork = Int32(work[0])
|
||||
work = [Float](repeating: 0, count: Int(lwork))
|
||||
|
||||
sgesvd_(&jobU, &jobVT, &m32, &n32, &A, &lda, &s, &u, &ldu, &vt, &ldvt, &work, &lwork, &info)
|
||||
guard info == 0 else { throw PCAError.svdFailed }
|
||||
|
||||
var components: [[Float]] = []
|
||||
var explainedVariance: [Float] = []
|
||||
for i in 0..<k {
|
||||
var comp = [Float](repeating: 0, count: d)
|
||||
for j in 0..<d {
|
||||
comp[j] = vt[i + j * d]
|
||||
}
|
||||
components.append(comp)
|
||||
explainedVariance.append(s[i] * s[i] / Float(n - 1))
|
||||
}
|
||||
|
||||
return PCA(
|
||||
inputDimension: d,
|
||||
outputDimension: k,
|
||||
mean: mean,
|
||||
components: components,
|
||||
explainedVariance: explainedVariance,
|
||||
whiteningEnabled: whitening,
|
||||
sampleCount: n,
|
||||
trainingData: data
|
||||
)
|
||||
}
|
||||
|
||||
public func incrementalUpdate(newSamples: [[Float]]) throws -> PCA {
|
||||
let combined = trainingData + newSamples
|
||||
return try PCA.train(data: combined, outputDimension: outputDimension, whitening: whiteningEnabled)
|
||||
}
|
||||
|
||||
public func partialFit(sample: [Float]) throws -> PCA {
|
||||
return try incrementalUpdate(newSamples: [sample])
|
||||
}
|
||||
}
|
||||
|
||||
public enum PCAError: Error, LocalizedError {
|
||||
case noData
|
||||
case invalidDimension
|
||||
case dimensionMismatch
|
||||
case svdFailed
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .noData: return "No data provided for PCA training"
|
||||
case .invalidDimension: return "Invalid output dimension"
|
||||
case .dimensionMismatch: return "Input dimension does not match model"
|
||||
case .svdFailed: return "SVD computation failed"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
import Foundation
|
||||
|
||||
public enum PoolingMethod: String, Codable, Sendable {
|
||||
case mean
|
||||
case last
|
||||
case cls
|
||||
case max
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import Foundation
|
||||
|
||||
public struct TextEmbeddingConfig: Codable, Sendable {
|
||||
public var modelType: String
|
||||
public var poolingMethod: PoolingMethod
|
||||
public var normalize: Bool
|
||||
|
||||
public init(
|
||||
modelType: String = "gemma-2b",
|
||||
poolingMethod: PoolingMethod = .mean,
|
||||
normalize: Bool = true
|
||||
) {
|
||||
self.modelType = modelType
|
||||
self.poolingMethod = poolingMethod
|
||||
self.normalize = normalize
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
import Foundation
|
||||
|
||||
public final class TextEmbeddingModel: @unchecked Sendable {
|
||||
private let model: E4BModel
|
||||
private let engine: MarkBaseEngine
|
||||
private let config: TextEmbeddingConfig
|
||||
private var pca: PCA?
|
||||
private let tokenizer: Tokenizer
|
||||
|
||||
public init(modelDir: String, engine: MarkBaseEngine, config: TextEmbeddingConfig) throws {
|
||||
self.engine = engine
|
||||
self.config = config
|
||||
self.model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||||
}
|
||||
|
||||
public func embed(text: String) throws -> [Float] {
|
||||
let tokens = tokenizer.encode(text: text)
|
||||
guard !tokens.isEmpty else { return [] }
|
||||
|
||||
model.kvCaches.forEach { $0.reset() }
|
||||
|
||||
let hiddenSize = model.hiddenSize
|
||||
var allHiddenStates: [[Float]] = []
|
||||
|
||||
for (pos, tokenId) in tokens.enumerated() {
|
||||
_ = try model.forward(tokenId: tokenId, position: pos, debug: false)
|
||||
let hs = engine.readFloats(from: model.temps.io, count: hiddenSize)
|
||||
allHiddenStates.append(hs)
|
||||
}
|
||||
|
||||
var result = pool(allHiddenStates)
|
||||
if config.normalize {
|
||||
let norm = sqrt(result.reduce(0) { $0 + $1 * $1 })
|
||||
if norm > 0 {
|
||||
for i in 0..<result.count {
|
||||
result[i] /= norm
|
||||
}
|
||||
}
|
||||
}
|
||||
if let pca = pca {
|
||||
result = try pca.transform(result)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
public func embedBatch(texts: [String]) throws -> [[Float]] {
|
||||
try texts.map { try embed(text: $0) }
|
||||
}
|
||||
|
||||
public func trainPCA(texts: [String], outputDimension: Int, whitening: Bool = false) throws {
|
||||
let embeddings = try embedBatch(texts: texts)
|
||||
pca = try PCA.train(data: embeddings, outputDimension: outputDimension, whitening: whitening)
|
||||
}
|
||||
|
||||
public func savePCA(to url: URL) throws {
|
||||
guard let pca = pca else { throw PCAError.noData }
|
||||
try pca.save(to: url)
|
||||
}
|
||||
|
||||
public func loadPCA(from url: URL) throws {
|
||||
pca = try PCA.load(from: url)
|
||||
}
|
||||
|
||||
private func pool(_ states: [[Float]]) -> [Float] {
|
||||
guard !states.isEmpty else { return [] }
|
||||
switch config.poolingMethod {
|
||||
case .last:
|
||||
return states.last ?? states[0]
|
||||
case .cls:
|
||||
return states[0]
|
||||
case .mean:
|
||||
let count = states.count
|
||||
let dim = states[0].count
|
||||
var result = [Float](repeating: 0, count: dim)
|
||||
for i in 0..<count {
|
||||
for j in 0..<dim {
|
||||
result[j] += states[i][j]
|
||||
}
|
||||
}
|
||||
for j in 0..<dim {
|
||||
result[j] /= Float(count)
|
||||
}
|
||||
return result
|
||||
case .max:
|
||||
let count = states.count
|
||||
let dim = states[0].count
|
||||
var result = states[0]
|
||||
for i in 1..<count {
|
||||
for j in 0..<dim {
|
||||
result[j] = max(result[j], states[i][j])
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user