import Foundation
import MarkBase

print("=== 测试 Gemma-4 26B 模型加载 ===")

let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-26b-a4b-mxfp4"

print("\n步骤 1: 检查文件")
let fm = FileManager.default

let files = [
    "config.json",
    "tokenizer.json",
    "model-00001-of-00003.safetensors",
    "model-00002-of-00003.safetensors",
    "model-00003-of-00003.safetensors",
    "model.safetensors.index.json"
]

for file in files {
    let path = modelDir + "/" + file
    if fm.fileExists(atPath: path) {
        let size = try! fm.attributesOfItem(atPath: path)[.size] as! Int64
        print("  ✓ \(file): \(size / 1024 / 1024) MB")
    } else {
        print("  ✗ \(file): NOT FOUND")
    }
}

print("\n步骤 2: 加载 Engine")
do {
    let engine = try E4BEngine(autoCompile: true)
    print("  ✓ Engine created")
    
    print("\n步骤 3: 加载 Model (26B)")
    print("  警告: 这可能需要几分钟...")
    
    // Try loading with smaller context to save memory
    let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
    
    print("  ✓ Model loaded!")
    print("  Layers: \(model.numHiddenLayers)")
    print("  Hidden size: \(model.hiddenSize)")
    print("  Vocab size: \(model.vocabSize)")
    
    print("\n步骤 4: 加载 Tokenizer")
    let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
    print("  ✓ Tokenizer loaded")
    
    print("\n步骤 5: 测试推理")
    let tokens = tokenizer.encode(text: "Hello")
    print("  Test tokens: \(tokens)")
    
    print("\n=== 26B 模型加载成功 ===")
    
} catch {
    print("\n✗ 错误: \(error)")
    print("\n可能的错误原因:")
    print("  1. 权重命名不匹配 (language_model.model.layers vs layers)")
    print("  2. MoE 结构不支持 (需要专家路由)")
    print("  3. Memory 不足 (26B 需要 ~16GB)")
    print("  4. 配置格式不兼容")
    
    print("\n建议:")
    print("  1. 检查权重命名并适配")
    print("  2. 实现 MoE 支持")
    print("  3. 转换为我们的标准格式")
}