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
2026-07-05 13:29:25 +08:00
commit 8a66b9086a
90 changed files with 22252 additions and 0 deletions
@@ -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
+        }
+    }
+}