MarkBase Admin
81 lines
3.3 KiB
Swift
81 lines
3.3 KiB
Swift
import XCTest
|
|
@testable import MarkBase
|
|
|
|
final class EmbeddingTest: XCTestCase {
|
|
var engine: MarkBaseEngine!
|
|
var model: E4BModel!
|
|
var embeddingModel: TextEmbeddingModel!
|
|
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
|
|
|
override func setUp() {
|
|
super.setUp()
|
|
guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
|
|
return
|
|
}
|
|
engine = try? MarkBaseEngine(autoCompile: true)
|
|
model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
|
embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig())
|
|
}
|
|
|
|
func testEmbeddingDimension() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let embedding = try embeddingModel.embed(text: "Hello world")
|
|
XCTAssertEqual(embedding.count, 2560, "Embedding dimension should be 2560")
|
|
}
|
|
|
|
func testEmbeddingNormalized() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let embedding = try embeddingModel.embed(text: "Test text")
|
|
let norm = sqrt(embedding.reduce(0) { $0 + $1 * $1 })
|
|
XCTAssertEqual(norm, 1.0, accuracy: 0.001, "Embedding should be L2 normalized")
|
|
}
|
|
|
|
func testSimilarSentences() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let e1 = try embeddingModel.embed(text: "The cat is sitting on the mat")
|
|
let e2 = try embeddingModel.embed(text: "A cat rests on a rug")
|
|
let e3 = try embeddingModel.embed(text: "The stock market crashed today")
|
|
|
|
let sim12 = cosineSimilarity(e1, e2)
|
|
let sim13 = cosineSimilarity(e1, e3)
|
|
print("Similar(cat, cat): \(sim12)")
|
|
print("Similar(cat, stock): \(sim13)")
|
|
XCTAssertGreaterThan(sim12, sim13, "Similar sentences should have higher cosine similarity")
|
|
}
|
|
|
|
func testDifferentLengths() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let e1 = try embeddingModel.embed(text: "Hi")
|
|
let e2 = try embeddingModel.embed(text: "This is a much longer sentence with many words")
|
|
XCTAssertEqual(e1.count, e2.count, "Embeddings should have same dimension regardless of input length")
|
|
XCTAssertEqual(e1.count, 2560)
|
|
}
|
|
|
|
func testEmptyInput() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let embedding = try embeddingModel.embed(text: "")
|
|
XCTAssertEqual(embedding.count, 0, "Empty input should return empty embedding")
|
|
}
|
|
|
|
func testBatchEmbedding() throws {
|
|
try XCTSkipIf(embeddingModel == nil, "Model not found")
|
|
let texts = ["Hello", "World", "Test"]
|
|
let embeddings = try embeddingModel.embedBatch(texts: texts)
|
|
XCTAssertEqual(embeddings.count, 3)
|
|
for embedding in embeddings {
|
|
XCTAssertEqual(embedding.count, 2560)
|
|
}
|
|
}
|
|
|
|
private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
|
|
guard a.count == b.count, !a.isEmpty else { return 0 }
|
|
var dot: Float = 0, normA: Float = 0, normB: Float = 0
|
|
for i in 0..<a.count {
|
|
dot += a[i] * b[i]
|
|
normA += a[i] * a[i]
|
|
normB += b[i] * b[i]
|
|
}
|
|
return dot / (sqrt(normA) * sqrt(normB))
|
|
}
|
|
}
|