markbaseengine/Tests/MarkBaseTests/NonProgrammingTest.swift

import XCTest
@testable import MarkBase

class NonProgrammingTest: XCTestCase {
    
    func testTextUnderstanding() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Text Understanding Test")
        print("═════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        
        let prompt = textPrompts[0]
        print("Prompt: \(prompt.prompt)")
        
        let tokens = tokenizer.encode(text: prompt.prompt)
        var generatedTokens: [Int] = []
        
        for i in 0..<50 {
            let pos = tokens.count + i - 1
            let logits = try model.forwardOptimized(tokenId: tokens.last ?? 0, position: pos)
            let nextToken = logits.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
            generatedTokens.append(nextToken)
            
            if tokenizer.eosTokenIds.contains(nextToken) { break }
        }
        
        let response = tokenizer.decode(tokens: generatedTokens)
        print("Response: \(response)")
        
        let hasKeywords = prompt.expectedKeywords.contains { keyword in
            response.lowercased().contains(keyword.lowercased())
        }
        
        if hasKeywords {
            print("✅ PASS")
        } else {
            print("⚠ FAIL - Missing expected keywords")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
    
    func testMathReasoning() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Math Reasoning Test")
        print("═════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        
        let prompt = mathPrompts[0]
        print("Prompt: \(prompt.prompt)")
        
        let tokens = tokenizer.encode(text: prompt.prompt)
        var generatedTokens: [Int] = []
        
        for i in 0..<50 {
            let pos = tokens.count + i - 1
            let logits = try model.forwardOptimized(tokenId: tokens.last ?? 0, position: pos)
            let nextToken = logits.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
            generatedTokens.append(nextToken)
            
            if tokenizer.eosTokenIds.contains(nextToken) { break }
        }
        
        let response = tokenizer.decode(tokens: generatedTokens)
        print("Response: \(response)")
        
        let hasKeywords = prompt.expectedKeywords.contains { keyword in
            response.contains(keyword)
        }
        
        if hasKeywords {
            print("✅ PASS")
        } else {
            print("⚠ FAIL - Incorrect math answer")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
}