ac75faa0cc
Some checks failed
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
2086 lines
94 KiB
Swift
2086 lines
94 KiB
Swift
import XCTest
|
||
import CoreImage
|
||
@testable import MarkBase
|
||
|
||
final class E4BSimpleInferenceTest: XCTestCase {
|
||
|
||
func testKVCacheDebug() throws {
|
||
// Load model
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
print("Loading engine...")
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
print("Loading model...")
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine)
|
||
|
||
print("\n=== KV Cache Debug Test ===")
|
||
print("Running forward at position 0...")
|
||
|
||
// Process position 0 and 1 with debug output
|
||
let logits0 = try model.forward(tokenId: 2, position: 0, debug: true) // BOS
|
||
print("Position 0 logits: max=\(logits0.max() ?? 0), min=\(logits0.min() ?? 0)")
|
||
|
||
// Check KV cache for layer 0
|
||
let cache0 = model.kvCaches[0]
|
||
print("Layer 0 cache after pos0: currentLength=\(cache0.currentLength)")
|
||
let k0Ptr = cache0.buffer.contents().bindMemory(to: Float.self, capacity: 20)
|
||
let k0vals = [k0Ptr[0], k0Ptr[1], k0Ptr[2], k0Ptr[3], k0Ptr[4], k0Ptr[5], k0Ptr[6], k0Ptr[7], k0Ptr[8], k0Ptr[9]]
|
||
print("K[0, 0:10] = \(k0vals)")
|
||
|
||
print("---")
|
||
let logits1 = try model.forward(tokenId: 9259, position: 1, debug: true) // "Hello"
|
||
print("Position 1 logits: max=\(logits1.max() ?? 0), min=\(logits1.min() ?? 0)")
|
||
|
||
print("Layer 0 cache after pos1: currentLength=\(cache0.currentLength)")
|
||
let k1vals = [k0Ptr[0], k0Ptr[1], k0Ptr[2], k0Ptr[3], k0Ptr[4], k0Ptr[5], k0Ptr[6], k0Ptr[7], k0Ptr[8], k0Ptr[9]]
|
||
let k1posvals = [k0Ptr[256], k0Ptr[257], k0Ptr[258], k0Ptr[259], k0Ptr[260], k0Ptr[261], k0Ptr[262], k0Ptr[263], k0Ptr[264], k0Ptr[265]]
|
||
print("K[0, 0:10] = \(k1vals)")
|
||
print("K[1, 0:10] = \(k1posvals)")
|
||
|
||
// Check shared layer's source
|
||
if let src = model.kvSourceMap[24] {
|
||
print("Layer 24 (shared) reads from layer \(src)")
|
||
let cacheSrc = model.kvCaches[src]
|
||
print("Layer \(src) cache: currentLength=\(cacheSrc.currentLength)")
|
||
print("Layer 24 cache: currentLength=\(model.kvCaches[24].currentLength)")
|
||
}
|
||
|
||
// Check all kvSourceMap entries
|
||
print("\nKV source map (shared layers):")
|
||
for layer in 24..<42 {
|
||
let src = model.kvSourceMap[layer] ?? -1
|
||
print(" Layer \(layer) → reads from layer \(src)")
|
||
}
|
||
|
||
print("---")
|
||
let _ = try model.forward(tokenId: 22470, position: 2, debug: true) // "Character"
|
||
}
|
||
|
||
func testLongerPrompt() throws {
|
||
Swift.print("\n=== Longer Prompt Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
let sampler = Sampler()
|
||
|
||
// Process 10 positions (BOS + 9 random tokens)
|
||
var tokens: [Int] = [2] // BOS
|
||
for i in 0..<9 {
|
||
let tokenId = 1000 + i // Some random tokens
|
||
tokens.append(tokenId)
|
||
}
|
||
|
||
Swift.print("Processing \(tokens.count) tokens...")
|
||
for (i, token) in tokens.enumerated() {
|
||
let _ = try model.forward(tokenId: token, position: i)
|
||
}
|
||
|
||
// Check cache lengths after processing
|
||
Swift.print("\nCache lengths after processing 10 tokens:")
|
||
for i in [0, 5, 11, 17, 23] { // Full attention owners
|
||
let cache = model.kvCaches[i]
|
||
Swift.print(" Layer \(i) (full) cache: currentLength=\(cache.currentLength)")
|
||
}
|
||
for i in [1, 10, 22] { // Sliding attention owners
|
||
let cache = model.kvCaches[i]
|
||
Swift.print(" Layer \(i) (sliding) cache: currentLength=\(cache.currentLength)")
|
||
}
|
||
|
||
// Get logits at position 9 (last position)
|
||
let logits = try model.forward(tokenId: tokens.last ?? 2, position: tokens.count - 1)
|
||
Swift.print("\nTop 10 logits at position \(tokens.count - 1):")
|
||
let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
for (idx, val) in indexed {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
}
|
||
|
||
func testSimpleTokenGeneration() throws {
|
||
print("\n=== Simple Token Generation Test ===")
|
||
print("Testing E4B model with proper sampling\n")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
|
||
// Load engine and model
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
|
||
// Load tokenizer
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
// Create sampler
|
||
let sampler = Sampler()
|
||
|
||
print("Model loaded, vocab size: \(tokenizer.vocabSize)")
|
||
print("Sampler settings: temperature=0.7, topK=50, topP=0.95")
|
||
|
||
var lastLogits: [Float]? = nil
|
||
|
||
// Test 1: Simple BOS → generation
|
||
print("\n--- Test 1: BOS only ---")
|
||
var tokens: [Int] = [2] // BOS
|
||
// Process BOS and capture logits
|
||
lastLogits = try model.forward(tokenId: 2, position: 0)
|
||
// Sample first generated token from logits at position 0
|
||
if let logits = lastLogits {
|
||
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(sampled)
|
||
}
|
||
// Generate continuation
|
||
for _ in 1..<15 {
|
||
guard let lastToken = tokens.last else { break }
|
||
if lastToken == 1 || lastToken == 106 { break } // EOS
|
||
let logits = try model.forward(tokenId: lastToken, position: tokens.count - 1)
|
||
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(sampled)
|
||
if sampled == 1 || sampled == 106 { break }
|
||
}
|
||
print("Generated: '\(tokenizer.decode(tokens: tokens))'")
|
||
|
||
// Test 2: Hello prompt
|
||
print("\n--- Test 2: 'Hello' prompt ---")
|
||
let model2 = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
tokens = [2, 9259] // BOS + "Hello"
|
||
// Process prompt and capture logits
|
||
lastLogits = nil
|
||
for (i, token) in tokens.enumerated() {
|
||
lastLogits = try model2.forward(tokenId: token, position: i)
|
||
}
|
||
// Sample first generated token from logits at position P-1
|
||
if let logits = lastLogits {
|
||
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(sampled)
|
||
}
|
||
// Generate continuation
|
||
for _ in 1..<20 {
|
||
guard let lastToken = tokens.last else { break }
|
||
if lastToken == 1 || lastToken == 106 { break }
|
||
let logits = try model2.forward(tokenId: lastToken, position: tokens.count - 1)
|
||
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(sampled)
|
||
if sampled == 1 || sampled == 106 { break }
|
||
}
|
||
print("Generated: '\(tokenizer.decode(tokens: tokens))'")
|
||
|
||
// Test 3: Full prompt format
|
||
print("\n--- Test 3: Full prompt ---")
|
||
let model3 = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let prompt = "<start_of_turn>user\nWhat is the capital of France?<end_of_turn>\n<start_of_turn>model\n"
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
print("Prompt: \(promptTokens.count) tokens")
|
||
print("Prompt text: '\(prompt)'")
|
||
// Process prompt
|
||
lastLogits = nil
|
||
for (i, token) in promptTokens.enumerated() {
|
||
lastLogits = try model3.forward(tokenId: token, position: i)
|
||
}
|
||
// Generate continuation
|
||
tokens = promptTokens
|
||
// Sample first generated token from logits at position P-1
|
||
if let logits = lastLogits {
|
||
print("Top 10 logits at position \(promptTokens.count - 1):")
|
||
let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
for (idx, val) in indexed {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
}
|
||
}
|
||
|
||
func testShortPrompt() throws {
|
||
Swift.print("\n=== Short Prompt Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
// Simple prompt: "Hello"
|
||
let prompt = "Hello"
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
Swift.print("Prompt: \(promptTokens.count) tokens, tokens: \(promptTokens)")
|
||
|
||
// Process prompt
|
||
for (i, token) in promptTokens.enumerated() {
|
||
let _ = try model.forward(tokenId: token, position: i)
|
||
}
|
||
|
||
// Get logits at last position
|
||
let logits = try model.forward(tokenId: promptTokens.last ?? 2, position: promptTokens.count - 1)
|
||
Swift.print("\nTop 10 logits at position \(promptTokens.count - 1):")
|
||
let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
for (idx, val) in indexed {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Decode full output
|
||
var tokens = promptTokens
|
||
let sampler = Sampler()
|
||
|
||
Swift.print("\n=== Generation trace ===")
|
||
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(sampled)
|
||
let decoded0 = tokenizer.decode(tokens: [sampled])
|
||
Swift.print("Position \(tokens.count - 1): sampled token \(sampled) ('\(decoded0)')")
|
||
|
||
for _ in 1..<10 {
|
||
guard let lastToken = tokens.last else { break }
|
||
if lastToken == 1 || lastToken == 106 { break }
|
||
let pos = tokens.count - 1
|
||
let newLogits = try model.forward(tokenId: lastToken, position: pos)
|
||
|
||
// Show top 5 logits at each position
|
||
let top5 = newLogits.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Position \(pos): top logits:")
|
||
for (idx, val) in top5 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
let newSampled = sampler.sample(logits: newLogits, temperature: 0.7, topK: 50, topP: 0.95)
|
||
tokens.append(newSampled)
|
||
let decodedNew = tokenizer.decode(tokens: [newSampled])
|
||
Swift.print(" Sampled: token \(newSampled) ('\(decodedNew)')")
|
||
if newSampled == 1 || newSampled == 106 { break }
|
||
}
|
||
Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokens))'")
|
||
}
|
||
|
||
func testPosition1ForwardPass() throws {
|
||
Swift.print("\n=== Position 1 Forward Pass Verification ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// Process position 0 (BOS)
|
||
Swift.print("\nPosition 0: BOS token (2)")
|
||
let logits0 = try model.forward(tokenId: 2, position: 0)
|
||
let top5_0 = logits0.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Top 5 logits:")
|
||
for (idx, val) in top5_0 {
|
||
Swift.print(" Token \(idx): \(val)")
|
||
}
|
||
|
||
// Check KV cache state after position 0
|
||
Swift.print("\nKV cache state after position 0:")
|
||
for i in 0..<5 {
|
||
let cache = model.kvCaches[i]
|
||
Swift.print(" Layer \(i) cache: currentLength=\(cache.currentLength), maxLength=\(cache.maxLength)")
|
||
}
|
||
|
||
// Process position 1 (token 568 = '(')
|
||
Swift.print("\nPosition 1: Token 568 ('(')")
|
||
let logits1 = try model.forward(tokenId: 568, position: 1)
|
||
let top5_1 = logits1.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Top 5 logits:")
|
||
for (idx, val) in top5_1 {
|
||
Swift.print(" Token \(idx): \(val)")
|
||
}
|
||
|
||
// Check KV cache state after position 1
|
||
Swift.print("\nKV cache state after position 1:")
|
||
for i in 0..<5 {
|
||
let cache = model.kvCaches[i]
|
||
Swift.print(" Layer \(i) cache: currentLength=\(cache.currentLength), maxLength=\(cache.maxLength)")
|
||
}
|
||
|
||
// Process position 2 (greedy choice)
|
||
let greedy1 = logits1.enumerated().max(by: { $0.element < $1.element })!.offset
|
||
Swift.print("\nPosition 2: Token \(greedy1)")
|
||
let logits2 = try model.forward(tokenId: greedy1, position: 2)
|
||
let top5_2 = logits2.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Top 5 logits:")
|
||
for (idx, val) in top5_2 {
|
||
Swift.print(" Token \(idx): \(val)")
|
||
}
|
||
|
||
// Check if there's a repeating pattern in logits
|
||
Swift.print("\nComparing logits across positions:")
|
||
Swift.print(" Position 0 top token: \(top5_0.first!.offset)")
|
||
Swift.print(" Position 1 top token: \(top5_1.first!.offset)")
|
||
Swift.print(" Position 2 top token: \(top5_2.first!.offset)")
|
||
|
||
// Check magnitude of hidden states
|
||
Swift.print("\nHidden state magnitudes (checking for explosion):")
|
||
// We can check by looking at the logits values - if they're close to softcapping limit (30), that's normal
|
||
let maxLogit0 = logits0.max()!
|
||
let maxLogit1 = logits1.max()!
|
||
let maxLogit2 = logits2.max()!
|
||
Swift.print(" Position 0 max logit: \(maxLogit0) (softcapped to ~30)")
|
||
Swift.print(" Position 1 max logit: \(maxLogit1)")
|
||
Swift.print(" Position 2 max logit: \(maxLogit2)")
|
||
|
||
// Check if model enters repeating pattern
|
||
if top5_1.first!.offset == top5_2.first!.offset {
|
||
Swift.print("\n⚠️ WARNING: Same token predicted at position 1 and 2 - potential loop!")
|
||
}
|
||
|
||
// Check for abnormal patterns
|
||
let logitRange0 = logits0.max()! - logits0.min()!
|
||
let logitRange1 = logits1.max()! - logits1.min()!
|
||
let logitRange2 = logits2.max()! - logits2.min()!
|
||
Swift.print("\nLogit ranges:")
|
||
Swift.print(" Position 0: \(logitRange0)")
|
||
Swift.print(" Position 1: \(logitRange1)")
|
||
Swift.print(" Position 2: \(logitRange2)")
|
||
|
||
// Very narrow logit range could indicate degenerate state
|
||
if logitRange2 < 1.0 {
|
||
Swift.print("\n⚠️ WARNING: Very narrow logit range at position 2 - degenerate state!")
|
||
}
|
||
}
|
||
|
||
func testRealImageProperPreprocessing() throws {
|
||
Swift.print("\n=== Real Image with Proper Preprocessing ===")
|
||
|
||
let imagePath = "/Users/accusys/MarkBase/tests/images/test_image.png"
|
||
|
||
// Load and preprocess image using CoreImage
|
||
Swift.print("Loading image: \(imagePath)")
|
||
guard let image = CIImage(contentsOf: URL(fileURLWithPath: imagePath)) else {
|
||
Swift.print("ERROR: Cannot load image")
|
||
return
|
||
}
|
||
|
||
Swift.print(" Original image extent: \(image.extent)")
|
||
|
||
// Resize to 224x224 (standard ViT input size)
|
||
let targetSize = CGSize(width: 224, height: 224)
|
||
let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
|
||
resizeFilter.setValue(image, forKey: kCIInputImageKey)
|
||
resizeFilter.setValue(224.0 / image.extent.width, forKey: kCIInputScaleKey)
|
||
resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
|
||
|
||
guard let resizedImage = resizeFilter.outputImage else {
|
||
Swift.print("ERROR: Cannot resize image")
|
||
return
|
||
}
|
||
|
||
Swift.print(" Resized image: \(resizedImage.extent)")
|
||
|
||
// Convert to pixel data
|
||
let context = CIContext()
|
||
let bitmap = context.createCGImage(resizedImage, from: resizedImage.extent)!
|
||
|
||
// Extract RGB pixel values
|
||
let dataProvider = bitmap.dataProvider!
|
||
let pixelData = dataProvider.data!
|
||
let ptr = CFDataGetBytePtr(pixelData)!
|
||
let length = CFDataGetLength(pixelData)
|
||
|
||
Swift.print(" Pixel data length: \(length) bytes")
|
||
Swift.print(" Expected: \(224 * 224 * 4) bytes (RGBA)")
|
||
|
||
// Convert to normalized floats [0, 1]
|
||
// Extract RGB channels (skip alpha)
|
||
var normalizedPixels = [Float](repeating: 0, count: 224 * 224 * 3)
|
||
|
||
// Config says standardize=false, so use raw [0,1] pixels
|
||
for i in 0..<224*224 {
|
||
let offset = i * 4 // RGBA
|
||
let r = Float(ptr[offset]) / 255.0
|
||
let g = Float(ptr[offset + 1]) / 255.0
|
||
let b = Float(ptr[offset + 2]) / 255.0
|
||
|
||
normalizedPixels[i * 3] = r
|
||
normalizedPixels[i * 3 + 1] = g
|
||
normalizedPixels[i * 3 + 2] = b
|
||
}
|
||
|
||
Swift.print(" ✓ Extracted \(normalizedPixels.count) pixel values (standardize=false)")
|
||
Swift.print(" Sample: r=\(normalizedPixels[0]), g=\(normalizedPixels[1]), b=\(normalizedPixels[2])")
|
||
|
||
// Create patch embeddings
|
||
let patchSize = 16
|
||
let numPatchesPerRow = 224 / patchSize // 14
|
||
let numPatches = numPatchesPerRow * numPatchesPerRow // 196
|
||
let hiddenSize = 768
|
||
|
||
Swift.print("\nCreating \(numPatches) patch embeddings (16x16 patches):")
|
||
|
||
// Each patch: 16*16 pixels * 3 channels = 768 floats
|
||
// This matches hiddenSize of vision tower!
|
||
var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
|
||
|
||
for patchIdx in 0..<numPatches {
|
||
let patchRow = patchIdx / numPatchesPerRow
|
||
let patchCol = patchIdx % numPatchesPerRow
|
||
|
||
for y in 0..<patchSize {
|
||
for x in 0..<patchSize {
|
||
// Global pixel coordinates
|
||
let globalY = patchRow * patchSize + y
|
||
let globalX = patchCol * patchSize + x
|
||
let pixelIdx = globalY * 224 + globalX
|
||
|
||
// RGB values
|
||
let r = normalizedPixels[pixelIdx * 3]
|
||
let g = normalizedPixels[pixelIdx * 3 + 1]
|
||
let b = normalizedPixels[pixelIdx * 3 + 2]
|
||
|
||
// Embedding position
|
||
let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
|
||
patchEmbeddings[embedIdx] = r
|
||
patchEmbeddings[embedIdx + 1] = g
|
||
patchEmbeddings[embedIdx + 2] = b
|
||
}
|
||
}
|
||
}
|
||
|
||
Swift.print(" ✓ Created \(patchEmbeddings.count) patch embedding values")
|
||
Swift.print(" First patch sample: \(patchEmbeddings[0..<10]))")
|
||
|
||
// Load model and run inference
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("\nProcessing through VisionTower...")
|
||
|
||
// Create buffers
|
||
let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
|
||
|
||
if let tower = multimodalModel.visionTowerFull {
|
||
try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
Swift.print(" ✓ VisionTower forward complete")
|
||
|
||
// Pool embeddings
|
||
Swift.print("\nPooling vision features...")
|
||
let visionOutputPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
var pooledEmbedding = [Float](repeating: 0, count: 2560)
|
||
|
||
for i in 0..<2560 {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionOutputPtr[p * 2560 + i]
|
||
}
|
||
pooledEmbedding[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
let mag = sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 })
|
||
Swift.print(" ✓ Pooled embedding magnitude: \(mag)")
|
||
|
||
// Normalize to match text embedding magnitude (~5)
|
||
// Text embeddings have norm ~5, vision has norm ~2130
|
||
// Scale down by factor of ~400 to match
|
||
let targetMag: Float = 5.0
|
||
let scale = targetMag / mag
|
||
Swift.print(" Normalizing vision features: scaling by \(scale)")
|
||
|
||
for i in 0..<2560 {
|
||
pooledEmbedding[i] *= scale
|
||
}
|
||
|
||
let normalizedMag = sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 })
|
||
Swift.print(" ✓ Normalized magnitude: \(normalizedMag)")
|
||
|
||
// Generate response
|
||
let prompt = "Describe what you see in this image"
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
|
||
Swift.print("\nPrompt: '\(prompt)'")
|
||
Swift.print("Generating response...")
|
||
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooledEmbedding, length: pooledEmbedding.count * 4)!
|
||
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: promptTokens,
|
||
precomputedVisionEmbedding: pooledBuffer,
|
||
maxTokens: 50
|
||
)
|
||
|
||
let decoded = tokenizer.decode(tokens: generatedTokens)
|
||
Swift.print("\nFull output: '\(decoded)'")
|
||
|
||
// Extract response
|
||
let responseStart = promptTokens.count + 4 // Skip prompt + BOI + IMAGE + EOI
|
||
if generatedTokens.count > responseStart {
|
||
let responseTokens = Array(generatedTokens[responseStart...])
|
||
let response = tokenizer.decode(tokens: responseTokens)
|
||
Swift.print("\nResponse: '\(response)'")
|
||
|
||
// Check if response makes sense
|
||
Swift.print("\n=== Analysis ===")
|
||
Swift.print("Vision features magnitude: \(mag)")
|
||
Swift.print("Generated tokens count: \(responseTokens.count)")
|
||
Swift.print("Note: Response quality depends on:")
|
||
Swift.print(" - Proper vision preprocessing")
|
||
Swift.print(" - Learned patch projection (we used raw pixels)")
|
||
Swift.print(" - Model training quality")
|
||
}
|
||
} else {
|
||
Swift.print("ERROR: Vision tower not loaded")
|
||
}
|
||
}
|
||
|
||
func testRealImageInference() throws {
|
||
Swift.print("\n=== Real Image Multimodal Inference Test ===")
|
||
|
||
let imagePath = "/Users/accusys/MarkBase/tests/images/test_image.png"
|
||
Swift.print("Loading image: \(imagePath)")
|
||
|
||
// Check if image exists
|
||
guard FileManager.default.fileExists(atPath: imagePath) else {
|
||
Swift.print("ERROR: Test image not found at \(imagePath)")
|
||
return
|
||
}
|
||
|
||
// Load image and create patch embeddings
|
||
Swift.print("\nProcessing image into patch embeddings...")
|
||
|
||
// Simple patch embedding: resize to 224x224, split into 16x16 patches
|
||
// For Gemma-4 vision: patch_size=16, hidden_size=768
|
||
let patchSize = 16
|
||
let imageWidth = 224
|
||
let imageHeight = 224
|
||
let numPatches = (imageWidth / patchSize) * (imageHeight / patchSize) // 14*14 = 196
|
||
let hiddenSize = 768
|
||
|
||
Swift.print(" Image size: \(imageWidth)x\(imageHeight)")
|
||
Swift.print(" Patch size: \(patchSize)x\(patchSize)")
|
||
Swift.print(" Num patches: \(numPatches)")
|
||
Swift.print(" Hidden size: \(hiddenSize)")
|
||
|
||
// Create simple patch embeddings (normalized pixel values)
|
||
// In a real implementation, this would use a learned projection
|
||
// For testing, we'll use raw pixel patches flattened and normalized
|
||
var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
|
||
|
||
// Simulate patch extraction: each patch has 16*16*3 = 768 values
|
||
for patchIdx in 0..<numPatches {
|
||
let patchRow = patchIdx / 14
|
||
let patchCol = patchIdx % 14
|
||
|
||
// Each patch: 16x16 pixels x 3 channels = 768 floats
|
||
for y in 0..<patchSize {
|
||
for x in 0..<patchSize {
|
||
for c in 0..<3 {
|
||
let idx = patchIdx * hiddenSize + (y * patchSize + x) * 3 + c
|
||
// Simple pattern: use patch position as feature
|
||
patchEmbeddings[idx] = Float(patchRow * patchSize + y + patchCol * patchSize + x) / 224.0
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Swift.print(" ✓ Created \(numPatches) patch embeddings")
|
||
|
||
// Load multimodal model
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("\nVision tower: \(multimodalModel.visionTowerFull != nil ? "✓ loaded" : "✗ not loaded")")
|
||
|
||
// Generate with vision input
|
||
let prompt = "What do you see in this image?"
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
|
||
Swift.print("\nPrompt: '\(prompt)'")
|
||
Swift.print("Generating response...")
|
||
|
||
// Process through vision tower
|
||
let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * hiddenSize * 4)!
|
||
|
||
// Run VisionTower
|
||
Swift.print("Running vision tower forward pass...")
|
||
if let tower = multimodalModel.visionTowerFull {
|
||
try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
Swift.print(" ✓ Vision tower forward complete")
|
||
|
||
// Pool embeddings: mean of all 196 patches → single embedding
|
||
Swift.print("\nPooling vision embeddings...")
|
||
let visionOutputPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
var pooledEmbedding = [Float](repeating: 0, count: hiddenSize)
|
||
|
||
// Mean pooling
|
||
for i in 0..<hiddenSize {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionOutputPtr[p * hiddenSize + i]
|
||
}
|
||
pooledEmbedding[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
Swift.print(" ✓ Pooled 196 embeddings into 1")
|
||
Swift.print(" Pooled embedding magnitude: \(sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 }))")
|
||
|
||
// Create buffer for pooled embedding
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooledEmbedding, length: pooledEmbedding.count * 4)!
|
||
|
||
// Use pooled embedding with MultimodalInference
|
||
Swift.print("\nGenerating with pooled vision embedding...")
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: promptTokens,
|
||
precomputedVisionEmbedding: pooledBuffer,
|
||
maxTokens: 30
|
||
)
|
||
|
||
// Decode result
|
||
let decoded = tokenizer.decode(tokens: generatedTokens)
|
||
Swift.print("\nGenerated: '\(decoded)'")
|
||
|
||
// Extract response (skip prompt + BOI + IMAGE + EOI = 4 tokens)
|
||
let newTextStart = promptTokens.count + 4
|
||
if generatedTokens.count > newTextStart {
|
||
let newText = tokenizer.decode(tokens: Array(generatedTokens[newTextStart...]))
|
||
Swift.print("\nResponse: '\(newText)'")
|
||
}
|
||
} else {
|
||
Swift.print("ERROR: Vision tower not loaded")
|
||
}
|
||
}
|
||
|
||
func testMultimodalVisionInference() throws {
|
||
Swift.print("\n=== Multimodal Vision Inference Test ===")
|
||
|
||
// Check if MultimodalModel is available
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
|
||
Swift.print("\nLoading multimodal model...")
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("\nModel components:")
|
||
Swift.print(" ✓ Text model: \(multimodalModel.textModel.numHiddenLayers) layers")
|
||
Swift.print(" ✓ Vision tower full: \(multimodalModel.visionTowerFull != nil ? "loaded (\(multimodalModel.visionTowerFull!.config.numHiddenLayers) layers)" : "not loaded")")
|
||
Swift.print(" ✓ Vision tower 12B: \(multimodalModel.visionTower != nil ? "loaded" : "not loaded")")
|
||
Swift.print(" ✓ Audio tower full: \(multimodalModel.audioTowerFull != nil ? "loaded" : "not loaded")")
|
||
Swift.print(" ✓ Audio tower 12B: \(multimodalModel.audioTower != nil ? "loaded" : "not loaded")")
|
||
Swift.print(" ✓ BOI token: \(multimodalModel.boiTokenId)")
|
||
Swift.print(" ✓ IMAGE token: \(multimodalModel.imageTokenId)")
|
||
Swift.print(" ✓ EOI token: \(multimodalModel.eoiTokenId)")
|
||
|
||
// Create a simple test image (dummy patches)
|
||
// In reality, this would come from actual image processing
|
||
// For testing, we'll create synthetic patch embeddings
|
||
Swift.print("\nCreating synthetic vision input...")
|
||
|
||
let hiddenSize = multimodalModel.textModel.hiddenSize
|
||
let numPatches = 1 // Single patch for testing
|
||
|
||
// Create a simple synthetic embedding
|
||
// This simulates what the vision tower would output
|
||
var patchEmbedding = [Float](repeating: 0.01, count: hiddenSize)
|
||
|
||
// Add some variation to make it more realistic
|
||
for i in 0..<hiddenSize {
|
||
patchEmbedding[i] = Float.random(in: -0.1...0.1)
|
||
}
|
||
|
||
// Test prompt: "Describe this image"
|
||
let promptTokens = tokenizer.encode(text: "Describe this image")
|
||
Swift.print("Prompt: 'Describe this image'")
|
||
Swift.print("Tokens: \(promptTokens)")
|
||
|
||
// Generate with vision input
|
||
Swift.print("\nGenerating with vision conditioning...")
|
||
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: promptTokens,
|
||
imagePatches: patchEmbedding,
|
||
numImagePatches: numPatches,
|
||
maxTokens: 20
|
||
)
|
||
|
||
// Decode result
|
||
let decoded = tokenizer.decode(tokens: generatedTokens)
|
||
Swift.print("\nGenerated tokens: \(generatedTokens)")
|
||
Swift.print("Decoded: '\(decoded)'")
|
||
|
||
// Check if generation makes more sense than text-only
|
||
let newTextTokens = Array(generatedTokens.dropFirst(promptTokens.count + 3)) // Skip prompt + BOI/IMAGE/EOI
|
||
let newText = tokenizer.decode(tokens: newTextTokens)
|
||
Swift.print("\nNew text (after multimodal tokens): '\(newText)'")
|
||
|
||
Swift.print("\n=== Summary ===")
|
||
Swift.print("Multimodal inference pipeline executed successfully")
|
||
Swift.print("Vision conditioning was injected into text generation")
|
||
Swift.print("Note: Synthetic embeddings were used - real test needs actual image")
|
||
}
|
||
|
||
func testRealVisionPipeline() throws {
|
||
Swift.print("\n=== Real Vision Pipeline Test ===")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("\n✓ Multimodal model loaded")
|
||
|
||
// Create test image (red 224x224)
|
||
let imageData = """
|
||
iVBORw0KGgoAAAANSUhEUgAAAOAAAADgCAIAAACVT/22AAACaUlEQVR4nO3SMQEAIAzAsIF/z+CAlx6Jgh5dZ6Br/w6AF4OSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGZcouY7ECvxTQrjoAAAAASUVORK5CYII=
|
||
"""
|
||
|
||
let data = Data(base64Encoded: imageData)!
|
||
Swift.print("\n✓ Test image loaded (red 224x224): \(data.count) bytes")
|
||
|
||
// Preprocess image (create patch embeddings)
|
||
|
||
guard let ciImage = CIImage(data: data) else {
|
||
Swift.print("✗ Failed to create CIImage")
|
||
return
|
||
}
|
||
|
||
// Resize to 224x224
|
||
let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
|
||
resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
|
||
let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
|
||
resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
|
||
resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
|
||
|
||
guard let resized = resizeFilter.outputImage else {
|
||
Swift.print("✗ Failed to resize image")
|
||
return
|
||
}
|
||
|
||
// Convert to CGImage
|
||
let context = CIContext()
|
||
guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
|
||
Swift.print("✗ Failed to create CGImage")
|
||
return
|
||
}
|
||
|
||
// Extract RGB pixels
|
||
let dataProvider = cgImage.dataProvider!
|
||
let pixelData = dataProvider.data!
|
||
let ptr = CFDataGetBytePtr(pixelData)!
|
||
|
||
Swift.print("✓ First pixel RGB: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
|
||
|
||
// Create patch embeddings (16x16 patches)
|
||
let patchSize = 16
|
||
let numPatches = 14 * 14 // 196
|
||
let hiddenSize = 768
|
||
|
||
var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
|
||
|
||
for patchIdx in 0..<numPatches {
|
||
let patchRow = patchIdx / 14
|
||
let patchCol = patchIdx % 14
|
||
|
||
for y in 0..<patchSize {
|
||
for x in 0..<patchSize {
|
||
let globalY = patchRow * patchSize + y
|
||
let globalX = patchCol * patchSize + x
|
||
|
||
if globalY >= 224 || globalX >= 224 { continue }
|
||
|
||
let pixelIdx = globalY * 224 + globalX
|
||
let offset = pixelIdx * 4 // RGBA
|
||
|
||
let r = Float(ptr[offset]) / 255.0
|
||
let g = Float(ptr[offset + 1]) / 255.0
|
||
let b = Float(ptr[offset + 2]) / 255.0
|
||
|
||
let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
|
||
if embedIdx + 2 < patchEmbeddings.count {
|
||
patchEmbeddings[embedIdx] = r
|
||
patchEmbeddings[embedIdx + 1] = g
|
||
patchEmbeddings[embedIdx + 2] = b
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Swift.print("✓ Patch embeddings created: \(patchEmbeddings.count) floats")
|
||
|
||
// Vision tower forward pass
|
||
let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
|
||
|
||
guard let tower = multimodalModel.visionTowerFull else {
|
||
Swift.print("✗ Vision tower not available")
|
||
return
|
||
}
|
||
|
||
try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
Swift.print("✓ Vision tower forward pass complete")
|
||
|
||
// Pool embeddings (196 patches → 1)
|
||
let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
var pooled = [Float](repeating: 0, count: 2560)
|
||
|
||
for i in 0..<2560 {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionPtr[p * 2560 + i]
|
||
}
|
||
pooled[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
Swift.print("✓ Pooled vision embedding magnitude: \(mag)")
|
||
|
||
// Normalize to match text embeddings (magnitude ~5)
|
||
let scaleNorm: Float = 5.0 / mag
|
||
for i in 0..<2560 {
|
||
pooled[i] *= scaleNorm
|
||
}
|
||
|
||
let magAfterNorm = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
Swift.print("✓ Normalized magnitude: \(magAfterNorm)")
|
||
|
||
// Multimodal inference
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
|
||
|
||
let promptTokens = tokenizer.encode(text: "What color is this image?")
|
||
Swift.print("\nPrompt: 'What color is this image?'")
|
||
Swift.print("Tokens: \(promptTokens)")
|
||
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: promptTokens,
|
||
precomputedVisionEmbedding: pooledBuffer,
|
||
maxTokens: 20
|
||
)
|
||
|
||
let decoded = tokenizer.decode(tokens: generatedTokens)
|
||
Swift.print("\nGenerated: '\(decoded)'")
|
||
|
||
// Extract response (skip prompt + BOI + IMAGE + EOI)
|
||
let responseStart = promptTokens.count + 4
|
||
if generatedTokens.count > responseStart {
|
||
let responseTokens = Array(generatedTokens[responseStart...])
|
||
let response = tokenizer.decode(tokens: responseTokens)
|
||
Swift.print("Response: '\(response)'")
|
||
}
|
||
|
||
Swift.print("\n=== Real Vision Pipeline Test Complete ===")
|
||
}
|
||
|
||
func test26BModelLoading() throws {
|
||
Swift.print("\n=== 测试 Gemma-4 26B A4B 真正 4-bit 模型 ===")
|
||
|
||
// 真正的 4-bit A4B model (不是 MXFP4)
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
|
||
|
||
// Check if model exists
|
||
guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
|
||
Swift.print("✗ 转换后的26B模型未找到")
|
||
Swift.print(" 请运行: python3 convert_mlx_26b.py")
|
||
return
|
||
}
|
||
|
||
Swift.print("✓ 转换后的26B模型找到")
|
||
|
||
// Check files
|
||
let configFile = modelDir + "/config.json"
|
||
let tokenizerFile = modelDir + "/tokenizer.json"
|
||
let weightsFile1 = modelDir + "/model-00001-of-00003.safetensors"
|
||
|
||
for (file, desc) in [(configFile, "Config"), (tokenizerFile, "Tokenizer"), (weightsFile1, "Weights shard 1")] {
|
||
if FileManager.default.fileExists(atPath: file) {
|
||
let attrs = try FileManager.default.attributesOfItem(atPath: file)
|
||
let size = (attrs[.size] as? Int64 ?? 0) / 1024 / 1024
|
||
Swift.print(" ✓ \(desc): \(size) MB")
|
||
} else {
|
||
Swift.print(" ✗ \(desc): NOT FOUND")
|
||
return
|
||
}
|
||
}
|
||
|
||
// Check total weights size
|
||
let totalWeights: Int64 = try FileManager.default.contentsOfDirectory(atPath: modelDir)
|
||
.filter { $0.hasSuffix(".safetensors") }
|
||
.reduce(0) { sum, file in
|
||
let attrs = try FileManager.default.attributesOfItem(atPath: modelDir + "/" + file)
|
||
return sum + (attrs[.size] as? Int64 ?? 0)
|
||
}
|
||
Swift.print(" ✓ Total weights: \(totalWeights / 1024 / 1024 / 1024) GB")
|
||
|
||
Swift.print("\n步骤 2: 创建 Engine")
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
Swift.print(" ✓ Engine created")
|
||
|
||
Swift.print("\n步骤 3: 加载 Model")
|
||
Swift.print(" 26B 转换后模型 (~15GB):")
|
||
Swift.print(" - Hidden size: 2816")
|
||
Swift.print(" - Layers: 42")
|
||
Swift.print(" - Vocab: 262144")
|
||
Swift.print(" 警告: 可能需要1-2分钟...")
|
||
|
||
do {
|
||
// Use smaller context for testing
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
|
||
|
||
Swift.print(" ✓✓✓ Model loaded!")
|
||
Swift.print(" Layers: \(model.numHiddenLayers)")
|
||
Swift.print(" Hidden size: \(model.hiddenSize)")
|
||
Swift.print(" Vocab size: \(model.vocabSize)")
|
||
|
||
// Test tokenizer
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
Swift.print(" ✓ Tokenizer loaded")
|
||
|
||
// Test simple encoding
|
||
let tokens = tokenizer.encode(text: "Hello")
|
||
Swift.print(" Test tokens: \(tokens)")
|
||
|
||
Swift.print("\n=== 26B 测试成功 ===")
|
||
Swift.print("✓ 转换后的26B模型可以加载！")
|
||
Swift.print("✓ Hidden size: \(model.hiddenSize) (比12B的2560大)")
|
||
Swift.print("✓ 准备运行推理...")
|
||
|
||
// Try simple forward pass (optional)
|
||
Swift.print("\n尝试推理测试...")
|
||
do {
|
||
let logits = try model.forward(tokenId: tokens[0], position: 0)
|
||
Swift.print(" ✓ Forward pass成功!")
|
||
Swift.print(" Logits size: \(logits.count)")
|
||
|
||
// Check top predictions
|
||
let sorted = logits.enumerated().sorted { $0.element > $1.element }
|
||
let top5 = sorted.prefix(5)
|
||
Swift.print(" Top 5 predictions:")
|
||
for (idx, val) in top5 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
Swift.print("\n🎉 26B模型完全工作！")
|
||
|
||
// ── Quick generation test (3 tokens) ──
|
||
Swift.print("\n尝试简短的生成测试...")
|
||
do {
|
||
let genConfig = GenerationConfig(maxTokens: 3, temperature: 0.7)
|
||
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||
let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
|
||
Swift.print(" Generated: '\(response)'")
|
||
XCTAssertGreaterThan(response.count, 0, "Should generate text")
|
||
} catch {
|
||
Swift.print(" ⚠️ Generation test失败: \(error)")
|
||
}
|
||
|
||
} catch {
|
||
Swift.print(" ⚠️ Forward pass失败: \(error)")
|
||
Swift.print(" 可能需要MoE支持或额外适配")
|
||
}
|
||
|
||
} catch {
|
||
Swift.print("\n✗ 加载失败: \(error)")
|
||
|
||
let nsError = error as NSError
|
||
Swift.print("\n错误详情:")
|
||
Swift.print(" Domain: \(nsError.domain)")
|
||
Swift.print(" Code: \(nsError.code)")
|
||
Swift.print(" Description: \(nsError.localizedDescription)")
|
||
|
||
if let userInfo = nsError.userInfo as? [String: String] {
|
||
for (key, value) in userInfo {
|
||
Swift.print(" \(key): \(value)")
|
||
}
|
||
}
|
||
|
||
Swift.print("\n可能原因:")
|
||
Swift.print(" 1. 权重命名仍不完全匹配")
|
||
Swift.print(" 2. MoE结构需要实现")
|
||
Swift.print(" 3. scales转换有问题")
|
||
Swift.print(" 4. Memory不足 (需要~17GB)")
|
||
|
||
Swift.print("\n建议:")
|
||
Swift.print(" - 检查权重命名是否正确")
|
||
Swift.print(" - 检查Memory是否充足")
|
||
Swift.print(" - 查看详细错误日志")
|
||
}
|
||
}
|
||
|
||
func testNaturalImageInference() throws {
|
||
Swift.print("\n=== Natural Image Inference Test ===")
|
||
Swift.print("Testing with sky gradient + sun (more realistic)")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("✓ Model loaded")
|
||
|
||
guard let imageData = try? String(contentsOfFile: "/tmp/test_natural_image.txt") else {
|
||
Swift.print("✗ Natural image not found")
|
||
return
|
||
}
|
||
|
||
let data = Data(base64Encoded: imageData)!
|
||
Swift.print("✓ Natural image: \(data.count) bytes (sky + sun)")
|
||
|
||
let outputFile = "/tmp/natural_inference_output.txt"
|
||
var outputs: [String] = ["=== Natural Image Inference Test ==="]
|
||
|
||
// Preprocess
|
||
guard let ciImage = CIImage(data: data) else {
|
||
outputs.append("✗ Failed to create CIImage")
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
return
|
||
}
|
||
|
||
outputs.append("✓ Image size: \(Int(ciImage.extent.width))x\(Int(ciImage.extent.height))")
|
||
|
||
let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
|
||
resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
|
||
let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
|
||
resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
|
||
resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
|
||
|
||
guard let resized = resizeFilter.outputImage else { return }
|
||
|
||
let context = CIContext()
|
||
guard let cgImage = context.createCGImage(resized, from: resized.extent) else { return }
|
||
|
||
let ptr = CFDataGetBytePtr(cgImage.dataProvider!.data!)!
|
||
|
||
outputs.append("✓ Top-left pixel: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
|
||
outputs.append("✓ Sun center (~40,40): (\(ptr[(40*224+40)*4]), \(ptr[(40*224+40)*4+1]), \(ptr[(40*224+40)*4+2]))")
|
||
|
||
// Vision processing
|
||
let numPatches = 196
|
||
let hiddenSize = 768
|
||
var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
|
||
|
||
for patchIdx in 0..<numPatches {
|
||
let row = patchIdx / 14
|
||
let col = patchIdx % 14
|
||
|
||
for y in 0..<16 {
|
||
for x in 0..<16 {
|
||
let gy = row * 16 + y
|
||
let gx = col * 16 + x
|
||
|
||
if gy >= 224 || gx >= 224 { continue }
|
||
|
||
let offset = (gy * 224 + gx) * 4
|
||
let r = Float(ptr[offset]) / 255.0
|
||
let g = Float(ptr[offset + 1]) / 255.0
|
||
let b = Float(ptr[offset + 2]) / 255.0
|
||
|
||
let idx = patchIdx * hiddenSize + (y * 16 + x) * 3
|
||
if idx + 2 < patchEmbeddings.count {
|
||
patchEmbeddings[idx] = r
|
||
patchEmbeddings[idx + 1] = g
|
||
patchEmbeddings[idx + 2] = b
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
outputs.append("✓ Patch embeddings: \(patchEmbeddings.count) floats")
|
||
|
||
// Vision tower
|
||
let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
|
||
try multimodalModel.visionTowerFull!.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
|
||
outputs.append("✓ Vision tower forward complete")
|
||
|
||
// Pool
|
||
let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
var pooled = [Float](repeating: 0, count: 2560)
|
||
|
||
for i in 0..<2560 {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionPtr[p * 2560 + i]
|
||
}
|
||
pooled[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
outputs.append("✓ Pooled magnitude: \(mag)")
|
||
|
||
let normScale: Float = 5.0 / mag
|
||
for i in 0..<2560 {
|
||
pooled[i] *= normScale
|
||
}
|
||
|
||
outputs.append("✓ Normalized magnitude: \(sqrt(pooled.reduce(0) { $0 + $1 * $1 }))")
|
||
|
||
// Test multiple prompts
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
|
||
|
||
// Prompt 1: Describe scene
|
||
let p1 = tokenizer.encode(text: "Describe what you see in this image")
|
||
outputs.append("\n[Test 1] Prompt: 'Describe what you see in this image'")
|
||
let g1 = try inference.generate(textTokens: p1, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
|
||
outputs.append("Generated: '\(tokenizer.decode(tokens: g1))'")
|
||
|
||
// Prompt 2: Colors
|
||
let p2 = tokenizer.encode(text: "What colors are present?")
|
||
outputs.append("\n[Test 2] Prompt: 'What colors are present?'")
|
||
let g2 = try inference.generate(textTokens: p2, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
|
||
outputs.append("Generated: '\(tokenizer.decode(tokens: g2))'")
|
||
|
||
// Prompt 3: Scene type
|
||
let p3 = tokenizer.encode(text: "Is this outdoor or indoor?")
|
||
outputs.append("\n[Test 3] Prompt: 'Is this outdoor or indoor?'")
|
||
let g3 = try inference.generate(textTokens: p3, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
|
||
outputs.append("Generated: '\(tokenizer.decode(tokens: g3))'")
|
||
|
||
outputs.append("\n=== Natural Image Test Complete ===")
|
||
outputs.append("Note: Compare outputs with model's expected behavior")
|
||
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
Swift.print("✓ Results: \(outputFile)")
|
||
|
||
XCTAssertTrue(FileManager.default.fileExists(atPath: outputFile))
|
||
}
|
||
|
||
func testGradientImageInference() throws {
|
||
Swift.print("\n=== Gradient Image Inference Test ===")
|
||
Swift.print("Testing with more complex gradient pattern (not pure red)")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
Swift.print("✓ Multimodal model loaded")
|
||
|
||
// Create gradient image (224x224)
|
||
guard let imageData = try? String(contentsOfFile: "/tmp/test_gradient_image.txt") else {
|
||
Swift.print("✗ Gradient image file not found, skipping test")
|
||
return
|
||
}
|
||
|
||
let data = Data(base64Encoded: imageData)!
|
||
Swift.print("✓ Gradient image loaded: \(data.count) bytes")
|
||
|
||
// Save outputs to file for verification
|
||
let outputFile = "/tmp/gradient_inference_output.txt"
|
||
var outputs: [String] = []
|
||
outputs.append("=== Gradient Image Inference Test ===")
|
||
|
||
// Preprocess
|
||
guard let ciImage = CIImage(data: data) else {
|
||
outputs.append("✗ Failed to create CIImage")
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
return
|
||
}
|
||
|
||
let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
|
||
resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
|
||
let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
|
||
resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
|
||
resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
|
||
|
||
guard let resized = resizeFilter.outputImage else {
|
||
outputs.append("✗ Failed to resize")
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
return
|
||
}
|
||
|
||
let context = CIContext()
|
||
guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
|
||
outputs.append("✗ Failed to create CGImage")
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
return
|
||
}
|
||
|
||
let dataProvider = cgImage.dataProvider!
|
||
let pixelData = dataProvider.data!
|
||
let ptr = CFDataGetBytePtr(pixelData)!
|
||
|
||
outputs.append("✓ First pixel RGB: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
|
||
|
||
// Create patch embeddings
|
||
let patchSize = 16
|
||
let numPatches = 14 * 14
|
||
let hiddenSize = 768
|
||
var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
|
||
|
||
for patchIdx in 0..<numPatches {
|
||
let patchRow = patchIdx / 14
|
||
let patchCol = patchIdx % 14
|
||
|
||
for y in 0..<patchSize {
|
||
for x in 0..<patchSize {
|
||
let globalY = patchRow * patchSize + y
|
||
let globalX = patchCol * patchSize + x
|
||
|
||
if globalY >= 224 || globalX >= 224 { continue }
|
||
|
||
let pixelIdx = globalY * 224 + globalX
|
||
let offset = pixelIdx * 4
|
||
|
||
let r = Float(ptr[offset]) / 255.0
|
||
let g = Float(ptr[offset + 1]) / 255.0
|
||
let b = Float(ptr[offset + 2]) / 255.0
|
||
|
||
let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
|
||
if embedIdx + 2 < patchEmbeddings.count {
|
||
patchEmbeddings[embedIdx] = r
|
||
patchEmbeddings[embedIdx + 1] = g
|
||
patchEmbeddings[embedIdx + 2] = b
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
outputs.append("✓ Patch embeddings created: \(patchEmbeddings.count) floats")
|
||
|
||
// Vision tower forward
|
||
let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
|
||
try multimodalModel.visionTowerFull!.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
|
||
outputs.append("✓ Vision tower forward complete")
|
||
|
||
// Pool and normalize
|
||
let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
var pooled = [Float](repeating: 0, count: 2560)
|
||
|
||
for i in 0..<2560 {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionPtr[p * 2560 + i]
|
||
}
|
||
pooled[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
outputs.append("✓ Pooled magnitude: \(mag)")
|
||
|
||
let scaleNorm: Float = 5.0 / mag
|
||
for i in 0..<2560 {
|
||
pooled[i] *= scaleNorm
|
||
}
|
||
|
||
let magAfter = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
outputs.append("✓ Normalized magnitude: \(magAfter)")
|
||
|
||
// Multimodal inference with different prompts
|
||
let inference = try MultimodalInference(model: multimodalModel)
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
|
||
|
||
// Test 1: Simple question
|
||
let prompt1 = tokenizer.encode(text: "What do you see?")
|
||
outputs.append("\n[Test 1] Prompt: 'What do you see?'")
|
||
let gen1 = try inference.generate(textTokens: prompt1, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
|
||
let dec1 = tokenizer.decode(tokens: gen1)
|
||
outputs.append("Generated: '\(dec1)'")
|
||
|
||
// Extract response
|
||
if gen1.count > prompt1.count + 4 {
|
||
let respTokens = Array(gen1[(prompt1.count + 4)...])
|
||
let resp = tokenizer.decode(tokens: respTokens)
|
||
outputs.append("Response: '\(resp)'")
|
||
}
|
||
|
||
// Test 2: Describe request
|
||
let prompt2 = tokenizer.encode(text: "Describe this image")
|
||
outputs.append("\n[Test 2] Prompt: 'Describe this image'")
|
||
let gen2 = try inference.generate(textTokens: prompt2, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
|
||
let dec2 = tokenizer.decode(tokens: gen2)
|
||
outputs.append("Generated: '\(dec2)'")
|
||
|
||
if gen2.count > prompt2.count + 4 {
|
||
let respTokens = Array(gen2[(prompt2.count + 4)...])
|
||
let resp = tokenizer.decode(tokens: respTokens)
|
||
outputs.append("Response: '\(resp)'")
|
||
}
|
||
|
||
// Test 3: Color question
|
||
let prompt3 = tokenizer.encode(text: "What colors are in this image?")
|
||
outputs.append("\n[Test 3] Prompt: 'What colors are in this image?'")
|
||
let gen3 = try inference.generate(textTokens: prompt3, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
|
||
let dec3 = tokenizer.decode(tokens: gen3)
|
||
outputs.append("Generated: '\(dec3)'")
|
||
|
||
if gen3.count > prompt3.count + 4 {
|
||
let respTokens = Array(gen3[(prompt3.count + 4)...])
|
||
let resp = tokenizer.decode(tokens: respTokens)
|
||
outputs.append("Response: '\(resp)'")
|
||
}
|
||
|
||
outputs.append("\n=== Gradient Image Test Complete ===")
|
||
outputs.append("Note: Output quality should be analyzed for coherence")
|
||
|
||
// Write outputs to file
|
||
try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
|
||
Swift.print("✓ Results saved to \(outputFile)")
|
||
|
||
// Verify file was created
|
||
XCTAssertTrue(FileManager.default.fileExists(atPath: outputFile), "Output file should be created")
|
||
}
|
||
|
||
func testMultimodalTextOnly() throws {
|
||
Swift.print("\n=== Multimodal Model Text-Only Test ===")
|
||
Swift.print("Note: E4B-MarkBase is Gemma4ForConditionalGeneration (multimodal)")
|
||
Swift.print("Text-only generation may not work properly without vision/audio conditioning")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
|
||
// Check if multimodal inference is available
|
||
Swift.print("\nChecking multimodal components:")
|
||
let mmInferenceAvailable = NSClassFromString("G12B.MultimodalInference") != nil
|
||
Swift.print(" MultimodalInference class: \(mmInferenceAvailable ? "available" : "NOT available")")
|
||
|
||
// Check if vision/audio weights are loaded
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
let sampler = Sampler()
|
||
|
||
Swift.print("\nModel components loaded:")
|
||
Swift.print(" ✓ 42 text layers")
|
||
Swift.print(" Vision/Audio towers should be separate components")
|
||
|
||
// The issue is: pure text generation without multimodal conditioning
|
||
// The model expects vision/audio features to be interleaved with text
|
||
Swift.print("\n=== Testing with multimodal placeholder tokens ===")
|
||
|
||
// Try using image_token or audio_token as conditioning
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// Test 1: Add image token before text (258880 = image_token)
|
||
Swift.print("\nTest 1: Prompt with image token")
|
||
let imageToken = 258880
|
||
var tokens = [2, imageToken] // BOS + IMAGE
|
||
|
||
Swift.print("Processing IMAGE token at position 0...")
|
||
// Note: This won't work properly because we need vision features
|
||
// But let's see what happens
|
||
|
||
// Generate a few tokens
|
||
for i in 0..<5 {
|
||
let pos = tokens.count - 1
|
||
let logits = try model.forward(tokenId: tokens.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print(" Position \(i): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Test 2: Check special token IDs
|
||
Swift.print("\n\nSpecial multimodal token IDs:")
|
||
let specialTokens = [
|
||
(258880, "image_token"),
|
||
(258881, "audio_token"),
|
||
(258882, "end_of_image"),
|
||
(258883, "end_of_audio"),
|
||
(255999, "boi"),
|
||
(256000, "boa"),
|
||
]
|
||
|
||
for (id, name) in specialTokens {
|
||
let decoded = tokenizer.decode(tokens: [id])
|
||
Swift.print(" Token \(id) (\(name)): '\(decoded)'")
|
||
}
|
||
|
||
// Test 3: What happens with end_of_turn token?
|
||
Swift.print("\n\nTest 3: Check if model responds to format tokens")
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// <|turn> token is 105
|
||
let turnStart = 105
|
||
let turnEnd = 106
|
||
|
||
// BOS + <|turn> + text + <turn|>
|
||
let prompt = "Describe this image"
|
||
var promptTokens = tokenizer.encode(text: prompt)
|
||
Swift.print("Prompt tokens: \(promptTokens)")
|
||
|
||
// We need to insert <|turn> and <turn|> correctly
|
||
// Actually, the tokenizer should handle these if they're in vocabulary
|
||
Swift.print("Note: Proper multimodal inference requires vision/audio features")
|
||
Swift.print(" Text-only generation is not the intended use case")
|
||
|
||
Swift.print("\n=== Summary ===")
|
||
Swift.print("E4B-MarkBase is a multimodal model (Gemma4ForConditionalGeneration)")
|
||
Swift.print("It requires:")
|
||
Swift.print(" 1. Vision input (image) processed through vision_tower")
|
||
Swift.print(" 2. Audio input (audio) processed through audio_tower")
|
||
Swift.print(" 3. Multimodal embedding interleaved with text tokens")
|
||
Swift.print("\nFor text-only generation, use a pure text model like Gemma-4-4B-IT")
|
||
}
|
||
|
||
func testContinuationMode() throws {
|
||
Swift.print("\n=== Continuation Mode Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
let sampler = Sampler()
|
||
|
||
// Test 1: Simple text continuation (no special format tokens)
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// "The capital of France is" - expecting "Paris"
|
||
let prompt1 = "The capital of France is"
|
||
var tokens1 = tokenizer.encode(text: prompt1)
|
||
Swift.print("\nPrompt: '\(prompt1)'")
|
||
Swift.print("Tokens: \(tokens1)")
|
||
Swift.print("Decoded: '\(tokenizer.decode(tokens: tokens1))'")
|
||
|
||
// Process prompt
|
||
for (pos, token) in tokens1.enumerated() {
|
||
let logits = try model.forward(tokenId: token, position: pos)
|
||
if pos == tokens1.count - 1 {
|
||
// Show top 10 predictions at last prompt position
|
||
let top10 = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("\nTop 10 predictions at prompt end:")
|
||
for (idx, val) in top10 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
}
|
||
}
|
||
|
||
// Generate continuation
|
||
Swift.print("\nGenerating continuation:")
|
||
for i in 0..<10 {
|
||
let pos = tokens1.count - 1
|
||
let logits = try model.forward(tokenId: tokens1.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens1.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print(" Token \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nContinuation: '\(tokenizer.decode(tokens: tokens1))'")
|
||
|
||
// Test 2: Another continuation
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
let prompt2 = "Hello, my name is"
|
||
var tokens2 = tokenizer.encode(text: prompt2)
|
||
Swift.print("\n\nPrompt 2: '\(prompt2)'")
|
||
|
||
// Process and generate
|
||
for (pos, token) in tokens2.enumerated() {
|
||
let logits = try model.forward(tokenId: token, position: pos)
|
||
if pos == tokens2.count - 1 {
|
||
let top5 = logits.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Top 5 at prompt end:")
|
||
for (idx, val) in top5 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx): \(val) ('\(tokenStr)')")
|
||
}
|
||
}
|
||
}
|
||
|
||
Swift.print("\nGenerating:")
|
||
for i in 0..<10 {
|
||
let pos = tokens2.count - 1
|
||
let logits = try model.forward(tokenId: tokens2.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens2.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print(" \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nResult: '\(tokenizer.decode(tokens: tokens2))'")
|
||
|
||
// Test 3: Check if "Paris" appears in predictions for "The capital of France is"
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
Swift.print("\n\n=== Checking 'Paris' prediction ===")
|
||
let prompt3 = "The capital of France is"
|
||
let tokens3 = tokenizer.encode(text: prompt3)
|
||
|
||
// Process prompt
|
||
for (pos, token) in tokens3.enumerated() {
|
||
let logits = try model.forward(tokenId: token, position: pos)
|
||
if pos == tokens3.count - 1 {
|
||
// Find "Paris" token
|
||
// Try different variations
|
||
let parisVariants = [
|
||
tokenizer.encode(text: "Paris"),
|
||
tokenizer.encode(text: " paris"),
|
||
tokenizer.encode(text: "▁Paris"),
|
||
]
|
||
Swift.print("Paris token variants: \(parisVariants)")
|
||
|
||
// Check if any variant is in top 50
|
||
let top50 = logits.enumerated().sorted { $0.element > $1.element }.prefix(50)
|
||
for (idx, val) in top50 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
if tokenStr.lowercased().contains("paris") {
|
||
Swift.print("Found 'Paris' variant at position \(idx) in top 50: rank \(top50.firstIndex(where: { $0.offset == idx }) ?? -1), logit=\(val)")
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func testKernelSelection() throws {
|
||
Swift.print("\n=== Kernel Selection Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
|
||
// Check if SIMD kernel is available
|
||
let simdAvailable = (try? engine.pipeline(named: "full_attention_simd")) != nil
|
||
Swift.print("full_attention_simd available: \(simdAvailable)")
|
||
|
||
let regularAvailable = (try? engine.pipeline(named: "full_attention")) != nil
|
||
Swift.print("full_attention available: \(regularAvailable)")
|
||
|
||
let simdSlidingAvailable = (try? engine.pipeline(named: "sliding_attention_simd")) != nil
|
||
Swift.print("sliding_attention_simd available: \(simdSlidingAvailable)")
|
||
|
||
let slidingAvailable = (try? engine.pipeline(named: "sliding_attention")) != nil
|
||
Swift.print("sliding_attention available: \(slidingAvailable)")
|
||
|
||
// Test with model
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
let sampler = Sampler()
|
||
|
||
Swift.print("\nModel layer types:")
|
||
for i in [0, 5, 6, 11, 17, 23, 29, 35, 41] {
|
||
let layer = model.layers[i]
|
||
Swift.print(" Layer \(i): \(layer.config.isSliding ? "sliding" : "full") attention")
|
||
}
|
||
|
||
// Generate with temperature=0.1
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
Swift.print("\nGenerating with temperature=0.1:")
|
||
var tokens = [2]
|
||
for i in 0..<20 {
|
||
let pos = tokens.count - 1
|
||
let logits = try model.forward(tokenId: tokens.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens))'")
|
||
}
|
||
|
||
func testSimpleBOSGeneration() throws {
|
||
Swift.print("\n=== Simple BOS Generation Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
let sampler = Sampler()
|
||
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// Generate from BOS only
|
||
Swift.print("\nGenerating from BOS token only (position 0):")
|
||
let logits0 = try model.forward(tokenId: 2, position: 0)
|
||
|
||
let top10 = logits0.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Generate 20 tokens from BOS
|
||
var tokens = [2]
|
||
for i in 0..<20 {
|
||
let pos = tokens.count - 1
|
||
let logits = try model.forward(tokenId: tokens.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens))'")
|
||
|
||
// Test with very low temperature (near greedy)
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
Swift.print("\n\nGenerating with temperature=0.1 (near-greedy):")
|
||
var tokens2 = [2]
|
||
for i in 0..<20 {
|
||
let pos = tokens2.count - 1
|
||
let logits = try model.forward(tokenId: tokens2.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens2.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens2))'")
|
||
|
||
// Test with top-K=1 (pure greedy)
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
Swift.print("\n\nGenerating with top-K=1 (greedy):")
|
||
var tokens3 = [2]
|
||
for i in 0..<20 {
|
||
let pos = tokens3.count - 1
|
||
let logits = try model.forward(tokenId: tokens3.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 1.0, topK: 1, filterUnusedTokens: true)
|
||
tokens3.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens3))'")
|
||
}
|
||
|
||
func testTokenizerEncoding() throws {
|
||
Swift.print("\n=== Tokenizer Encoding Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
// Test 1: Simple prompt without special tokens
|
||
let prompt1 = "What is the capital of France?"
|
||
let tokens1 = tokenizer.encode(text: prompt1)
|
||
let decoded1 = tokenizer.decode(tokens: tokens1)
|
||
Swift.print("\nTest 1: Simple prompt")
|
||
Swift.print(" Prompt: '\(prompt1)'")
|
||
Swift.print(" Tokens: \(tokens1)")
|
||
Swift.print(" Decoded: '\(decoded1)'")
|
||
|
||
// Check each token
|
||
Swift.print(" Token breakdown:")
|
||
for (i, token) in tokens1.enumerated() {
|
||
let decodedToken = tokenizer.decode(tokens: [token])
|
||
Swift.print(" \(i): Token \(token) -> '\(decodedToken)'")
|
||
}
|
||
|
||
// Test 2: Prompt with special format tokens
|
||
let prompt2 = "<bos><|turn>What is the capital of France?<turn|><|turn>"
|
||
let tokens2 = tokenizer.encode(text: prompt2)
|
||
let decoded2 = tokenizer.decode(tokens: tokens2)
|
||
Swift.print("\nTest 2: Prompt with format tokens")
|
||
Swift.print(" Prompt: '\(prompt2)'")
|
||
Swift.print(" Tokens: \(tokens2)")
|
||
Swift.print(" Decoded: '\(decoded2)'")
|
||
|
||
// Check first 20 tokens
|
||
Swift.print(" Token breakdown (first 20):")
|
||
for (i, token) in tokens2.enumerated().prefix(20) {
|
||
let decodedToken = tokenizer.decode(tokens: [token])
|
||
Swift.print(" \(i): Token \(token) -> '\(decodedToken)'")
|
||
}
|
||
|
||
// Test 3: Check if spaces are preserved
|
||
let prompt3 = "Hello World"
|
||
let tokens3 = tokenizer.encode(text: prompt3)
|
||
let decoded3 = tokenizer.decode(tokens: tokens3)
|
||
Swift.print("\nTest 3: Space preservation")
|
||
Swift.print(" Prompt: '\(prompt3)'")
|
||
Swift.print(" Tokens: \(tokens3)")
|
||
Swift.print(" Decoded: '\(decoded3)'")
|
||
Swift.print(" Are spaces preserved? \(decoded3.contains(" ") ? "YES" : "NO")")
|
||
|
||
// Test 4: Check underscore prefix tokens
|
||
// In Gemma, words usually start with ▁ (underscore) prefix
|
||
Swift.print("\nTest 4: Check for ▁ prefix in vocabulary")
|
||
// Manually check some vocabulary entries
|
||
// Token IDs for common words should have ▁ prefix
|
||
let commonTokens = [506, 532, 496, 107] // the, and, a, newline
|
||
for token in commonTokens {
|
||
let decodedToken = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Token \(token): '\(decodedToken)' - has prefix? \(decodedToken.hasPrefix("▁") || decodedToken.hasPrefix(" ") ? "YES" : "NO")")
|
||
}
|
||
}
|
||
|
||
func testSamplingStrategies() throws {
|
||
Swift.print("\n=== Sampling Strategies Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
// Process token 126019 at position 0, then check next token predictions
|
||
Swift.print("\nProcessing token 126019 at position 0...")
|
||
let logits = try model.forward(tokenId: 126019, position: 0)
|
||
|
||
// Test different sampling strategies
|
||
let top5 = logits.enumerated().sorted { $0.element > $1.element }.prefix(5)
|
||
Swift.print("Top 5 logits:")
|
||
for (idx, val) in top5 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Test 1: Greedy decoding (temperature=0, take argmax)
|
||
let sampler = Sampler()
|
||
let argmax = sampler.greedySample(logits: logits)
|
||
Swift.print("\nGreedy decoding: Token \(argmax) ('\(tokenizer.decode(tokens: [argmax]))')")
|
||
|
||
// Test 2: Temperature=0.1 (near-greedy)
|
||
let samples01 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95) }
|
||
Swift.print("\nTemperature=0.1, 5 samples:")
|
||
for (i, token) in samples01.enumerated() {
|
||
let tokenStr = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Sample \(i): Token \(token) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Test 3: Temperature=0.7 (default)
|
||
let samples07 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95) }
|
||
Swift.print("\nTemperature=0.7, 5 samples:")
|
||
for (i, token) in samples07.enumerated() {
|
||
let tokenStr = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Sample \(i): Token \(token) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Test 4: Temperature=1.0 (more random)
|
||
let samples10 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 1.0, topK: 50, topP: 0.95) }
|
||
Swift.print("\nTemperature=1.0, 5 samples:")
|
||
for (i, token) in samples10.enumerated() {
|
||
let tokenStr = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Sample \(i): Token \(token) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Test 5: Top-K=1 (force argmax)
|
||
let samplesK1 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 1.0, topK: 1) }
|
||
Swift.print("\nTop-K=1 (force argmax), 5 samples:")
|
||
for (i, token) in samplesK1.enumerated() {
|
||
let tokenStr = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Sample \(i): Token \(token) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Check if all top 5 tokens are unused tokens
|
||
let top5AreUnused = top5.allSatisfy { (idx, _) in idx >= 258000 && idx < 259000 }
|
||
Swift.print("\nAll top 5 are unused tokens (258000-258999): \(top5AreUnused)")
|
||
|
||
// Test with filtering enabled
|
||
Swift.print("\n=== Testing with unused token filtering enabled ===")
|
||
let samplesFiltered = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true) }
|
||
Swift.print("Temperature=0.7 with filtering, 5 samples:")
|
||
for (i, token) in samplesFiltered.enumerated() {
|
||
let tokenStr = tokenizer.decode(tokens: [token])
|
||
Swift.print(" Sample \(i): Token \(token) ('\(tokenStr)')")
|
||
}
|
||
|
||
// Test generation with filtering
|
||
Swift.print("\n=== Generation test with filtering ===")
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
var tokens2: [Int] = [2] // BOS
|
||
for pos in 0..<10 {
|
||
let nextLogits = try model.forward(tokenId: tokens2.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: nextLogits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens2.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Position \(pos + 1): sampled token \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
Swift.print("\nGenerated text: '\(tokenizer.decode(tokens: tokens2))'")
|
||
|
||
// Test with meaningful prompt
|
||
Swift.print("\n=== Generation with meaningful prompt ===")
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// Encode prompt
|
||
let prompt = "What is the capital of France?"
|
||
var tokens3 = tokenizer.encode(text: prompt)
|
||
Swift.print("Prompt tokens: \(tokens3)")
|
||
Swift.print("Prompt: '\(tokenizer.decode(tokens: tokens3))'")
|
||
|
||
// Process prompt
|
||
for (pos, token) in tokens3.enumerated() {
|
||
let logits = try model.forward(tokenId: token, position: pos)
|
||
if pos == tokens3.count - 1 {
|
||
// Sample from last position
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens3.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("First generated token: \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
}
|
||
|
||
// Generate more tokens
|
||
for i in 0..<20 {
|
||
let pos = tokens3.count - 1
|
||
let logits = try model.forward(tokenId: tokens3.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens3.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Generated token \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
|
||
Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokens3))'")
|
||
|
||
// Test with proper format tokens
|
||
Swift.print("\n=== Generation with proper format tokens ===")
|
||
for cache in model.kvCaches {
|
||
cache.reset()
|
||
}
|
||
|
||
// Use proper Gemma-4 format: <|turn> user question <turn|> <|turn> model response <turn|>
|
||
// Format: <bos> <|turn> user prompt <turn|> <|turn> model response
|
||
let promptFormatted = "<bos><|turn>What is the capital of France?<turn|><|turn>"
|
||
|
||
Swift.print("Formatted prompt: '\(promptFormatted)'")
|
||
|
||
// Encode
|
||
var tokens4 = tokenizer.encode(text: promptFormatted)
|
||
Swift.print("Tokens: \(tokens4)")
|
||
Swift.print("Decoded: '\(tokenizer.decode(tokens: tokens4))'")
|
||
|
||
// Process
|
||
for (pos, token) in tokens4.enumerated() {
|
||
let logits = try model.forward(tokenId: token, position: pos)
|
||
if pos == tokens4.count - 1 {
|
||
// Sample from last position
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokens4.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("First generated token: \(nextToken) ('\(tokenStr)')")
|
||
}
|
||
}
|
||
|
||
// Generate 20 tokens
|
||
var tokensGenerated = tokens4
|
||
for i in 0..<20 {
|
||
let pos = tokensGenerated.count - 1
|
||
let logits = try model.forward(tokenId: tokensGenerated.last!, position: pos)
|
||
let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
|
||
tokensGenerated.append(nextToken)
|
||
let tokenStr = tokenizer.decode(tokens: [nextToken])
|
||
Swift.print("Generated token \(i + 1): \(nextToken) ('\(tokenStr)')")
|
||
|
||
// Stop if we hit end of turn or EOS
|
||
if nextToken == 106 || nextToken == 1 {
|
||
break
|
||
}
|
||
}
|
||
|
||
Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokensGenerated))'")
|
||
}
|
||
|
||
func testCompareEmbedding126019() throws {
|
||
Swift.print("\n=== Compare Token 126019 Embedding: Swift vs Python ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
|
||
// Get embedding for token 126019
|
||
let embedWeight = model.embedWeight
|
||
|
||
// Token 126019
|
||
let tokenId = 126019
|
||
let numGroups = 40
|
||
let groupSize = 64
|
||
let hiddenDim = 2560
|
||
|
||
// Dequantize in Swift
|
||
var dequantized = [Float](repeating: 0, count: hiddenDim)
|
||
|
||
let weightPtr = embedWeight.weight.contents().assumingMemoryBound(to: UInt32.self)
|
||
let scalesPtr = embedWeight.scales.contents().assumingMemoryBound(to: Float.self)
|
||
let biasesPtr = embedWeight.biases.contents().assumingMemoryBound(to: Float.self)
|
||
|
||
for groupIdx in 0..<numGroups {
|
||
let scale = scalesPtr[tokenId * numGroups + groupIdx]
|
||
let bias = biasesPtr[tokenId * numGroups + groupIdx]
|
||
|
||
for u32Idx in 0..<8 {
|
||
let packed = weightPtr[tokenId * 320 + groupIdx * 8 + u32Idx]
|
||
|
||
for bitIdx in 0..<8 {
|
||
let val4bit = Int32((packed >> (UInt32(bitIdx) * 4)) & 0xF)
|
||
let signed = val4bit > 7 ? val4bit - 16 : val4bit
|
||
let deqValue = Float(signed) * scale + bias
|
||
|
||
let idx = groupIdx * groupSize + u32Idx * 8 + bitIdx
|
||
dequantized[idx] = deqValue
|
||
}
|
||
}
|
||
}
|
||
|
||
// Print first 20 values
|
||
Swift.print("Swift embedding for token 126019 (first 20 values):")
|
||
for i in 0..<20 {
|
||
Swift.print(" [\(i)] = \(dequantized[i])")
|
||
}
|
||
|
||
// Print some middle values
|
||
Swift.print("\nSwift embedding for token 126019 (values 1000-1020):")
|
||
for i in 1000..<1020 {
|
||
Swift.print(" [\(i)] = \(dequantized[i])")
|
||
}
|
||
|
||
// Print last 20 values
|
||
Swift.print("\nSwift embedding for token 126019 (last 20 values):")
|
||
for i in 2540..<2560 {
|
||
Swift.print(" [\(i)] = \(dequantized[i])")
|
||
}
|
||
|
||
// Compute norm
|
||
let norm = sqrt(dequantized.reduce(0) { $0 + $1 * $1 })
|
||
Swift.print("\nSwift embedding norm: \(norm)")
|
||
|
||
// Compare scales/biases
|
||
Swift.print("\nSwift scales (first 10):")
|
||
for i in 0..<10 {
|
||
Swift.print(" [\(i)] = \(scalesPtr[tokenId * numGroups + i])")
|
||
}
|
||
|
||
Swift.print("\nSwift biases (first 10):")
|
||
for i in 0..<10 {
|
||
Swift.print(" [\(i)] = \(biasesPtr[tokenId * numGroups + i])")
|
||
}
|
||
}
|
||
|
||
func testUnusedTokenEmbedding() throws {
|
||
Swift.print("\n=== Unused Token Embedding Test ===")
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
|
||
// Test unused token 258123 as input
|
||
Swift.print("\nProcessing unused token 258123 ('<unused2211>') at position 0...")
|
||
let logits = try model.forward(tokenId: 258123, position: 0)
|
||
|
||
let top10 = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Test unused token 258090
|
||
Swift.print("\nProcessing unused token 258090 ('<unused2178>') at position 0...")
|
||
let logits2 = try model.forward(tokenId: 258090, position: 0)
|
||
|
||
let top10_2 = logits2.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10_2 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Compare with BOS token
|
||
Swift.print("\nCompare with BOS token (2) at position 0...")
|
||
let logitsBOS = try model.forward(tokenId: 2, position: 0)
|
||
|
||
let top10_BOS = logitsBOS.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10_BOS {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Test token 126019 ('ccql') - the token that led to position 6 unused predictions
|
||
Swift.print("\nProcessing token 126019 ('ccql') at position 0...")
|
||
let logits_ccql = try model.forward(tokenId: 126019, position: 0)
|
||
|
||
let top10_ccql = logits_ccql.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10_ccql {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Test token 225448 ('globalMap') - another sampled token
|
||
Swift.print("\nProcessing token 225448 ('globalMap') at position 0...")
|
||
let logits_global = try model.forward(tokenId: 225448, position: 0)
|
||
|
||
let top10_global = logits_global.enumerated().sorted { $0.element > $1.element }.prefix(10)
|
||
Swift.print("Top 10 logits:")
|
||
for (idx, val) in top10_global {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
}
|
||
|
||
// ── 31B dense model test ─────────────────────────
|
||
|
||
func test31BModelLoading() throws {
|
||
Swift.print("\n=== Testing Gemma-4 31B (dense, 4-bit) ===")
|
||
|
||
let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
|
||
|
||
guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
|
||
Swift.print("✗ 31B model not found")
|
||
return
|
||
}
|
||
Swift.print("✓ 31B model found")
|
||
|
||
let totalWeights: Int64 = try FileManager.default.contentsOfDirectory(atPath: modelDir)
|
||
.filter { $0.hasSuffix(".safetensors") }
|
||
.reduce(0) { sum, file in
|
||
let attrs = try FileManager.default.attributesOfItem(atPath: modelDir + "/" + file)
|
||
return sum + (attrs[.size] as? Int64 ?? 0)
|
||
}
|
||
Swift.print(" Total weights: \(totalWeights / 1024 / 1024 / 1024) GB")
|
||
|
||
let engine = try MarkBaseEngine(autoCompile: true)
|
||
Swift.print("✓ Engine created")
|
||
|
||
Swift.print("\nLoading 31B model (~18 GB, may take 2-3 minutes)...")
|
||
Swift.print(" hidden_size=5376, 60 layers, 32 heads, 16 KV heads, intermediate=21504")
|
||
|
||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
|
||
|
||
Swift.print("✓✓✓ Model loaded!")
|
||
Swift.print(" Layers: \(model.numHiddenLayers)")
|
||
Swift.print(" Hidden size: \(model.hiddenSize)")
|
||
Swift.print(" Vocab size: \(model.vocabSize)")
|
||
let fullCount = model.layerTypesIsFull.filter { $0 }.count
|
||
let slideCount = model.layerTypesIsFull.filter { !$0 }.count
|
||
Swift.print(" Layer types: \(fullCount) full, \(slideCount) sliding")
|
||
|
||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
Swift.print("✓ Tokenizer loaded")
|
||
|
||
// Forward pass
|
||
Swift.print("\nForward pass test...")
|
||
let tokens = tokenizer.encode(text: "Hello")
|
||
Swift.print("Tokenized 'Hello': \(tokens)")
|
||
|
||
let logits = try model.forward(tokenId: tokens[0], position: 0)
|
||
Swift.print("✓ Forward pass: \(logits.count) logits, max=\(logits.max() ?? -999)")
|
||
|
||
let sorted = logits.enumerated().sorted { $0.element > $1.element }
|
||
let top5 = sorted.prefix(5)
|
||
Swift.print("Top 5:")
|
||
for (idx, val) in top5 {
|
||
let tokenStr = tokenizer.decode(tokens: [idx])
|
||
Swift.print(" Token \(idx) ('\(tokenStr)'): \(val)")
|
||
}
|
||
|
||
// Generation test (3 tokens)
|
||
Swift.print("\nGeneration test (3 tokens)...")
|
||
do {
|
||
let genConfig = GenerationConfig(maxTokens: 3, temperature: 0.7)
|
||
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||
let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
|
||
Swift.print("Generated: '\(response)'")
|
||
XCTAssertGreaterThan(response.count, 0, "Should generate text")
|
||
} catch {
|
||
Swift.print("⚠️ Generation failed: \(error)")
|
||
}
|
||
|
||
Swift.print("\n✅ 31B test complete!")
|
||
}
|
||
} |