Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled
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
This commit is contained in:
@@ -0,0 +1,217 @@
|
||||
import XCTest
|
||||
@testable import MarkBase
|
||||
|
||||
final class G12BTextGenerationTests: XCTestCase {
|
||||
|
||||
func testTextGeneration() throws {
|
||||
print("\n═══════════════════════════════════════")
|
||||
print(" 12B Text Generation Test")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
|
||||
|
||||
print("Step 1: Load tokenizer...")
|
||||
// Simple tokenizer test - check if tokenizer.json exists
|
||||
let tokenizerPath = modelDir + "/tokenizer.json"
|
||||
let tokenizerExists = FileManager.default.fileExists(atPath: tokenizerPath)
|
||||
print(" Tokenizer exists: \(tokenizerExists)")
|
||||
if !tokenizerExists {
|
||||
print(" ⚠️ No tokenizer found - using simple token IDs\n")
|
||||
}
|
||||
|
||||
print("Step 2: Initialize engine and model...")
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
print(" ✓ Model loaded\n")
|
||||
|
||||
print("Step 3: Test generation with simple prompts...")
|
||||
|
||||
// Test 1: Generate from token 0 (usually special token)
|
||||
print("\nTest 1: Start token (ID=0)")
|
||||
var generatedTokens: [Int] = []
|
||||
let startToken = 0
|
||||
|
||||
for position in 0..<10 {
|
||||
let logits = try model.forward(tokenId: startToken, position: position)
|
||||
|
||||
// Find max logit (greedy sampling)
|
||||
var maxLogit = logits[0]
|
||||
var maxIdx = 0
|
||||
for i in 1..<logits.count {
|
||||
if logits[i] > maxLogit {
|
||||
maxLogit = logits[i]
|
||||
maxIdx = i
|
||||
}
|
||||
}
|
||||
generatedTokens.append(maxIdx)
|
||||
print(" Position \(position): token=\(maxIdx), logit=\(maxLogit)")
|
||||
}
|
||||
|
||||
print(" Generated tokens: \(generatedTokens)")
|
||||
|
||||
// Test 2: Different start token
|
||||
print("\nTest 2: Token ID=1")
|
||||
generatedTokens = []
|
||||
for position in 0..<10 {
|
||||
let logits = try model.forward(tokenId: 1, position: position)
|
||||
var maxLogit = logits[0]
|
||||
var maxIdx = 0
|
||||
for i in 1..<logits.count {
|
||||
if logits[i] > maxLogit {
|
||||
maxLogit = logits[i]
|
||||
maxIdx = i
|
||||
}
|
||||
}
|
||||
generatedTokens.append(maxIdx)
|
||||
}
|
||||
print(" Generated tokens: \(generatedTokens)")
|
||||
|
||||
// Test 3: Sequence generation
|
||||
print("\nTest 3: Sequence generation (different tokens per position)")
|
||||
var sequence: [Int] = [1, 2, 3, 4, 5] // Start sequence
|
||||
print(" Input sequence: \(sequence)")
|
||||
|
||||
for i in 0..<5 {
|
||||
let logits = try model.forward(tokenId: sequence[i], position: i)
|
||||
var maxLogit = logits[0]
|
||||
var maxIdx = 0
|
||||
for j in 1..<logits.count {
|
||||
if logits[j] > maxLogit {
|
||||
maxLogit = logits[j]
|
||||
maxIdx = j
|
||||
}
|
||||
}
|
||||
sequence.append(maxIdx)
|
||||
}
|
||||
print(" Extended sequence: \(sequence)")
|
||||
|
||||
print("\n═══════════════════════════════════════")
|
||||
print("✓ Text generation test passed")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
// Verify diversity - tokens should not all be the same
|
||||
let uniqueTokens = Set(generatedTokens)
|
||||
XCTAssertGreaterThan(uniqueTokens.count, 1, "Generated tokens should have diversity")
|
||||
}
|
||||
|
||||
func testE4BInference() throws {
|
||||
print("\n═══════════════════════════════════════")
|
||||
print(" E4B-MarkBase Inference Test")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
|
||||
|
||||
print("Step 1: Load engine...")
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
print(" ✓ Engine initialized\n")
|
||||
|
||||
print("Step 2: Load tokenizer...")
|
||||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||||
print(" ✓ Tokenizer loaded: vocabSize=\(tokenizer.vocabSize)\n")
|
||||
|
||||
print("Step 3: Load model...")
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
print(" ✓ Model loaded: hiddenSize=\(model.hiddenSize), layers=\(model.numHiddenLayers)\n")
|
||||
|
||||
// Debug: check logits at position 0
|
||||
print("Debug: Single token test (BOS at pos 0)")
|
||||
let logits0 = try model.forward(tokenId: 2, position: 0)
|
||||
print(" logits count: \(logits0.count)")
|
||||
print(" min: \(logits0.min() ?? 0), max: \(logits0.max() ?? 0)")
|
||||
let sorted0 = logits0.enumerated().sorted { $0.element > $1.element }
|
||||
print(" Top 10 tokens:")
|
||||
for (i, (idx, val)) in sorted0.prefix(10).enumerated() {
|
||||
let text = tokenizer.decode(tokens: [idx])
|
||||
print(" \(i+1). token \(idx) '\(text)': \(val)")
|
||||
}
|
||||
print("")
|
||||
|
||||
print("Step 4: Create generator...")
|
||||
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||||
|
||||
print("Step 5: Text completion inference...")
|
||||
let prompt = "The capital of France is"
|
||||
print(" Prompt: \"\(prompt)\"")
|
||||
let config = GenerationConfig(maxTokens: 30, temperature: 1.0, topK: 40, topP: 0.9)
|
||||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||||
print(" Response: \"\(response)\"")
|
||||
|
||||
XCTAssertFalse(response.isEmpty, "Response should not be empty")
|
||||
print(" ✓ Text generation complete\n")
|
||||
|
||||
print("Step 6: Chat-style inference...")
|
||||
let chatPrompt = "<|turn>user\nWhat is 2+2?<turn|>\n<|turn>model\n"
|
||||
print(" Prompt: \"What is 2+2?\"")
|
||||
let chatResponse = try generator.generateComplete(prompt: chatPrompt, config: config)
|
||||
print(" Response: \"\(chatResponse)\"")
|
||||
|
||||
XCTAssertFalse(chatResponse.isEmpty, "Chat response should not be empty")
|
||||
print(" ✓ Chat generation complete\n")
|
||||
|
||||
print("Step 7: Chinese generation test...")
|
||||
let chinesePrompt = "<|turn>user\n用中文說你好<turn|>\n<|turn>model\n"
|
||||
print(" Prompt: \"用中文說你好\"")
|
||||
let chineseResponse = try generator.generateComplete(prompt: chinesePrompt, config: config)
|
||||
print(" Response: \"\(chineseResponse)\"")
|
||||
|
||||
print("\n═══════════════════════════════════════")
|
||||
print("✓ E4B inference test passed")
|
||||
print("═══════════════════════════════════════\n")
|
||||
}
|
||||
|
||||
func testGenerationQuality() throws {
|
||||
print("\n═══════════════════════════════════════")
|
||||
print(" 12B Generation Quality Test")
|
||||
print("═══════════════════════════════════════\n")
|
||||
|
||||
let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
|
||||
|
||||
print("Step 1: Load model...")
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
print(" ✓ Model loaded\n")
|
||||
|
||||
print("Step 2: Check logits distribution...")
|
||||
|
||||
// Generate logits for several tokens
|
||||
var allLogits: [[Float]] = []
|
||||
for tokenId in [0, 1, 100, 1000, 10000] {
|
||||
let logits = try model.forward(tokenId: tokenId, position: 0)
|
||||
allLogits.append(logits)
|
||||
|
||||
// Stats
|
||||
let slice = logits[0..<1000]
|
||||
let minVal = slice.min() ?? 0
|
||||
let maxVal = slice.max() ?? 0
|
||||
let mean = slice.reduce(0, +) / Float(slice.count)
|
||||
|
||||
print(" Token \(tokenId): min=\(minVal), max=\(maxVal), mean=\(mean)")
|
||||
}
|
||||
|
||||
print("\nStep 3: Verify generation diversity...")
|
||||
|
||||
// Generate 20 tokens and check if they're diverse
|
||||
var tokens: [Int] = []
|
||||
for i in 0..<20 {
|
||||
let logits = try model.forward(tokenId: i, position: i)
|
||||
|
||||
// Top-5 sampling (not just greedy)
|
||||
var sortedLogits = logits.enumerated().sorted { $0.element > $1.element }
|
||||
let top5 = sortedLogits[0..<5].map { $0.offset }
|
||||
|
||||
// Pick random from top-5 (simulate sampling)
|
||||
let selectedToken = top5[i % 5] // Deterministic for test
|
||||
tokens.append(selectedToken)
|
||||
}
|
||||
|
||||
print(" Generated tokens: \(tokens)")
|
||||
let uniqueTokens = Set(tokens)
|
||||
print(" Unique tokens: \(uniqueTokens.count)")
|
||||
|
||||
XCTAssertGreaterThan(uniqueTokens.count, 5, "Should have reasonable diversity")
|
||||
|
||||
print("\n═══════════════════════════════════════")
|
||||
print("✓ Generation quality test passed")
|
||||
print("═══════════════════════════════════════\n")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user