MarkBase Admin
8a66b9086a
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
1177 lines
46 KiB
Swift
1177 lines
46 KiB
Swift
import AVFoundation
|
||
import CoreImage
|
||
import Foundation
|
||
import MarkBase
|
||
import Hummingbird
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// OpenAI Compatible API Server with SSE Support
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public final class MarkBaseServer: @unchecked Sendable {
|
||
private let modelDir: String
|
||
private let modelId: String
|
||
private let maxContextLength: Int
|
||
|
||
// Runtime state
|
||
private let engine: MarkBaseEngine
|
||
private let model: E4BModel
|
||
private let tokenizer: Tokenizer
|
||
private let generator: StreamingGenerator
|
||
private let sampler: Sampler
|
||
private let multimodalModel: MultimodalModel?
|
||
|
||
public init(modelDir: String, modelId: String = "markbase-12b", maxContextLength: Int = 512) throws {
|
||
self.modelDir = modelDir
|
||
self.modelId = modelId
|
||
self.maxContextLength = maxContextLength
|
||
|
||
// Validate model path
|
||
try Validator.validateModelPath(modelDir)
|
||
|
||
print("Loading model from \(modelDir)...")
|
||
engine = try MarkBaseEngine(autoCompile: true)
|
||
model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxContextLength)
|
||
tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||
sampler = Sampler()
|
||
print("✓ Model loaded: \(modelId)")
|
||
// Skip multimodal loading for faster startup (it creates another E4BModel internally)
|
||
multimodalModel = nil
|
||
// multimodalModel = try? MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: maxContextLength)
|
||
// if multimodalModel != nil { print(" ✓ Multimodal model ready") }
|
||
print(" Layers: \(model.numHiddenLayers)")
|
||
print(" Vocab: \(model.vocabSize)")
|
||
}
|
||
|
||
public func start(port: Int = 8080) async throws {
|
||
print(" Port: \(port)")
|
||
|
||
// Create router
|
||
let router = Router()
|
||
|
||
// Middleware
|
||
router.middlewares.add(LogRequestsMiddleware(.info))
|
||
router.middlewares.add(CORSMiddleware(
|
||
allowOrigin: .all,
|
||
allowHeaders: [.contentType, .authorization],
|
||
allowMethods: [.get, .post, .options]
|
||
))
|
||
|
||
// Routes
|
||
router.get("/health") { _, _ in
|
||
return "OK"
|
||
}
|
||
|
||
router.get("/v1/models") { _, _ in
|
||
let models = [
|
||
ModelDetails(
|
||
id: self.modelId,
|
||
capabilities: ModelCapabilities(
|
||
vision: self.multimodalModel != nil,
|
||
audio: self.multimodalModel?.audioTowerFull != nil
|
||
),
|
||
parameters: ModelParameters(
|
||
context_length: self.maxContextLength,
|
||
num_hidden_layers: self.model.numHiddenLayers,
|
||
hidden_size: self.model.hiddenSize,
|
||
vocab_size: self.model.vocabSize,
|
||
num_attention_heads: 8,
|
||
num_kv_heads: 2
|
||
)
|
||
)
|
||
]
|
||
return models
|
||
}
|
||
|
||
router.post("/v1/chat/completions") { request, _ in
|
||
let buffer = try await request.body.collect(upTo: .max)
|
||
|
||
let req = try JSONDecoder().decode(ChatCompletionRequest.self, from: Data(buffer: buffer))
|
||
let response = try self.handleChatCompletion(messages: req.messages, config: req.toGenerationConfig())
|
||
|
||
return try ByteBuffer(data: JSONEncoder().encode(response))
|
||
}
|
||
|
||
router.post("/v1/multimodal/chat/completions") { request, _ in
|
||
print("[DEBUG ROUTER] Received multimodal request")
|
||
let buffer = try await request.body.collect(upTo: .max)
|
||
print("[DEBUG ROUTER] Body collected: \(buffer.readableBytes) bytes")
|
||
|
||
do {
|
||
let req = try JSONDecoder().decode(MultimodalChatCompletionRequest.self, from: Data(buffer: buffer))
|
||
print("[DEBUG ROUTER] Request decoded successfully, messages: \(req.messages.count)")
|
||
let response = try self.handleMultimodalChatCompletion(messages: req.messages, config: req.toGenerationConfig())
|
||
print("[DEBUG ROUTER] Response generated")
|
||
return try ByteBuffer(data: JSONEncoder().encode(response))
|
||
} catch {
|
||
print("[DEBUG ROUTER] Error: \(error)")
|
||
throw error
|
||
}
|
||
}
|
||
|
||
// Create Hummingbird app
|
||
let app = Application(
|
||
router: router,
|
||
configuration: .init(address: .hostname(port: port))
|
||
)
|
||
|
||
print("\nEndpoints:")
|
||
print(" GET /health")
|
||
print(" GET /v1/models")
|
||
print(" POST /v1/chat/completions")
|
||
print(" POST /v1/multimodal/chat/completions")
|
||
print("\nServer starting on port \(port)...")
|
||
|
||
try await app.run()
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Multimodal Handlers
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public func handleMultimodalChatCompletion(
|
||
messages: [MultimodalMessage],
|
||
config: GenerationConfig
|
||
) throws -> ChatCompletionResponse {
|
||
guard multimodalModel != nil else {
|
||
throw MarkBaseError.multimodalNotSupported
|
||
}
|
||
|
||
print("[DEBUG] handleMultimodalChatCompletion: Processing multimodal request")
|
||
|
||
// Build multimodal prompt
|
||
var textParts: [String] = []
|
||
var hasImage = false
|
||
var hasAudio = false
|
||
var imageData: Data? = nil
|
||
var audioData: Data? = nil
|
||
|
||
for message in messages {
|
||
print("[DEBUG] Message role: \(message.role), content parts: \(message.content.count)")
|
||
for part in message.content {
|
||
switch part {
|
||
case .text(let text):
|
||
print("[DEBUG] Text part: \(text)")
|
||
textParts.append(text)
|
||
case .imageUrl(let url):
|
||
print("[DEBUG] Image URL detected: \(url.url.prefix(50))...")
|
||
hasImage = true
|
||
// Load image from URL
|
||
if url.url.hasPrefix("data:image") {
|
||
// Base64 encoded
|
||
let parts = url.url.split(separator: ",")
|
||
if parts.count == 2 {
|
||
imageData = Data(base64Encoded: String(parts[1]))
|
||
print("[DEBUG] Base64 decoded, data size: \(imageData?.count ?? 0)")
|
||
}
|
||
} else if url.url.hasPrefix("file://") {
|
||
// Local file
|
||
imageData = try Data(contentsOf: URL(fileURLWithPath: String(url.url.dropFirst(7))))
|
||
print("[DEBUG] File loaded, data size: \(imageData?.count ?? 0)")
|
||
}
|
||
case .audioUrl(let url):
|
||
print("[DEBUG] Audio URL detected: \(url.url.prefix(50))...")
|
||
hasAudio = true
|
||
// Load audio from URL
|
||
if url.url.hasPrefix("data:audio") {
|
||
// Base64 encoded
|
||
let parts = url.url.split(separator: ",")
|
||
if parts.count == 2 {
|
||
audioData = Data(base64Encoded: String(parts[1]))
|
||
print("[DEBUG] Audio base64 decoded, data size: \(audioData?.count ?? 0)")
|
||
}
|
||
} else if url.url.hasPrefix("file://") {
|
||
// Local file
|
||
audioData = try Data(contentsOf: URL(fileURLWithPath: String(url.url.dropFirst(7))))
|
||
print("[DEBUG] Audio file loaded, data size: \(audioData?.count ?? 0)")
|
||
}
|
||
case .videoUrl:
|
||
// Video handling (future)
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
let prompt = textParts.joined(separator: " ")
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
|
||
// Process image if present
|
||
if hasImage, let data = imageData {
|
||
print("[DEBUG] Processing image data: \(data.count) bytes")
|
||
|
||
// Vision preprocessing
|
||
let visionFeatures = try processImageData(data)
|
||
print("[DEBUG] Vision features created, buffer length: \(visionFeatures.length)")
|
||
|
||
// Generate with vision conditioning
|
||
let response = try generateWithVision(
|
||
textTokens: promptTokens,
|
||
visionFeatures: visionFeatures,
|
||
maxTokens: config.maxTokens
|
||
)
|
||
print("[DEBUG] Generated response: \(response.prefix(100))")
|
||
|
||
return ChatCompletionResponse(
|
||
id: generateId("chatcmpl"),
|
||
object: "chat.completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
Choice(
|
||
index: 0,
|
||
message: ChatMessage(role: "assistant", content: response),
|
||
finish_reason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: promptTokens.count,
|
||
completionTokens: tokenizer.encode(text: response).count,
|
||
totalTokens: promptTokens.count + tokenizer.encode(text: response).count
|
||
)
|
||
)
|
||
} else if hasAudio, let data = audioData {
|
||
print("[DEBUG] Processing audio data: \(data.count) bytes")
|
||
|
||
// Audio preprocessing
|
||
let audioFeatures = try processAudioData(data)
|
||
print("[DEBUG] Audio features created")
|
||
|
||
// Generate with audio conditioning
|
||
// Note: Need to determine numFrames from audio length
|
||
// For now, use placeholder
|
||
let numFrames = 100 // Placeholder
|
||
|
||
let response = try generateWithAudio(
|
||
textTokens: promptTokens,
|
||
audioFeatures: audioFeatures,
|
||
numFrames: numFrames,
|
||
maxTokens: config.maxTokens
|
||
)
|
||
print("[DEBUG] Generated audio response: \(response.prefix(100))")
|
||
|
||
return ChatCompletionResponse(
|
||
id: generateId("chatcmpl"),
|
||
object: "chat.completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
Choice(
|
||
index: 0,
|
||
message: ChatMessage(role: "assistant", content: response),
|
||
finish_reason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: promptTokens.count,
|
||
completionTokens: tokenizer.encode(text: response).count,
|
||
totalTokens: promptTokens.count + tokenizer.encode(text: response).count
|
||
)
|
||
)
|
||
} else {
|
||
// Pure text generation
|
||
return try handleChatCompletion(
|
||
messages: messages.map { ChatMessage(role: $0.role, content: $0.content.map { part in
|
||
if case .text(let t) = part { return t }
|
||
return ""
|
||
}.joined()) },
|
||
config: config
|
||
)
|
||
}
|
||
}
|
||
|
||
private func processImageData(_ data: Data) throws -> MTLBuffer {
|
||
print("[VISION] Processing image data: \(data.count) bytes")
|
||
|
||
// Create CIImage from data
|
||
guard let ciImage = CIImage(data: data) else {
|
||
print("[VISION] Failed to create CIImage")
|
||
throw MarkBaseError.imageProcessingFailed
|
||
}
|
||
print("[VISION] Image size: \(ciImage.extent.width) x \(ciImage.extent.height)")
|
||
|
||
// 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 {
|
||
print("[VISION] Failed to resize image")
|
||
throw MarkBaseError.imageProcessingFailed
|
||
}
|
||
print("[VISION] Resized to: \(resized.extent.width) x \(resized.extent.height)")
|
||
|
||
// Convert to pixel data
|
||
let context = CIContext()
|
||
guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
|
||
print("[VISION] Failed to create CGImage")
|
||
throw MarkBaseError.imageProcessingFailed
|
||
}
|
||
|
||
// Extract RGB pixels
|
||
let dataProvider = cgImage.dataProvider!
|
||
let pixelData = dataProvider.data!
|
||
let ptr = CFDataGetBytePtr(pixelData)!
|
||
|
||
// Create patch embeddings (16x16 patches)
|
||
let patchSize = 16
|
||
let numPatches = 14 * 14 // 224/16 = 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
|
||
|
||
// Clamp to bounds
|
||
if globalY >= 224 || globalX >= 224 { continue }
|
||
|
||
let pixelIdx = globalY * 224 + globalX
|
||
let offset = pixelIdx * 4 // RGBA
|
||
|
||
// Normalize to [0, 1]
|
||
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
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Debug: Show first patch stats
|
||
var firstPatchSum: Float = 0
|
||
for i in 0..<768 {
|
||
firstPatchSum += patchEmbeddings[i]
|
||
}
|
||
print("[VISION] First patch RGB mean: R=\(patchEmbeddings[0]), G=\(patchEmbeddings[1]), B=\(patchEmbeddings[2]), sum=\(firstPatchSum/768)")
|
||
|
||
// Create buffer
|
||
let buffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
|
||
print("[VISION] Patch embeddings buffer created: \(buffer.length) bytes")
|
||
return buffer
|
||
}
|
||
|
||
private func generateWithVision(
|
||
textTokens: [Int],
|
||
visionFeatures: MTLBuffer,
|
||
maxTokens: Int
|
||
) throws -> String {
|
||
print("[VISION GEN] Starting vision-guided generation")
|
||
print("[VISION GEN] Text tokens: \(textTokens.count), max tokens: \(maxTokens)")
|
||
|
||
guard let mm = multimodalModel,
|
||
let tower = mm.visionTowerFull else {
|
||
print("[VISION GEN] Multimodal model not available")
|
||
throw MarkBaseError.multimodalNotSupported
|
||
}
|
||
print("[VISION GEN] Vision tower available")
|
||
|
||
let numPatches = 196
|
||
let hiddenSize = 2560
|
||
|
||
// Vision tower forward
|
||
print("[VISION GEN] Running vision tower forward pass...")
|
||
let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * hiddenSize * 4)!
|
||
try tower.forward(patchEmbeddings: visionFeatures, numPatches: numPatches, outputBuffer: visionOutputBuffer)
|
||
print("[VISION GEN] Vision tower forward completed")
|
||
|
||
// Pool embeddings
|
||
let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
|
||
|
||
// Debug: Check vision output stats
|
||
var outputMag: Float = 0
|
||
for p in 0..<min(5, numPatches) {
|
||
var patchMag: Float = 0
|
||
for i in 0..<hiddenSize {
|
||
let val = visionPtr[p * hiddenSize + i]
|
||
patchMag += val * val
|
||
}
|
||
outputMag += sqrt(patchMag)
|
||
}
|
||
print("[VISION GEN] Vision output magnitude (avg 5 patches): \(outputMag/5)")
|
||
|
||
var pooled = [Float](repeating: 0, count: hiddenSize)
|
||
|
||
for i in 0..<hiddenSize {
|
||
var sum: Float = 0
|
||
for p in 0..<numPatches {
|
||
sum += visionPtr[p * hiddenSize + i]
|
||
}
|
||
pooled[i] = sum / Float(numPatches)
|
||
}
|
||
|
||
// Normalize to match text embeddings (magnitude ~5)
|
||
let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
print("[VISION GEN] Pooled embedding magnitude before norm: \(mag)")
|
||
|
||
let scale: Float = 5.0 / mag
|
||
for i in 0..<hiddenSize {
|
||
pooled[i] *= scale
|
||
}
|
||
|
||
let magAfterNorm = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
print("[VISION GEN] Pooled embedding magnitude after norm: \(magAfterNorm)")
|
||
|
||
// Create inference
|
||
let inference = try MultimodalInference(model: mm)
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
|
||
|
||
// Generate
|
||
print("[VISION GEN] Starting multimodal inference...")
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: textTokens,
|
||
precomputedVisionEmbedding: pooledBuffer,
|
||
maxTokens: maxTokens
|
||
)
|
||
print("[VISION GEN] Generated \(generatedTokens.count) tokens")
|
||
|
||
// Decode response (skip prompt + BOI + IMAGE + EOI)
|
||
let responseStart = textTokens.count + 4
|
||
if generatedTokens.count > responseStart {
|
||
let responseTokens = Array(generatedTokens[responseStart...])
|
||
return tokenizer.decode(tokens: responseTokens)
|
||
}
|
||
|
||
return ""
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Audio Handlers
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
private func processAudioData(_ data: Data) throws -> MTLBuffer {
|
||
print("[AUDIO] Processing audio data: \(data.count) bytes")
|
||
|
||
// Create audio extractor
|
||
let extractor = AudioFeatureExtractor(
|
||
sampleRate: 16000,
|
||
nMels: 128,
|
||
nFft: 400,
|
||
hopLength: 160,
|
||
fMin: 0,
|
||
fMax: 8000
|
||
)
|
||
|
||
// Save data to temp file
|
||
let tempFile = "/tmp/audio_input.wav"
|
||
try data.write(to: URL(fileURLWithPath: tempFile))
|
||
|
||
// Load audio file
|
||
let audioSamples = try extractor.loadAudioFile(url: URL(fileURLWithPath: tempFile))
|
||
print("[AUDIO] Audio samples: \(audioSamples.count)")
|
||
|
||
// Extract mel spectrogram
|
||
let melSpec = extractor.extractMelSpectrogram(from: audioSamples)
|
||
print("[AUDIO] Mel spectrogram: \(melSpec.count) frames x \(melSpec[0].count) mels")
|
||
|
||
// Flatten to [frames, 128]
|
||
let numFrames = melSpec.count
|
||
let melDim = 128
|
||
var audioFeatures = [Float](repeating: 0, count: numFrames * melDim)
|
||
|
||
for frameIdx in 0..<numFrames {
|
||
for melIdx in 0..<melDim {
|
||
audioFeatures[frameIdx * melDim + melIdx] = melSpec[frameIdx][melIdx]
|
||
}
|
||
}
|
||
|
||
// Normalize
|
||
let mean = audioFeatures.reduce(0, +) / Float(audioFeatures.count)
|
||
let std = sqrt(audioFeatures.map { ($0 - mean) * ($0 - mean) }.reduce(0, +) / Float(audioFeatures.count))
|
||
for i in 0..<audioFeatures.count {
|
||
audioFeatures[i] = (audioFeatures[i] - mean) / max(std, 1e-6)
|
||
}
|
||
|
||
print("[AUDIO] Audio features normalized: mean=\(mean), std=\(std)")
|
||
|
||
// Create buffer
|
||
let buffer = engine.device.makeBuffer(bytes: audioFeatures, length: audioFeatures.count * 4)!
|
||
print("[AUDIO] Audio buffer created: \(buffer.length) bytes")
|
||
|
||
return buffer
|
||
}
|
||
|
||
private func generateWithAudio(
|
||
textTokens: [Int],
|
||
audioFeatures: MTLBuffer,
|
||
numFrames: Int,
|
||
maxTokens: Int
|
||
) throws -> String {
|
||
print("[AUDIO GEN] Starting audio-guided generation")
|
||
print("[AUDIO GEN] Text tokens: \(textTokens.count), frames: \(numFrames), max tokens: \(maxTokens)")
|
||
|
||
guard let mm = multimodalModel else {
|
||
print("[AUDIO GEN] Multimodal model not available")
|
||
throw MarkBaseError.audioProcessingFailed
|
||
}
|
||
|
||
// Check if audio tower is available (either AudioTower or AudioTower12B)
|
||
let hasAudioTower = mm.audioTowerFull != nil || mm.audioTower != nil
|
||
if !hasAudioTower {
|
||
print("[AUDIO GEN] Audio tower not available")
|
||
throw MarkBaseError.audioProcessingFailed
|
||
}
|
||
|
||
print("[AUDIO GEN] Audio tower available")
|
||
|
||
// Audio tower forward pass (simplified)
|
||
// Note: Full audio tower implementation would use audioOutputBuffer
|
||
// For now, pool directly from input features
|
||
|
||
// Pool audio features
|
||
let audioPtr = audioFeatures.contents().assumingMemoryBound(to: Float.self)
|
||
let audioLength = audioFeatures.length / 4 // Float size
|
||
var pooled = [Float](repeating: 0, count: 2560)
|
||
|
||
// Simple pooling: average across frames
|
||
let melDim = 128
|
||
for i in 0..<min(2560, melDim * 20) { // Take first 20 mel bands
|
||
var sum: Float = 0
|
||
for frame in 0..<numFrames {
|
||
let idx = frame * melDim + i
|
||
if idx < audioLength {
|
||
sum += audioPtr[idx]
|
||
}
|
||
}
|
||
pooled[i] = sum / Float(numFrames)
|
||
}
|
||
|
||
// Pad remaining
|
||
for i in (melDim * 20)..<2560 {
|
||
pooled[i] = 0
|
||
}
|
||
|
||
// Normalize to magnitude ~5
|
||
let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
print("[AUDIO GEN] Pooled audio magnitude: \(mag)")
|
||
|
||
let scale: Float = 5.0 / max(mag, 1e-6)
|
||
for i in 0..<2560 {
|
||
pooled[i] *= scale
|
||
}
|
||
|
||
let magAfter = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
|
||
print("[AUDIO GEN] Normalized magnitude: \(magAfter)")
|
||
|
||
// Multimodal inference
|
||
let inference = try MultimodalInference(model: mm)
|
||
let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
|
||
|
||
// Note: MultimodalInference needs audio support
|
||
// For now, reuse vision path with audio embedding
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: textTokens,
|
||
precomputedVisionEmbedding: pooledBuffer, // Use audio embedding
|
||
maxTokens: maxTokens
|
||
)
|
||
|
||
print("[AUDIO GEN] Generated \(generatedTokens.count) tokens")
|
||
|
||
// Decode response
|
||
let responseStart = textTokens.count + 4
|
||
if generatedTokens.count > responseStart {
|
||
let responseTokens = Array(generatedTokens[responseStart...])
|
||
return tokenizer.decode(tokens: responseTokens)
|
||
}
|
||
|
||
return ""
|
||
}
|
||
|
||
public func handleChatCompletion(
|
||
messages: [ChatMessage],
|
||
config: GenerationConfig
|
||
) throws -> ChatCompletionResponse {
|
||
let prompt = buildChatPrompt(messages: messages)
|
||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||
|
||
return ChatCompletionResponse(
|
||
id: generateId("chatcmpl"),
|
||
object: "chat.completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
Choice(
|
||
index: 0,
|
||
message: ChatMessage(role: "assistant", content: response),
|
||
finish_reason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: tokenizer.encode(text: prompt).count,
|
||
completionTokens: tokenizer.encode(text: response).count,
|
||
totalTokens: tokenizer.encode(text: prompt + response).count
|
||
)
|
||
)
|
||
}
|
||
|
||
public func handleStreamChatCompletion(
|
||
messages: [ChatMessage],
|
||
config: GenerationConfig
|
||
) throws -> [SSEEvent] {
|
||
let prompt = buildChatPrompt(messages: messages)
|
||
let id = generateId("chatcmpl")
|
||
|
||
// Get tokens for streaming
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
var generatedTokens: [Int] = []
|
||
var position = promptTokens.count
|
||
var lastLogits: [Float] = []
|
||
|
||
// Pre-fill
|
||
for (i, tokenId) in promptTokens.enumerated() {
|
||
lastLogits = try model.forward(tokenId: tokenId, position: i)
|
||
}
|
||
|
||
var events: [SSEEvent] = []
|
||
var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
|
||
|
||
// Stream tokens
|
||
for _ in 0..<config.maxTokens {
|
||
let nextToken = sampler.sample(
|
||
logits: lastLogits,
|
||
temperature: config.temperature,
|
||
topK: config.topK,
|
||
topP: config.topP
|
||
)
|
||
|
||
if tokenizer.eosTokenIds.contains(nextToken) {
|
||
break
|
||
}
|
||
|
||
generatedTokens.append(nextToken)
|
||
let tokenText = streamDecoder.consume(tokenId: nextToken)
|
||
|
||
if !tokenText.isEmpty {
|
||
events.append(SSEStream.chatChunk(
|
||
id: id,
|
||
model: modelId,
|
||
content: tokenText
|
||
))
|
||
}
|
||
|
||
lastLogits = try model.forward(tokenId: nextToken, position: position)
|
||
position += 1
|
||
}
|
||
|
||
// Final chunk
|
||
events.append(SSEStream.chatChunk(
|
||
id: id,
|
||
model: modelId,
|
||
finishReason: "stop"
|
||
))
|
||
|
||
// Done
|
||
events.append(SSEStream.done())
|
||
|
||
return events
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Text Completion Handlers
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public func handleCompletion(
|
||
prompt: String,
|
||
config: GenerationConfig
|
||
) throws -> CompletionResponse {
|
||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||
|
||
return CompletionResponse(
|
||
id: generateId("cmpl"),
|
||
object: "text_completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
CompletionChoice(
|
||
index: 0,
|
||
text: response,
|
||
finishReason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: tokenizer.encode(text: prompt).count,
|
||
completionTokens: tokenizer.encode(text: response).count,
|
||
totalTokens: tokenizer.encode(text: prompt + response).count
|
||
)
|
||
)
|
||
}
|
||
|
||
public func handleStreamCompletion(
|
||
prompt: String,
|
||
config: GenerationConfig
|
||
) throws -> [SSEEvent] {
|
||
let id = generateId("cmpl")
|
||
|
||
// Get tokens
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
var generatedTokens: [Int] = []
|
||
var position = promptTokens.count
|
||
var lastLogits: [Float] = []
|
||
|
||
// Pre-fill
|
||
for (i, tokenId) in promptTokens.enumerated() {
|
||
lastLogits = try model.forward(tokenId: tokenId, position: i)
|
||
}
|
||
|
||
var events: [SSEEvent] = []
|
||
var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
|
||
|
||
// Stream tokens
|
||
for _ in 0..<config.maxTokens {
|
||
let nextToken = sampler.sample(
|
||
logits: lastLogits,
|
||
temperature: config.temperature,
|
||
topK: config.topK,
|
||
topP: config.topP
|
||
)
|
||
|
||
if tokenizer.eosTokenIds.contains(nextToken) {
|
||
break
|
||
}
|
||
|
||
generatedTokens.append(nextToken)
|
||
let tokenText = streamDecoder.consume(tokenId: nextToken)
|
||
|
||
if !tokenText.isEmpty {
|
||
events.append(SSEStream.textChunk(
|
||
id: id,
|
||
model: modelId,
|
||
text: tokenText
|
||
))
|
||
}
|
||
|
||
lastLogits = try model.forward(tokenId: nextToken, position: position)
|
||
position += 1
|
||
}
|
||
|
||
// Final chunk with finish reason
|
||
events.append(SSEStream.textChunk(
|
||
id: id,
|
||
model: modelId,
|
||
text: "",
|
||
finishReason: "stop"
|
||
))
|
||
|
||
// Done
|
||
events.append(SSEStream.done())
|
||
|
||
return events
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Multimodal Handlers
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public func handleMultimodalChat(
|
||
messages: [MultimodalMessage],
|
||
config: GenerationConfig
|
||
) throws -> ChatCompletionResponse {
|
||
// Validate messages
|
||
try Validator.validateMultimodalMessages(messages)
|
||
|
||
// Build prompt with multimodal content
|
||
let prompt = try buildMultimodalPrompt(messages: messages)
|
||
|
||
// Process images/audio if present
|
||
for message in messages {
|
||
for imageUrl in message.imageUrls {
|
||
let _ = try MediaProcessor.loadImage(from: imageUrl.url)
|
||
// TODO: Pass to vision tower
|
||
}
|
||
for audioUrl in message.audioUrls {
|
||
let _ = try MediaProcessor.loadAudio(from: audioUrl.url)
|
||
// TODO: Pass to audio tower
|
||
}
|
||
}
|
||
|
||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||
|
||
return ChatCompletionResponse(
|
||
id: generateId("chatcmpl"),
|
||
object: "chat.completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
Choice(
|
||
index: 0,
|
||
message: ChatMessage(role: "assistant", content: response),
|
||
finish_reason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: tokenizer.encode(text: prompt).count,
|
||
completionTokens: tokenizer.encode(text: response).count,
|
||
totalTokens: tokenizer.encode(text: prompt + response).count
|
||
)
|
||
)
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Utilities
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
private func buildMultimodalPrompt(messages: [MultimodalMessage]) throws -> String {
|
||
var prompt = ""
|
||
|
||
for message in messages {
|
||
switch message.role {
|
||
case "system":
|
||
prompt += "<start_of_turn>user\nSystem: \(message.textContent)<end_of_turn>\n"
|
||
case "user":
|
||
prompt += "<start_of_turn>user\n"
|
||
|
||
// Add image tags
|
||
for _ in message.imageUrls {
|
||
prompt += "<image>\n"
|
||
}
|
||
|
||
// Add audio tags
|
||
for _ in message.audioUrls {
|
||
prompt += "<audio>\n"
|
||
}
|
||
|
||
// Add video tags
|
||
for _ in message.videoUrls {
|
||
prompt += "<video>\n"
|
||
}
|
||
|
||
// Add text
|
||
if !message.textContent.isEmpty {
|
||
prompt += "\(message.textContent)\n"
|
||
}
|
||
|
||
prompt += "<end_of_turn>\n"
|
||
case "assistant":
|
||
prompt += "<start_of_turn>model\n\(message.textContent)<end_of_turn>\n"
|
||
default:
|
||
prompt += "\(message.textContent)\n"
|
||
}
|
||
}
|
||
|
||
prompt += "<start_of_turn>model\n"
|
||
return prompt
|
||
}
|
||
|
||
private func buildChatPrompt(messages: [ChatMessage]) -> String {
|
||
var prompt = ""
|
||
for message in messages {
|
||
let role = message.role == "assistant" ? "model" : message.role
|
||
prompt += "<|turn>\(role)\n\(message.content ?? "")<turn|>\n"
|
||
}
|
||
prompt += "<|turn>model\n"
|
||
return prompt
|
||
}
|
||
|
||
private func generateId(_ prefix: String) -> String {
|
||
let uuid = UUID().uuidString.replacingOccurrences(of: "-", with: "")
|
||
return "\(prefix)-\(uuid.prefix(29))"
|
||
}
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Request Models
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public struct ChatMessage: Codable, Sendable {
|
||
public let role: String
|
||
public let content: String?
|
||
public let tool_calls: [ToolCall]?
|
||
public let name: String?
|
||
|
||
public init(
|
||
role: String,
|
||
content: String? = nil,
|
||
tool_calls: [ToolCall]? = nil,
|
||
name: String? = nil
|
||
) {
|
||
self.role = role
|
||
self.content = content
|
||
self.tool_calls = tool_calls
|
||
self.name = name
|
||
}
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Response Models
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
public struct ChatCompletionResponse: Codable, Sendable {
|
||
public let id: String
|
||
public let object: String
|
||
public let created: Int
|
||
public let model: String
|
||
public let choices: [Choice]
|
||
public let usage: Usage
|
||
}
|
||
|
||
public struct CompletionResponse: Codable, Sendable {
|
||
public let id: String
|
||
public let object: String
|
||
public let created: Int
|
||
public let model: String
|
||
public let choices: [CompletionChoice]
|
||
public let usage: Usage
|
||
}
|
||
|
||
public struct Choice: Codable, Sendable {
|
||
public let index: Int
|
||
public let message: ChatMessage
|
||
public let finish_reason: String
|
||
|
||
public init(
|
||
index: Int,
|
||
message: ChatMessage,
|
||
finish_reason: String
|
||
) {
|
||
self.index = index
|
||
self.message = message
|
||
self.finish_reason = finish_reason
|
||
}
|
||
}
|
||
|
||
public struct CompletionChoice: Codable, Sendable {
|
||
public let index: Int
|
||
public let text: String
|
||
public let finishReason: String
|
||
}
|
||
|
||
public struct Usage: Codable, Sendable {
|
||
public let promptTokens: Int
|
||
public let completionTokens: Int
|
||
public let totalTokens: Int
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Embeddings Handler
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
extension MarkBaseServer {
|
||
/// Get model details
|
||
public func getModelDetails(modelId: String) -> ModelDetails {
|
||
ModelDetails(
|
||
id: modelId,
|
||
capabilities: ModelCapabilities(
|
||
text: true,
|
||
vision: true,
|
||
audio: true,
|
||
embeddings: true,
|
||
streaming: true
|
||
),
|
||
parameters: ModelParameters(
|
||
context_length: maxContextLength,
|
||
num_hidden_layers: model.numHiddenLayers,
|
||
hidden_size: model.hiddenSize,
|
||
vocab_size: model.vocabSize,
|
||
num_attention_heads: 16,
|
||
num_kv_heads: 8
|
||
)
|
||
)
|
||
}
|
||
|
||
/// Handle embeddings request
|
||
public func handleEmbeddings(
|
||
request: EmbeddingsRequest
|
||
) throws -> EmbeddingsResponse {
|
||
var embeddings: [EmbeddingData] = []
|
||
var totalTokens = 0
|
||
|
||
// Process input
|
||
switch request.input {
|
||
case .string(let text):
|
||
let embedding = try generateEmbedding(text: text)
|
||
let tokens = tokenizer.encode(text: text)
|
||
totalTokens += tokens.count
|
||
embeddings.append(EmbeddingData(index: 0, embedding: embedding))
|
||
|
||
case .strings(let texts):
|
||
for (index, text) in texts.enumerated() {
|
||
let embedding = try generateEmbedding(text: text)
|
||
let tokens = tokenizer.encode(text: text)
|
||
totalTokens += tokens.count
|
||
embeddings.append(EmbeddingData(index: index, embedding: embedding))
|
||
}
|
||
|
||
case .tokens(let tokens):
|
||
let text = tokenizer.decode(tokens: tokens)
|
||
let embedding = try generateEmbedding(text: text)
|
||
totalTokens += tokens.count
|
||
embeddings.append(EmbeddingData(index: 0, embedding: embedding))
|
||
|
||
case .tokensList(let tokensList):
|
||
for (index, tokens) in tokensList.enumerated() {
|
||
let text = tokenizer.decode(tokens: tokens)
|
||
let embedding = try generateEmbedding(text: text)
|
||
totalTokens += tokens.count
|
||
embeddings.append(EmbeddingData(index: index, embedding: embedding))
|
||
}
|
||
}
|
||
|
||
return EmbeddingsResponse(
|
||
data: embeddings,
|
||
model: modelId,
|
||
usage: EmbeddingUsage(
|
||
prompt_tokens: totalTokens,
|
||
total_tokens: totalTokens
|
||
)
|
||
)
|
||
}
|
||
|
||
/// Generate embedding for text
|
||
private func generateEmbedding(text: String) throws -> [Float] {
|
||
// Tokenize
|
||
let tokens = tokenizer.encode(text: text)
|
||
|
||
// Run forward pass for all tokens
|
||
var lastHidden: [Float] = []
|
||
for (position, tokenId) in tokens.enumerated() {
|
||
lastHidden = try model.forward(tokenId: tokenId, position: position)
|
||
}
|
||
|
||
// Use final hidden state as embedding
|
||
// For a proper embedding, we should use the mean of all token embeddings
|
||
// but for simplicity, we use the last hidden state
|
||
return lastHidden
|
||
}
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────
|
||
// Video Analysis Handler
|
||
// ─────────────────────────────────────────────────────────────
|
||
|
||
extension MarkBaseServer {
|
||
/// Analyze a video file with frame pooling: frames → per-frame VisionTower → mean pool → single IMAGE token.
|
||
public func handleVideoAnalysis(videoURL: URL, config: GenerationConfig) async throws -> ChatCompletionResponse {
|
||
// 1. Extract frames and audio
|
||
print("Video: Processing \(videoURL.lastPathComponent)...")
|
||
let videoConfig = VideoProcessor.Config(maxFrames: 32, sceneThreshold: 0.15)
|
||
let videoData = try await VideoProcessor.process(url: videoURL, config: videoConfig)
|
||
print(" Frames: \(videoData.frames.count), Audio: \(videoData.audioSamples.count) samples")
|
||
|
||
guard let mm = multimodalModel else {
|
||
throw MarkBaseError.invalidParameter(parameter: "video", message: "Multimodal model not loaded")
|
||
}
|
||
guard mm.visionTowerFull != nil || mm.visionTower != nil else {
|
||
throw MarkBaseError.invalidParameter(parameter: "video", message: "Vision tower not loaded")
|
||
}
|
||
|
||
let hiddenSize = model.hiddenSize
|
||
let patchSize = 16
|
||
let targetSize = 224
|
||
let device = engine.device
|
||
|
||
// 2. Per-frame: patchify → VisionTower → mean pool
|
||
var pooledFrames: [Float] = []
|
||
|
||
for (i, frame) in videoData.frames.enumerated() {
|
||
guard let resized = VideoProcessor.resizePixelBuffer(frame.pixelBuffer,
|
||
targetWidth: targetSize,
|
||
targetHeight: targetSize) else { continue }
|
||
|
||
let (embeddings, numPatches, _) = VideoProcessor.frameToPatchEmbeddings(resized, patchSize: patchSize)
|
||
guard numPatches > 0 else { continue }
|
||
|
||
// Create GPU buffer and run vision tower
|
||
let inputBuf = device.makeBuffer(bytes: embeddings,
|
||
length: embeddings.count * MemoryLayout<Float>.stride)!
|
||
let outputBuf = device.makeBuffer(length: numPatches * hiddenSize * MemoryLayout<Float>.stride)!
|
||
if let vt = mm.visionTowerFull {
|
||
try vt.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outputBuf)
|
||
} else if let vt = mm.visionTower {
|
||
try vt.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outputBuf)
|
||
}
|
||
|
||
// Read back and mean pool across patches
|
||
let outputPtr = outputBuf.contents().assumingMemoryBound(to: Float.self)
|
||
let outputFloats = Array(UnsafeBufferPointer(start: outputPtr, count: numPatches * hiddenSize))
|
||
var frameEmbedding = [Float](repeating: 0, count: hiddenSize)
|
||
for p in 0..<numPatches {
|
||
for h in 0..<hiddenSize {
|
||
frameEmbedding[h] += outputFloats[p * hiddenSize + h]
|
||
}
|
||
}
|
||
for h in 0..<hiddenSize {
|
||
frameEmbedding[h] /= Float(numPatches)
|
||
}
|
||
pooledFrames.append(contentsOf: frameEmbedding)
|
||
print(" Frame \(i): \(numPatches) patches → pooled [\(hiddenSize)]")
|
||
}
|
||
|
||
guard !pooledFrames.isEmpty else {
|
||
throw MarkBaseError.invalidParameter(parameter: "video", message: "No frames processed")
|
||
}
|
||
|
||
// 3. Cross-frame mean pool → single embedding
|
||
let numFrames = videoData.frames.count
|
||
var videoEmbedding = [Float](repeating: 0, count: hiddenSize)
|
||
for f in 0..<numFrames {
|
||
for h in 0..<hiddenSize {
|
||
videoEmbedding[h] += pooledFrames[f * hiddenSize + h]
|
||
}
|
||
}
|
||
for h in 0..<hiddenSize {
|
||
videoEmbedding[h] /= Float(numFrames)
|
||
}
|
||
|
||
let precomputedBuf = device.makeBuffer(bytes: videoEmbedding,
|
||
length: hiddenSize * MemoryLayout<Float>.stride)!
|
||
print(" ✓ Pooled \(numFrames) frames → 1×\(hiddenSize) embedding")
|
||
|
||
// 4. Extract audio mel spectrogram
|
||
var audioFeatures: [[Float]] = []
|
||
if !videoData.audioSamples.isEmpty {
|
||
let extractor = AudioFeatureExtractor()
|
||
audioFeatures = extractor.extractMelSpectrogram(from: videoData.audioSamples)
|
||
print(" Audio frames (mel): \(audioFeatures.count)")
|
||
}
|
||
|
||
// 5. Tokenize prompt
|
||
let prompt = "Describe this video in detail, including visual content and audio."
|
||
let promptTokens = tokenizer.encode(text: prompt)
|
||
|
||
// 6. Run multimodal inference with precomputed embedding
|
||
let inference = try MultimodalInference(model: mm)
|
||
let generatedTokens = try inference.generate(
|
||
textTokens: promptTokens,
|
||
audioFeatures: audioFeatures.isEmpty ? nil : audioFeatures,
|
||
precomputedVisionEmbedding: precomputedBuf,
|
||
maxTokens: config.maxTokens
|
||
)
|
||
|
||
// 7. Decode
|
||
let generatedText = tokenizer.decode(tokens: Array(generatedTokens[promptTokens.count...]))
|
||
|
||
return ChatCompletionResponse(
|
||
id: generateId("chatcmpl"),
|
||
object: "chat.completion",
|
||
created: Int(Date().timeIntervalSince1970),
|
||
model: modelId,
|
||
choices: [
|
||
Choice(
|
||
index: 0,
|
||
message: ChatMessage(role: "assistant", content: generatedText),
|
||
finish_reason: "stop"
|
||
)
|
||
],
|
||
usage: Usage(
|
||
promptTokens: promptTokens.count,
|
||
completionTokens: generatedTokens.count - promptTokens.count,
|
||
totalTokens: generatedTokens.count
|
||
)
|
||
)
|
||
}
|
||
}
|