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:
MarkBase Admin
@@ -0,0 +1,239 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Generation Configuration
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public struct GenerationConfig: Sendable {
|
||||
public let maxTokens: Int
|
||||
public let temperature: Float
|
||||
public let topK: Int?
|
||||
public let topP: Float?
|
||||
public let stopTokens: [Int]?
|
||||
|
||||
public init(
|
||||
maxTokens: Int = 100,
|
||||
temperature: Float = 1.0,
|
||||
topK: Int? = nil,
|
||||
topP: Float? = nil,
|
||||
stopTokens: [Int]? = nil
|
||||
) {
|
||||
self.maxTokens = maxTokens
|
||||
self.temperature = temperature
|
||||
self.topK = topK
|
||||
self.topP = topP
|
||||
self.stopTokens = stopTokens
|
||||
}
|
||||
|
||||
// Default configuration
|
||||
public static let defaultConfig = GenerationConfig(maxTokens: 100, temperature: 1.0)
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Streaming Generator - Token-by-token generation
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public final class StreamingGenerator: @unchecked Sendable {
|
||||
private let model: E4BModel
|
||||
private let tokenizer: Tokenizer
|
||||
private let engine: MarkBaseEngine
|
||||
private let sampler: Sampler
|
||||
|
||||
public init(model: E4BModel, tokenizer: Tokenizer, engine: MarkBaseEngine) {
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
self.engine = engine
|
||||
self.sampler = Sampler()
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Stream Generate - AsyncStream for token-by-token output
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public func generate(
|
||||
prompt: String,
|
||||
config: GenerationConfig = .defaultConfig
|
||||
) -> AsyncStream<String> {
|
||||
return AsyncStream { continuation in
|
||||
Task {
|
||||
do {
|
||||
// Encode prompt
|
||||
let promptTokens = tokenizer.encode(text: prompt)
|
||||
|
||||
// Pre-fill KV cache with prompt tokens
|
||||
var lastLogits: [Float] = []
|
||||
for (position, tokenId) in promptTokens.enumerated() {
|
||||
lastLogits = try model.forward(tokenId: tokenId, position: position)
|
||||
}
|
||||
|
||||
// Generate tokens
|
||||
var generatedTokens: [Int] = []
|
||||
var position = promptTokens.count
|
||||
var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
|
||||
|
||||
for _ in 0..<config.maxTokens {
|
||||
// Sample next token
|
||||
let nextToken = sampler.sample(
|
||||
logits: lastLogits,
|
||||
temperature: config.temperature,
|
||||
topK: config.topK,
|
||||
topP: config.topP
|
||||
)
|
||||
|
||||
// Debug: print selected token
|
||||
if generatedTokens.count < 5 {
|
||||
print("[DEBUG] Generated token \(generatedTokens.count): ID=\(nextToken), raw='\(tokenizer.rawToken(for: nextToken) ?? "nil")'")
|
||||
fflush(stdout)
|
||||
}
|
||||
|
||||
// Check stop tokens
|
||||
if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
// Check EOS (handle multiple EOS tokens)
|
||||
if tokenizer.eosTokenIds.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
// Add token
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
// Decode using streaming decoder (handles multi-byte UTF-8 correctly)
|
||||
let tokenText = streamDecoder.consume(tokenId: nextToken)
|
||||
if !tokenText.isEmpty {
|
||||
continuation.yield(tokenText)
|
||||
}
|
||||
|
||||
// Forward pass for next token
|
||||
lastLogits = try model.forward(tokenId: nextToken, position: position)
|
||||
position += 1
|
||||
}
|
||||
|
||||
continuation.finish()
|
||||
|
||||
} catch {
|
||||
continuation.finish()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Complete Generate - Returns full text
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public func generateComplete(
|
||||
prompt: String,
|
||||
config: GenerationConfig = .defaultConfig
|
||||
) throws -> String {
|
||||
print("[GEN COMPLETE] Starting generation for prompt: '\(prompt)'")
|
||||
fflush(stdout)
|
||||
|
||||
// Encode prompt
|
||||
print("[GEN COMPLETE] Encoding prompt...")
|
||||
fflush(stdout)
|
||||
let promptTokens = tokenizer.encode(text: prompt)
|
||||
print("[GEN COMPLETE] Encoded to \(promptTokens.count) tokens: \(promptTokens)")
|
||||
fflush(stdout)
|
||||
|
||||
// Pre-fill KV cache with prompt tokens
|
||||
print("[GEN COMPLETE] Starting forward pass for prompt tokens...")
|
||||
fflush(stdout)
|
||||
var lastLogits: [Float] = []
|
||||
for (position, tokenId) in promptTokens.enumerated() {
|
||||
print("[GEN COMPLETE] Forward pass for token \(tokenId) at position \(position)")
|
||||
fflush(stdout)
|
||||
lastLogits = try model.forward(tokenId: tokenId, position: position)
|
||||
print("[GEN COMPLETE] Forward pass completed, logits count: \(lastLogits.count)")
|
||||
fflush(stdout)
|
||||
}
|
||||
print("[GEN COMPLETE] All prompt tokens processed")
|
||||
fflush(stdout)
|
||||
|
||||
// Generate tokens
|
||||
var generatedTokens: [Int] = []
|
||||
var position = promptTokens.count
|
||||
|
||||
for _ in 0..<config.maxTokens {
|
||||
let nextToken = sampler.sample(
|
||||
logits: lastLogits,
|
||||
temperature: config.temperature,
|
||||
topK: config.topK,
|
||||
topP: config.topP
|
||||
)
|
||||
|
||||
// Debug: print selected token
|
||||
if generatedTokens.count < 5 {
|
||||
print("[DEBUG generateComplete] Token \(generatedTokens.count): ID=\(nextToken), raw='\(tokenizer.rawToken(for: nextToken) ?? "nil")'")
|
||||
fflush(stdout)
|
||||
// Print logits stats
|
||||
let maxLogit = lastLogits.max() ?? 0
|
||||
let minLogit = lastLogits.min() ?? 0
|
||||
let maxIdx = lastLogits.indices.filter { lastLogits[$0] == maxLogit }.first ?? -1
|
||||
print("[DEBUG] Logits: max=\(maxLogit) at idx=\(maxIdx), min=\(minLogit)")
|
||||
fflush(stdout)
|
||||
}
|
||||
|
||||
if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
if tokenizer.eosTokenIds.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
// Forward pass for next token
|
||||
lastLogits = try model.forward(tokenId: nextToken, position: position)
|
||||
position += 1
|
||||
}
|
||||
|
||||
// Decode full response
|
||||
return tokenizer.decode(tokens: generatedTokens)
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Generate with Token IDs - Returns token sequence
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public func generateTokens(
|
||||
promptTokens: [Int],
|
||||
config: GenerationConfig = .defaultConfig
|
||||
) throws -> [Int] {
|
||||
// Pre-fill KV cache with prompt tokens
|
||||
var lastLogits: [Float] = []
|
||||
for (position, tokenId) in promptTokens.enumerated() {
|
||||
lastLogits = try model.forward(tokenId: tokenId, position: position)
|
||||
}
|
||||
|
||||
var generatedTokens: [Int] = []
|
||||
var position = promptTokens.count
|
||||
|
||||
for _ in 0..<config.maxTokens {
|
||||
let nextToken = sampler.sample(
|
||||
logits: lastLogits,
|
||||
temperature: config.temperature,
|
||||
topK: config.topK,
|
||||
topP: config.topP
|
||||
)
|
||||
|
||||
if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
if tokenizer.eosTokenIds.contains(nextToken) {
|
||||
break
|
||||
}
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
// Forward pass for next token
|
||||
lastLogits = try model.forward(tokenId: nextToken, position: position)
|
||||
position += 1
|
||||
}
|
||||
|
||||
return generatedTokens
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user