markbaseengine/Tests/MarkBaseTests/VisionTowerIndependentTest.swift

import XCTest
@testable import MarkBase

final class VisionTowerIndependentTest: XCTestCase {
    
    func test12BVisionIndependent() throws {
        print("\n═══════════════════════════════════════")
        print("  12B Vision INDEPENDENT Test (No TEXT)")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load VisionTower12B ONLY (no TEXT model)...")
        let loadStart = Date()
        let visionTower = try VisionTower12B.load(modelDir: modelDir, engine: engine)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Vision tower loaded in \(loadTime) ms")
        print("  Config: hiddenDim=\(visionTower.hiddenDim), patchDim=\(visionTower.patchDim)")
        print("  Output: \(visionTower.outputDim)")
        
        print("\nStep 3: Generate synthetic patch embeddings...")
        let numPatches = 256
        let patchDim = visionTower.patchDim
        var patches: [Float] = []
        for i in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Patch shape: [\(numPatches), \(patchDim)]")
        
        print("\nStep 4: Process vision...")
        let patchBuffer = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: numPatches * visionTower.outputDim * 4)!
        
        let forwardStart = Date()
        try visionTower.forward(patchEmbeddings: patchBuffer, numPatches: numPatches, outputBuffer: outputBuffer)
        let forwardTime = Date().timeIntervalSince(forwardStart) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: numPatches * visionTower.outputDim))
        
        print("  Output shape: [\(numPatches), \(visionTower.outputDim)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        print("  Mean: \(output.reduce(0, +) / Float(output.count))")
        print("  Time: \(forwardTime) ms")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "Vision output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ 12B Vision INDEPENDENT test passed")
        print("  Load time: \(loadTime) ms (vs 3-4 min for TEXT model)")
        print("  Forward time: \(forwardTime) ms")
        print("═══════════════════════════════════════\n")
    }
    
    func testE4BVisionIndependent() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Vision INDEPENDENT Test (No TEXT)")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load VisionTower (E4B Full) ONLY (no TEXT model)...")
        let loadStart = Date()
        let vcfg = loadVisionConfig(modelDir: modelDir)
        guard let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") else {
            throw NSError(domain: "Vision", code: -1, userInfo: [NSLocalizedDescriptionKey: "Failed to create safetensors reader"])
        }
        let visionTower = try loadVisionTower(reader: reader, config: vcfg, engine: engine)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Vision tower loaded in \(loadTime) ms")
        print("  Config: hiddenSize=\(visionTower.config.hiddenSize), layers=\(visionTower.config.numHiddenLayers)")
        
        print("\nStep 3: Generate synthetic patch embeddings...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for i in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Patch shape: [\(numPatches), \(patchDim)]")
        
        print("\nStep 4: Process vision...")
        let patchBuffer = engine.device.makeBuffer(bytes: patches, length: patches.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: numPatches * visionTower.config.outputProjDims * 4)!
        
        let forwardStart = Date()
        try visionTower.forward(patchEmbeddings: patchBuffer, numPatches: numPatches, outputBuffer: outputBuffer)
        let forwardTime = Date().timeIntervalSince(forwardStart) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: numPatches * visionTower.config.outputProjDims))
        
        print("  Output shape: [\(numPatches), \(visionTower.config.outputProjDims)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        print("  Mean: \(output.reduce(0, +) / Float(output.count))")
        print("  Time: \(forwardTime) ms")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "Vision output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ E4B Vision INDEPENDENT test passed")
        print("  Load time: \(loadTime) ms")
        print("  Forward time: \(forwardTime) ms")
        print("═══════════════════════════════════════\n")
    }
}