markbaseengine/Sources/MarkBase/Embedding/TextEmbeddingModel.swift

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
        }
    }
}