markbaseengine/Tests/MarkBaseTests/E4BvsE2BComparisonTest.swift

import XCTest
@testable import MarkBase

class E4BvsE2BComparisonTest: XCTestCase {
    
    func testFullComparison() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  E4B-MarkBase vs E2B Full Comparison Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        // E4B Model Path
        let e4bPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let e2bPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
        
        guard FileManager.default.fileExists(atPath: e4bPath),
              FileManager.default.fileExists(atPath: e2bPath) else {
            print("⚠ Models not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // ===== 1. Architecture Comparison =====
        print("1. Architecture Comparison:")
        
        // E4B
        print("\n  E4B-MarkBase:")
        let e4bModel = try E4BModel(modelDir: e4bPath, engine: engine, maxContextLength: 128)
        print("    TEXT Layers: \(e4bModel.numHiddenLayers)")
        print("    Hidden Size: \(e4bModel.hiddenSize)")
        print("    Vocab Size: \(e4bModel.vocabSize)")
        
        // Check multimodal components
        let e4bReader = try SafeTensorsReader(path: "\(e4bPath)/model.safetensors")
        let e4bTensors = e4bReader.allTensors
        let e4bAudioTensors = e4bTensors.filter { $0.name.contains("audio_tower") }
        let e4bVisionTensors = e4bTensors.filter { $0.name.contains("vision_tower") }
        print("    Audio Tower Tensors: \(e4bAudioTensors.count)")
        print("    Vision Tower Tensors: \(e4bVisionTensors.count)")
        
        // E2B (sharded)
        print("\n  E2B:")
        let e2bModel = try E4BModel(modelDir: e2bPath, engine: engine, maxContextLength: 128)
        print("    TEXT Layers: \(e2bModel.numHiddenLayers)")
        print("    Hidden Size: \(e2bModel.hiddenSize)")
        print("    Vocab Size: \(e2bModel.vocabSize)")
        
        // E2B has index file
        let e2bIndex = try SafeTensorsIndex(modelDir: e2bPath)
        var e2bReaders: [String: SafeTensorsReader] = [:]
        for shardFile in Set(e2bIndex.weightMap.values) {
            e2bReaders[shardFile] = try SafeTensorsReader(path: "\(e2bPath)/\(shardFile)")
        }
        let e2bTensors = e2bReaders.values.flatMap { $0.allTensors }
        let e2bPerLayerTensors = e2bTensors.filter { $0.name.contains("per_layer") }
        print("    Per-layer Tensors: \(e2bPerLayerTensors.count)")
        
        // ===== 2. TEXT Performance Comparison =====
        print("\n2. TEXT Performance Comparison:")
        
        // Warmup
        _ = try e4bModel.forwardOptimized(tokenId: 2, position: 0)
        _ = try e2bModel.forwardOptimized(tokenId: 2, position: 0)
        
        // E4B TEXT Speed
        print("\n  E4B TEXT Speed (10 tokens):")
        let e4bStart = Date()
        var token = 2
        for i in 0..<10 {
            let result = try e4bModel.forwardOptimized(tokenId: token, position: i)
            token = result.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
        }
        let e4bTime = Date().timeIntervalSince(e4bStart) * 1000 / 10.0
        print("    Latency: \(String(format: "%.1f", e4bTime))ms/token")
        print("    Throughput: \(String(format: "%.1f", 1000.0 / e4bTime)) tok/s")
        
        // E2B TEXT Speed
        print("\n  E2B TEXT Speed (10 tokens):")
        let e2bStart = Date()
        token = 2
        for i in 0..<10 {
            let result = try e2bModel.forwardOptimized(tokenId: token, position: i)
            token = result.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
        }
        let e2bTime = Date().timeIntervalSince(e2bStart) * 1000 / 10.0
        print("    Latency: \(String(format: "%.1f", e2bTime))ms/token")
        print("    Throughput: \(String(format: "%.1f", 1000.0 / e2bTime)) tok/s")
        
        // ===== 3. NaN Stability Comparison =====
        print("\n3. NaN Stability Comparison (tokenIds 0-10):")
        
        // E4B NaN Test
        var e4bNaNCount = 0
        for tokenId in 0..<10 {
            let result = try e4bModel.forwardOptimized(tokenId: tokenId, position: 0)
            e4bNaNCount += result.filter { $0.isNaN }.count
        }
        print("  E4B NaN total: \(e4bNaNCount)")
        
        // E2B NaN Test
        var e2bNaNCount = 0
        for tokenId in 0..<10 {
            let result = try e2bModel.forwardOptimized(tokenId: tokenId, position: 0)
            e2bNaNCount += result.filter { $0.isNaN }.count
        }
        print("  E2B NaN total: \(e2bNaNCount)")
        
        // ===== 4. Scales Quality Comparison =====
        print("\n4. Scales Quality Comparison:")
        
        // E4B Scales
        let e4bScales = e4bTensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = e4bScales {
            let data = try e4bReader.read(tensor: s)
            let scales = data.withUnsafeBytes { ptr in
                Array(ptr.assumingMemoryBound(to: Float.self).prefix(20))
            }
            let negCount = scales.filter { $0 < 0 }.count
            print("  E4B Scales shape: \(s.shape)")
            print("  E4B Negative scales: \(negCount)")
        }
        
        // E2B Scales
        let e2bScales = e2bTensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = e2bScales {
            let shardFile = e2bIndex.weightMap[s.name] ?? "model-00001-of-00002.safetensors"
            let reader = e2bReaders[shardFile]!
            let data = try reader.read(tensor: s)
            let scales = data.withUnsafeBytes { ptr in
                Array(ptr.assumingMemoryBound(to: Float.self).prefix(20))
            }
            let negCount = scales.filter { $0 < 0 }.count
            print("  E2B Scales shape: \(s.shape)")
            print("  E2B Negative scales: \(negCount)")
        }
        
        // ===== 5. Long Context Comparison =====
        print("\n5. Long Context Performance (position 500-510):")
        
        // E4B Long Context
        let e4bLongStart = Date()
        for i in 500..<510 {
            _ = try e4bModel.forwardOptimized(tokenId: 2, position: i)
        }
        let e4bLongTime = Date().timeIntervalSince(e4bLongStart) * 1000 / 10.0
        let e4bDegradation = ((e4bLongTime - e4bTime) / e4bTime) * 100.0
        print("  E4B Position 500: \(String(format: "%.1f", e4bLongTime))ms, degradation: \(String(format: "%.1f", e4bDegradation))%")
        
        // E2B Long Context
        let e2bLongStart = Date()
        for i in 500..<510 {
            _ = try e2bModel.forwardOptimized(tokenId: 2, position: i)
        }
        let e2bLongTime = Date().timeIntervalSince(e2bLongStart) * 1000 / 10.0
        let e2bDegradation = ((e2bLongTime - e2bTime) / e2bTime) * 100.0
        print("  E2B Position 500: \(String(format: "%.1f", e2bLongTime))ms, degradation: \(String(format: "%.1f", e2bDegradation))%")
        
        // ===== 6. Feature Comparison =====
        print("\n6. Feature Comparison:")
        print("  E4B Features:")
        print("    ✓ TEXT inference")
        print("    ✓ Audio processing (\(e4bAudioTensors.count) tensors)")
        print("    ✓ Vision processing (\(e4bVisionTensors.count) tensors)")
        print("    ✓ Multimodal generation")
        
        print("\n  E2B Features:")
        print("    ✓ TEXT inference")
        print("    ✓ Per-layer embeddings (\(e2bPerLayerTensors.count) tensors)")
        print("    ✗ No audio tower")
        print("    ✗ No vision tower")
        
        // ===== 7. Summary =====
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Comparison Summary")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        print("TEXT Performance:")
        print("  E4B: \(String(format: "%.1f", e4bTime))ms, \(String(format: "%.1f", 1000.0/e4bTime)) tok/s")
        print("  E2B: \(String(format: "%.1f", e2bTime))ms, \(String(format: "%.1f", 1000.0/e2bTime)) tok/s")
        print("  Winner: \(e4bTime < e2bTime ? "E4B" : "E2B") (by \(String(format: "%.1f", abs(e4bTime - e2bTime)))ms)")
        
        print("\nNaN Stability:")
        print("  E4B: \(e4bNaNCount) NaN")
        print("  E2B: \(e2bNaNCount) NaN")
        print("  Winner: \(e4bNaNCount == 0 && e2bNaNCount == 0 ? "Both perfect" : (e4bNaNCount < e2bNaNCount ? "E4B" : "E2B"))")
        
        print("\nKV Cache Efficiency:")
        print("  E4B: \(String(format: "%.1f", e4bDegradation))% degradation")
        print("  E2B: \(String(format: "%.1f", e2bDegradation))% degradation")
        print("  Winner: \(abs(e4bDegradation) < abs(e2bDegradation) ? "E4B" : "E2B")")
        
        print("\nMultimodal:")
        print("  E4B: ✓ Full Audio+Vision+Text")
        print("  E2B: ✗ TEXT only")
        print("  Winner: E4B (multimodal support)")
        
        print("\nUse Case Recommendation:")
        print("  TEXT only: \(e2bTime < e4bTime ? "E2B" : "E4B")")
        print("  Multimodal: E4B")
        print("  Per-layer feature: E2B")
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}