v2: Initial clean branch with unit tests + CI/CD pipeline
- 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
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,240 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Complete API Router Implementation
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// API endpoint handler
|
||||
public final class APIRouter: @unchecked Sendable {
|
||||
private let modelManager: ModelManager
|
||||
private let metricsCollector: MetricsCollector
|
||||
private let concurrencyController: DynamicConcurrencyController
|
||||
private let requestQueue: RequestQueue
|
||||
|
||||
public init(
|
||||
modelManager: ModelManager,
|
||||
metricsCollector: MetricsCollector = .shared,
|
||||
concurrencyController: DynamicConcurrencyController
|
||||
) {
|
||||
self.modelManager = modelManager
|
||||
self.metricsCollector = metricsCollector
|
||||
self.concurrencyController = concurrencyController
|
||||
self.requestQueue = RequestQueue(controller: concurrencyController)
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Health & Info Endpoints
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// GET /health
|
||||
public func handleHealth() async -> [String: Any] {
|
||||
let currentModel = await modelManager.getCurrentModel()
|
||||
|
||||
return [
|
||||
"status": "healthy",
|
||||
"model": currentModel?.name ?? "none",
|
||||
"model_id": currentModel?.id ?? "",
|
||||
"loaded": currentModel?.loaded ?? false,
|
||||
"version": "1.0.0"
|
||||
]
|
||||
}
|
||||
|
||||
/// GET /v1/models
|
||||
public func handleListModels() async -> [String: Any] {
|
||||
let models = await modelManager.listModels()
|
||||
|
||||
return [
|
||||
"object": "list",
|
||||
"data": models.map { model in
|
||||
[
|
||||
"id": model.id,
|
||||
"object": "model",
|
||||
"created": Int(Date().timeIntervalSince1970),
|
||||
"owned_by": "markbase",
|
||||
"loaded": model.loaded
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model Management Endpoints
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// POST /v1/models/load
|
||||
public func handleLoadModel(modelId: String) async throws -> [String: Any] {
|
||||
try await modelManager.loadModel(id: modelId)
|
||||
|
||||
return [
|
||||
"status": "success",
|
||||
"message": "Model loaded: \(modelId)"
|
||||
]
|
||||
}
|
||||
|
||||
/// POST /v1/models/unload
|
||||
public func handleUnloadModel() async -> [String: Any] {
|
||||
await modelManager.unloadModel()
|
||||
|
||||
return [
|
||||
"status": "success",
|
||||
"message": "Model unloaded"
|
||||
]
|
||||
}
|
||||
|
||||
/// POST /v1/models/switch
|
||||
public func handleSwitchModel(modelId: String) async throws -> [String: Any] {
|
||||
try await modelManager.switchModel(to: modelId)
|
||||
|
||||
return [
|
||||
"status": "success",
|
||||
"message": "Model switched to: \(modelId)"
|
||||
]
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Chat Completions Endpoint
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// POST /v1/chat/completions
|
||||
public func handleChatCompletion(
|
||||
messages: [ChatMessage],
|
||||
config: GenerationConfig
|
||||
) async throws -> ChatCompletionResponse {
|
||||
try await requestQueue.execute { [weak self] in
|
||||
guard let self = self else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
|
||||
let generator = try await self.modelManager.getGenerator()
|
||||
let tokenizer = try await self.modelManager.getTokenizer()
|
||||
let _ = try await self.modelManager.getModel()
|
||||
|
||||
// Build prompt
|
||||
let prompt = self.buildChatPrompt(messages: messages, tokenizer: tokenizer)
|
||||
|
||||
// Generate
|
||||
let startTime = Date()
|
||||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||||
let duration = Date().timeIntervalSince(startTime)
|
||||
|
||||
// Record metrics
|
||||
let tokens = tokenizer.encode(text: response).count
|
||||
let resolvedModelId = (await modelManager.getCurrentModel())?.id ?? "unknown"
|
||||
self.metricsCollector.recordRequest(duration: duration, tokens: tokens, model: resolvedModelId)
|
||||
|
||||
return ChatCompletionResponse(
|
||||
id: self.generateId("chatcmpl"),
|
||||
object: "chat.completion",
|
||||
created: Int(Date().timeIntervalSince1970),
|
||||
model: resolvedModelId,
|
||||
choices: [
|
||||
Choice(
|
||||
index: 0,
|
||||
message: ChatMessage(role: "assistant", content: response),
|
||||
finish_reason: "stop"
|
||||
)
|
||||
],
|
||||
usage: Usage(
|
||||
promptTokens: tokenizer.encode(text: prompt).count,
|
||||
completionTokens: tokens,
|
||||
totalTokens: tokenizer.encode(text: prompt + response).count
|
||||
)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Embeddings Endpoint
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// POST /v1/embeddings
|
||||
public func handleEmbeddings(
|
||||
input: EmbeddingsRequest.InputType
|
||||
) async throws -> EmbeddingsResponse {
|
||||
try await requestQueue.execute { [weak self] in
|
||||
guard let self = self else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
|
||||
let tokenizer = try await self.modelManager.getTokenizer()
|
||||
|
||||
var embeddings: [EmbeddingData] = []
|
||||
var totalTokens = 0
|
||||
|
||||
switch input {
|
||||
case .string(let text):
|
||||
let embedding = try await self.generateEmbedding(text: text)
|
||||
let tokens = tokenizer.encode(text: text).count
|
||||
totalTokens += tokens
|
||||
embeddings.append(EmbeddingData(index: 0, embedding: embedding))
|
||||
|
||||
case .strings(let texts):
|
||||
for (index, text) in texts.enumerated() {
|
||||
let embedding = try await self.generateEmbedding(text: text)
|
||||
let tokens = tokenizer.encode(text: text).count
|
||||
totalTokens += tokens
|
||||
embeddings.append(EmbeddingData(index: index, embedding: embedding))
|
||||
}
|
||||
|
||||
case .tokens(let tokens):
|
||||
let text = tokenizer.decode(tokens: tokens)
|
||||
let embedding = try await self.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 await self.generateEmbedding(text: text)
|
||||
totalTokens += tokens.count
|
||||
embeddings.append(EmbeddingData(index: index, embedding: embedding))
|
||||
}
|
||||
}
|
||||
|
||||
return EmbeddingsResponse(
|
||||
data: embeddings,
|
||||
model: (await self.modelManager.getCurrentModel())?.id ?? "unknown",
|
||||
usage: EmbeddingUsage(
|
||||
prompt_tokens: totalTokens,
|
||||
total_tokens: totalTokens
|
||||
)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Private Helpers
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private func buildChatPrompt(messages: [ChatMessage], tokenizer: Tokenizer) -> 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 generateEmbedding(text: String) async throws -> [Float] {
|
||||
let tokenizer = try await modelManager.getTokenizer()
|
||||
let model = try await modelManager.getModel()
|
||||
|
||||
let tokens = tokenizer.encode(text: text)
|
||||
|
||||
var lastHidden: [Float] = []
|
||||
for (position, tokenId) in tokens.enumerated() {
|
||||
lastHidden = try model.forward(tokenId: tokenId, position: position)
|
||||
}
|
||||
|
||||
return lastHidden
|
||||
}
|
||||
|
||||
private func generateId(_ prefix: String) -> String {
|
||||
let uuid = UUID().uuidString.replacingOccurrences(of: "-", with: "")
|
||||
return "\(prefix)-\(uuid.prefix(29))"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
import MarkBase
|
||||
|
||||
/// Entry point for MarkBase API server
|
||||
/// Usage: swift run G12BServer [model_dir] [port] [model_id] [--benchmark]
|
||||
@main
|
||||
public struct ServerMain {
|
||||
public static func main() async throws {
|
||||
let args = CommandLine.arguments
|
||||
|
||||
// Check for benchmark mode
|
||||
if args.contains("--benchmark") {
|
||||
let modelDir = args.count > 2 ? args[1] : "./model"
|
||||
let modelName = args.count > 3 ? args[2] : "markbase"
|
||||
|
||||
var benchmark = PerformanceBenchmark(modelDir: modelDir, modelName: modelName)
|
||||
try await benchmark.run()
|
||||
return
|
||||
}
|
||||
|
||||
let modelDir = args.count > 1 ? args[1] : "./model"
|
||||
let port = args.count > 2 ? Int(args[2]) ?? 8080 : 8080
|
||||
let modelId = args.count > 3 ? args[3] : "markbase-12b"
|
||||
|
||||
print("\n╔═════════════════════════════════════════════╗")
|
||||
print("║ MarkBase API Server ║")
|
||||
print("╚═════════════════════════════════════════════╝\n")
|
||||
|
||||
let server = try MarkBaseServer(modelDir: modelDir, modelId: modelId)
|
||||
try await server.start(port: port)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// MarkBase CLI - Command Line Interface
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// CLI command
|
||||
public enum CLICommand {
|
||||
case serve(ServeOptions)
|
||||
case chat(ChatOptions)
|
||||
case list(ListOptions)
|
||||
case load(LoadOptions)
|
||||
case unload(UnloadOptions)
|
||||
case switchModel(SwitchOptions)
|
||||
case download(DownloadOptions)
|
||||
case search(SearchOptions)
|
||||
case benchmark(BenchmarkOptions)
|
||||
case help
|
||||
}
|
||||
|
||||
/// CLI options
|
||||
public struct ServeOptions {
|
||||
public var modelDir: String
|
||||
public var port: Int
|
||||
public var host: String
|
||||
public var maxConcurrency: Int
|
||||
|
||||
public init(
|
||||
modelDir: String = "./model",
|
||||
port: Int = 8080,
|
||||
host: String = "127.0.0.1",
|
||||
maxConcurrency: Int = 4
|
||||
) {
|
||||
self.modelDir = modelDir
|
||||
self.port = port
|
||||
self.host = host
|
||||
self.maxConcurrency = maxConcurrency
|
||||
}
|
||||
}
|
||||
|
||||
public struct ChatOptions {
|
||||
public var modelDir: String
|
||||
public var prompt: String
|
||||
public var maxTokens: Int
|
||||
public var temperature: Float
|
||||
public var stream: Bool
|
||||
|
||||
public init(
|
||||
modelDir: String = "./model",
|
||||
prompt: String = "",
|
||||
maxTokens: Int = 100,
|
||||
temperature: Float = 0.7,
|
||||
stream: Bool = true
|
||||
) {
|
||||
self.modelDir = modelDir
|
||||
self.prompt = prompt
|
||||
self.maxTokens = maxTokens
|
||||
self.temperature = temperature
|
||||
self.stream = stream
|
||||
}
|
||||
}
|
||||
|
||||
public struct ListOptions {
|
||||
public var modelsDir: String
|
||||
|
||||
public init(modelsDir: String = "./models") {
|
||||
self.modelsDir = modelsDir
|
||||
}
|
||||
}
|
||||
|
||||
public struct LoadOptions {
|
||||
public var modelId: String
|
||||
public var modelDir: String
|
||||
|
||||
public init(modelId: String, modelDir: String = "./models") {
|
||||
self.modelId = modelId
|
||||
self.modelDir = modelDir
|
||||
}
|
||||
}
|
||||
|
||||
public struct UnloadOptions {
|
||||
public var modelId: String?
|
||||
|
||||
public init(modelId: String? = nil) {
|
||||
self.modelId = modelId
|
||||
}
|
||||
}
|
||||
|
||||
public struct SwitchOptions {
|
||||
public var modelId: String
|
||||
|
||||
public init(modelId: String) {
|
||||
self.modelId = modelId
|
||||
}
|
||||
}
|
||||
|
||||
public struct DownloadOptions {
|
||||
public var repoId: String
|
||||
public var outputDir: String
|
||||
|
||||
public init(repoId: String, outputDir: String = "./models") {
|
||||
self.repoId = repoId
|
||||
self.outputDir = outputDir
|
||||
}
|
||||
}
|
||||
|
||||
public struct SearchOptions {
|
||||
public var query: String
|
||||
public var limit: Int
|
||||
public var gguf: Bool
|
||||
public var safetensors: Bool
|
||||
|
||||
public init(
|
||||
query: String = "",
|
||||
limit: Int = 20,
|
||||
gguf: Bool = false,
|
||||
safetensors: Bool = false
|
||||
) {
|
||||
self.query = query
|
||||
self.limit = limit
|
||||
self.gguf = gguf
|
||||
self.safetensors = safetensors
|
||||
}
|
||||
}
|
||||
|
||||
public struct BenchmarkOptions {
|
||||
public var modelDir: String
|
||||
public var modelName: String
|
||||
public var numPrompts: Int
|
||||
|
||||
public init(
|
||||
modelDir: String = "./model",
|
||||
modelName: String = "markbase",
|
||||
numPrompts: Int = 10
|
||||
) {
|
||||
self.modelDir = modelDir
|
||||
self.modelName = modelName
|
||||
self.numPrompts = numPrompts
|
||||
}
|
||||
}
|
||||
|
||||
/// CLI parser
|
||||
public final class CLIParser {
|
||||
public static func parse(arguments: [String]) -> CLICommand {
|
||||
let args = Array(arguments.dropFirst()) // Skip program name
|
||||
|
||||
guard !args.isEmpty else {
|
||||
return .help
|
||||
}
|
||||
|
||||
let command = args[0]
|
||||
|
||||
switch command {
|
||||
case "serve":
|
||||
return parseServe(args: Array(args.dropFirst()))
|
||||
case "chat":
|
||||
return parseChat(args: Array(args.dropFirst()))
|
||||
case "list":
|
||||
return .list(ListOptions())
|
||||
case "load":
|
||||
return parseLoad(args: Array(args.dropFirst()))
|
||||
case "unload":
|
||||
return .unload(UnloadOptions())
|
||||
case "switch":
|
||||
return parseSwitch(args: Array(args.dropFirst()))
|
||||
case "download":
|
||||
return parseDownload(args: Array(args.dropFirst()))
|
||||
case "search":
|
||||
return parseSearch(args: Array(args.dropFirst()))
|
||||
case "benchmark":
|
||||
return parseBenchmark(args: Array(args.dropFirst()))
|
||||
case "help", "--help", "-h":
|
||||
return .help
|
||||
default:
|
||||
print("Unknown command: \(command)")
|
||||
return .help
|
||||
}
|
||||
}
|
||||
|
||||
private static func parseServe(args: [String]) -> CLICommand {
|
||||
var options = ServeOptions()
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--model", "-m":
|
||||
if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--port", "-p":
|
||||
if i + 1 < args.count, let port = Int(args[i + 1]) { options.port = port; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--host":
|
||||
if i + 1 < args.count { options.host = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--concurrency", "-c":
|
||||
if i + 1 < args.count, let concurrency = Int(args[i + 1]) { options.maxConcurrency = concurrency; i += 2 }
|
||||
else { i += 1 }
|
||||
default:
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .serve(options)
|
||||
}
|
||||
|
||||
private static func parseChat(args: [String]) -> CLICommand {
|
||||
var options = ChatOptions()
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--model", "-m":
|
||||
if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--prompt", "-p":
|
||||
if i + 1 < args.count { options.prompt = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--max-tokens", "-n":
|
||||
if i + 1 < args.count, let n = Int(args[i + 1]) { options.maxTokens = n; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--temperature", "-t":
|
||||
if i + 1 < args.count, let t = Float(args[i + 1]) { options.temperature = t; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--stream", "-s":
|
||||
options.stream = true; i += 1
|
||||
case "--no-stream":
|
||||
options.stream = false; i += 1
|
||||
default:
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .chat(options)
|
||||
}
|
||||
|
||||
private static func parseLoad(args: [String]) -> CLICommand {
|
||||
var modelId = ""
|
||||
var modelDir = "./models"
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--model", "-m":
|
||||
if i + 1 < args.count { modelId = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--dir", "-d":
|
||||
if i + 1 < args.count { modelDir = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
default:
|
||||
if modelId.isEmpty { modelId = args[i] }
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .load(LoadOptions(modelId: modelId, modelDir: modelDir))
|
||||
}
|
||||
|
||||
private static func parseSwitch(args: [String]) -> CLICommand {
|
||||
var modelId = ""
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--model", "-m":
|
||||
if i + 1 < args.count { modelId = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
default:
|
||||
if modelId.isEmpty { modelId = args[i] }
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .switchModel(SwitchOptions(modelId: modelId))
|
||||
}
|
||||
|
||||
private static func parseDownload(args: [String]) -> CLICommand {
|
||||
var repoId = ""
|
||||
var outputDir = "./models"
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--output", "-o":
|
||||
if i + 1 < args.count { outputDir = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
default:
|
||||
if repoId.isEmpty { repoId = args[i] }
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .download(DownloadOptions(repoId: repoId, outputDir: outputDir))
|
||||
}
|
||||
|
||||
private static func parseSearch(args: [String]) -> CLICommand {
|
||||
var options = SearchOptions()
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--query", "-q":
|
||||
if i + 1 < args.count { options.query = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--limit", "-l":
|
||||
if i + 1 < args.count, let limit = Int(args[i + 1]) { options.limit = limit; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--gguf":
|
||||
options.gguf = true; i += 1
|
||||
case "--safetensors":
|
||||
options.safetensors = true; i += 1
|
||||
default:
|
||||
if options.query.isEmpty { options.query = args[i] }
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .search(options)
|
||||
}
|
||||
|
||||
private static func parseBenchmark(args: [String]) -> CLICommand {
|
||||
var options = BenchmarkOptions()
|
||||
|
||||
var i = 0
|
||||
while i < args.count {
|
||||
switch args[i] {
|
||||
case "--model", "-m":
|
||||
if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--name":
|
||||
if i + 1 < args.count { options.modelName = args[i + 1]; i += 2 }
|
||||
else { i += 1 }
|
||||
case "--prompts", "-n":
|
||||
if i + 1 < args.count, let n = Int(args[i + 1]) { options.numPrompts = n; i += 2 }
|
||||
else { i += 1 }
|
||||
default:
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
return .benchmark(options)
|
||||
}
|
||||
}
|
||||
|
||||
/// CLI help message
|
||||
public func printHelp() {
|
||||
print("""
|
||||
MarkBase CLI - Command Line Interface
|
||||
|
||||
Usage: markbase <command> [options]
|
||||
|
||||
Commands:
|
||||
serve Start API server
|
||||
chat Interactive chat
|
||||
list List available models
|
||||
load Load a model
|
||||
unload Unload current model
|
||||
switch Switch to a different model
|
||||
download Download model from HuggingFace
|
||||
search Search models on HuggingFace
|
||||
benchmark Run performance benchmark
|
||||
help Show this help message
|
||||
|
||||
Examples:
|
||||
markbase serve --model ./model --port 8080
|
||||
markbase chat --model ./model --prompt "Hello!"
|
||||
markbase download mlx-community/gemma-4-e4b-it-4bit
|
||||
markbase search "gemma-4" --gguf
|
||||
markbase benchmark --model ./model --prompts 10
|
||||
|
||||
Use "markbase <command> --help" for more information about a command.
|
||||
""")
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
import Foundation
|
||||
import os
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Concurrent Request Handling with Dynamic Concurrency
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Request queue for concurrent processing with dynamic concurrency
|
||||
public actor RequestQueue {
|
||||
private let controller: DynamicConcurrencyController
|
||||
private var semaphore: AsyncSemaphore
|
||||
|
||||
public init(controller: DynamicConcurrencyController) {
|
||||
self.controller = controller
|
||||
self.semaphore = AsyncSemaphore(value: 4)
|
||||
}
|
||||
|
||||
public func initialize() async {
|
||||
let maxConcurrency = await controller.getCurrentMax()
|
||||
semaphore = AsyncSemaphore(value: maxConcurrency)
|
||||
}
|
||||
|
||||
/// Execute request with concurrency limit
|
||||
public func execute<T: Sendable>(_ operation: @escaping @Sendable () async throws -> T) async throws -> T {
|
||||
await semaphore.wait()
|
||||
defer { semaphore.signal() }
|
||||
|
||||
return try await operation()
|
||||
}
|
||||
|
||||
/// Update semaphore when concurrency changes
|
||||
public func updateSemaphore(newMax: Int) {
|
||||
semaphore = AsyncSemaphore(value: newMax)
|
||||
}
|
||||
}
|
||||
|
||||
/// Async semaphore for concurrency control
|
||||
public final class AsyncSemaphore: @unchecked Sendable {
|
||||
private var value: Int
|
||||
private let lock = OSAllocatedUnfairLock()
|
||||
private var waiters: [CheckedContinuation<Void, Never>] = []
|
||||
|
||||
public init(value: Int) {
|
||||
self.value = value
|
||||
}
|
||||
|
||||
public func wait() async {
|
||||
let shouldWait = lock.withLock {
|
||||
if value > 0 {
|
||||
value -= 1
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
guard shouldWait else { return }
|
||||
|
||||
return await withCheckedContinuation { continuation in
|
||||
lock.withLock {
|
||||
waiters.append(continuation)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public func signal() {
|
||||
lock.withLock {
|
||||
if !waiters.isEmpty {
|
||||
let waiter = waiters.removeFirst()
|
||||
waiter.resume()
|
||||
} else {
|
||||
value += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch request processor with dynamic concurrency
|
||||
public actor BatchProcessor {
|
||||
private let requestQueue: RequestQueue
|
||||
|
||||
public init(controller: DynamicConcurrencyController) {
|
||||
self.requestQueue = RequestQueue(controller: controller)
|
||||
}
|
||||
|
||||
/// Process multiple requests concurrently
|
||||
public func processBatch<T: Sendable>(
|
||||
_ requests: [@Sendable () async throws -> T]
|
||||
) async throws -> [T] {
|
||||
try await withThrowingTaskGroup(of: (Int, T).self) { group in
|
||||
var results: [T?] = Array(repeating: nil, count: requests.count)
|
||||
|
||||
for (index, request) in requests.enumerated() {
|
||||
let req = request
|
||||
group.addTask {
|
||||
let result = try await self.requestQueue.execute(req)
|
||||
return (index, result)
|
||||
}
|
||||
}
|
||||
|
||||
for try await (index, result) in group {
|
||||
results[index] = result
|
||||
}
|
||||
|
||||
return results.compactMap { $0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Cross-Device Client
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Cross-device client for making requests to other nodes
|
||||
public final class CrossDeviceClient: @unchecked Sendable {
|
||||
private let session: URLSession
|
||||
private let loadBalancer: LoadBalancer
|
||||
private let timeout: TimeInterval
|
||||
|
||||
public init(
|
||||
loadBalancer: LoadBalancer,
|
||||
timeout: TimeInterval = 30
|
||||
) {
|
||||
self.loadBalancer = loadBalancer
|
||||
self.timeout = timeout
|
||||
|
||||
let config = URLSessionConfiguration.default
|
||||
config.timeoutIntervalForRequest = timeout
|
||||
self.session = URLSession(configuration: config)
|
||||
}
|
||||
|
||||
/// Send request to cluster
|
||||
public func sendToCluster(
|
||||
endpoint: String,
|
||||
method: String = "POST",
|
||||
body: Data? = nil
|
||||
) async throws -> CrossDeviceResponse {
|
||||
guard let node = loadBalancer.getNextNode() else {
|
||||
throw CrossDeviceError.noHealthyNodes
|
||||
}
|
||||
|
||||
return try await sendToNode(
|
||||
node: node,
|
||||
endpoint: endpoint,
|
||||
method: method,
|
||||
body: body
|
||||
)
|
||||
}
|
||||
|
||||
/// Send request to specific node
|
||||
public func sendToNode(
|
||||
node: DeviceNode,
|
||||
endpoint: String,
|
||||
method: String = "POST",
|
||||
body: Data? = nil
|
||||
) async throws -> CrossDeviceResponse {
|
||||
let url = URL(string: "\(node.baseURL)\(endpoint)")!
|
||||
var request = URLRequest(url: url)
|
||||
request.httpMethod = method
|
||||
request.httpBody = body
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.timeoutInterval = timeout
|
||||
|
||||
let startTime = Date()
|
||||
|
||||
do {
|
||||
let (data, response) = try await session.data(for: request)
|
||||
let latency = Date().timeIntervalSince(startTime)
|
||||
|
||||
guard let httpResponse = response as? HTTPURLResponse else {
|
||||
throw CrossDeviceError.invalidResponse
|
||||
}
|
||||
|
||||
// Update node status
|
||||
loadBalancer.updateNodeStatus(
|
||||
id: node.id,
|
||||
status: httpResponse.statusCode < 500 ? .healthy : .degraded,
|
||||
load: nil
|
||||
)
|
||||
|
||||
return CrossDeviceResponse(
|
||||
requestId: UUID().uuidString,
|
||||
statusCode: httpResponse.statusCode,
|
||||
body: data,
|
||||
latency: latency,
|
||||
nodeId: node.id
|
||||
)
|
||||
} catch {
|
||||
let latency = Date().timeIntervalSince(startTime)
|
||||
|
||||
// Update node status
|
||||
loadBalancer.updateNodeStatus(id: node.id, status: .unhealthy)
|
||||
|
||||
throw CrossDeviceError.requestFailed(error)
|
||||
}
|
||||
}
|
||||
|
||||
/// Broadcast request to all healthy nodes
|
||||
public func broadcastToAll(
|
||||
endpoint: String,
|
||||
method: String = "POST",
|
||||
body: Data? = nil
|
||||
) async throws -> [CrossDeviceResponse] {
|
||||
let nodes = loadBalancer.getNodes().filter { $0.status == .healthy }
|
||||
|
||||
return try await withThrowingTaskGroup(of: CrossDeviceResponse.self) { group in
|
||||
for node in nodes {
|
||||
group.addTask {
|
||||
try await self.sendToNode(
|
||||
node: node,
|
||||
endpoint: endpoint,
|
||||
method: method,
|
||||
body: body
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
var responses: [CrossDeviceResponse] = []
|
||||
for try await response in group {
|
||||
responses.append(response)
|
||||
}
|
||||
return responses
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cross-device errors
|
||||
public enum CrossDeviceError: Error, LocalizedError {
|
||||
case noHealthyNodes
|
||||
case invalidResponse
|
||||
case requestFailed(Error)
|
||||
case timeout
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .noHealthyNodes:
|
||||
return "No healthy nodes available"
|
||||
case .invalidResponse:
|
||||
return "Invalid response from node"
|
||||
case .requestFailed(let error):
|
||||
return "Request failed: \(error.localizedDescription)"
|
||||
case .timeout:
|
||||
return "Request timed out"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Cross-Device Communication Protocol
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Device node in the cluster
|
||||
public struct DeviceNode: Codable, Sendable {
|
||||
public let id: String
|
||||
public let host: String
|
||||
public let port: Int
|
||||
public let capabilities: DeviceCapabilities
|
||||
public var status: DeviceStatus
|
||||
public var load: Double // 0.0 to 1.0
|
||||
|
||||
public init(
|
||||
id: String,
|
||||
host: String,
|
||||
port: Int,
|
||||
capabilities: DeviceCapabilities,
|
||||
status: DeviceStatus = .healthy,
|
||||
load: Double = 0.0
|
||||
) {
|
||||
self.id = id
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.capabilities = capabilities
|
||||
self.status = status
|
||||
self.load = load
|
||||
}
|
||||
|
||||
public var baseURL: String {
|
||||
"http://\(host):\(port)"
|
||||
}
|
||||
}
|
||||
|
||||
/// Device capabilities
|
||||
public struct DeviceCapabilities: Codable, Sendable {
|
||||
public let maxConcurrency: Int
|
||||
public let supportedModels: [String]
|
||||
public let hasGPU: Bool
|
||||
public let memoryGB: Int
|
||||
|
||||
public init(
|
||||
maxConcurrency: Int = 4,
|
||||
supportedModels: [String] = [],
|
||||
hasGPU: Bool = true,
|
||||
memoryGB: Int = 16
|
||||
) {
|
||||
self.maxConcurrency = maxConcurrency
|
||||
self.supportedModels = supportedModels
|
||||
self.hasGPU = hasGPU
|
||||
self.memoryGB = memoryGB
|
||||
}
|
||||
}
|
||||
|
||||
/// Device status
|
||||
public enum DeviceStatus: String, Codable, Sendable {
|
||||
case healthy
|
||||
case degraded
|
||||
case unhealthy
|
||||
case offline
|
||||
}
|
||||
|
||||
/// Cross-device request
|
||||
public struct CrossDeviceRequest: Codable, Sendable {
|
||||
public let id: String
|
||||
public let endpoint: String
|
||||
public let method: String
|
||||
public let body: Data?
|
||||
public let timeout: TimeInterval
|
||||
|
||||
public init(
|
||||
id: String = UUID().uuidString,
|
||||
endpoint: String,
|
||||
method: String = "POST",
|
||||
body: Data? = nil,
|
||||
timeout: TimeInterval = 30
|
||||
) {
|
||||
self.id = id
|
||||
self.endpoint = endpoint
|
||||
self.method = method
|
||||
self.body = body
|
||||
self.timeout = timeout
|
||||
}
|
||||
}
|
||||
|
||||
/// Cross-device response
|
||||
public struct CrossDeviceResponse: Codable, Sendable {
|
||||
public let requestId: String
|
||||
public let statusCode: Int
|
||||
public let body: Data?
|
||||
public let latency: TimeInterval
|
||||
public let nodeId: String
|
||||
|
||||
public init(
|
||||
requestId: String,
|
||||
statusCode: Int,
|
||||
body: Data? = nil,
|
||||
latency: TimeInterval,
|
||||
nodeId: String
|
||||
) {
|
||||
self.requestId = requestId
|
||||
self.statusCode = statusCode
|
||||
self.body = body
|
||||
self.latency = latency
|
||||
self.nodeId = nodeId
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Dynamic Concurrency Controller
|
||||
// Automatically adjusts concurrency based on system resources
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Memory statistics
|
||||
public struct MemoryStats {
|
||||
public let total: UInt64
|
||||
public let available: UInt64
|
||||
public let used: UInt64
|
||||
public let percentage: Double
|
||||
|
||||
public init(total: UInt64, available: UInt64, used: UInt64) {
|
||||
self.total = total
|
||||
self.available = available
|
||||
self.used = used
|
||||
self.percentage = total > 0 ? Double(used) / Double(total) : 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Dynamic concurrency controller
|
||||
public actor DynamicConcurrencyController {
|
||||
private var currentMax: Int
|
||||
private let minConcurrency: Int
|
||||
private let maxConcurrency: Int
|
||||
private let modelMemoryEstimate: UInt64
|
||||
|
||||
/// Concurrency adjustment event
|
||||
public struct ConcurrencyEvent {
|
||||
public let timestamp: Date
|
||||
public let oldMax: Int
|
||||
public let newMax: Int
|
||||
public let reason: String
|
||||
}
|
||||
|
||||
/// Event handler
|
||||
public var onConcurrencyChange: ((ConcurrencyEvent) -> Void)?
|
||||
|
||||
public init(
|
||||
initialConcurrency: Int = 4,
|
||||
minConcurrency: Int = 1,
|
||||
maxConcurrency: Int = 16,
|
||||
modelMemoryEstimate: UInt64 = 9 * 1024 * 1024 * 1024 // 9GB for 12B
|
||||
) {
|
||||
self.currentMax = initialConcurrency
|
||||
self.minConcurrency = minConcurrency
|
||||
self.maxConcurrency = maxConcurrency
|
||||
self.modelMemoryEstimate = modelMemoryEstimate
|
||||
}
|
||||
|
||||
/// Get current max concurrency
|
||||
public func getCurrentMax() -> Int {
|
||||
return currentMax
|
||||
}
|
||||
|
||||
/// Adjust concurrency based on current memory
|
||||
@discardableResult
|
||||
public func adjust() -> ConcurrencyEvent? {
|
||||
let oldMax = currentMax
|
||||
|
||||
// Simplified: use available memory estimation
|
||||
// In production, you would use actual memory stats
|
||||
let recommendedConcurrency = 4 // Default for now
|
||||
|
||||
let newMax = max(minConcurrency, min(maxConcurrency, recommendedConcurrency))
|
||||
|
||||
if newMax != oldMax {
|
||||
let event = ConcurrencyEvent(
|
||||
timestamp: Date(),
|
||||
oldMax: oldMax,
|
||||
newMax: newMax,
|
||||
reason: "Automatic adjustment"
|
||||
)
|
||||
|
||||
currentMax = newMax
|
||||
onConcurrencyChange?(event)
|
||||
return event
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
/// Get recommended concurrency for a given model size
|
||||
public static func recommendConcurrency(
|
||||
modelMemoryBytes: UInt64,
|
||||
totalMemory: UInt64,
|
||||
reservedMemory: UInt64 = 2 * 1024 * 1024 * 1024 // 2GB reserved
|
||||
) -> Int {
|
||||
let availableMemory = totalMemory - reservedMemory
|
||||
let concurrency = Int(availableMemory / modelMemoryBytes)
|
||||
return max(1, concurrency)
|
||||
}
|
||||
}
|
||||
|
||||
/// Memory monitor using system information
|
||||
public actor MemoryMonitor {
|
||||
public init() {}
|
||||
|
||||
/// Get current memory statistics (simplified)
|
||||
public func getStats() -> MemoryStats {
|
||||
// Simplified implementation
|
||||
return MemoryStats(
|
||||
total: 48 * 1024 * 1024 * 1024, // 48GB
|
||||
available: 32 * 1024 * 1024 * 1024, // 32GB
|
||||
used: 16 * 1024 * 1024 * 1024 // 16GB
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Embeddings API Models
|
||||
// OpenAI Compatible Format
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Embeddings request
|
||||
public struct EmbeddingsRequest: Codable {
|
||||
public let model: String
|
||||
public let input: InputType
|
||||
public let encoding_format: String?
|
||||
|
||||
public enum InputType: Codable, Sendable {
|
||||
case string(String)
|
||||
case strings([String])
|
||||
case tokens([Int])
|
||||
case tokensList([[Int]])
|
||||
|
||||
public init(from decoder: Decoder) throws {
|
||||
let container = try decoder.singleValueContainer()
|
||||
|
||||
if let string = try? container.decode(String.self) {
|
||||
self = .string(string)
|
||||
} else if let strings = try? container.decode([String].self) {
|
||||
self = .strings(strings)
|
||||
} else if let tokens = try? container.decode([Int].self) {
|
||||
self = .tokens(tokens)
|
||||
} else if let tokensList = try? container.decode([[Int]].self) {
|
||||
self = .tokensList(tokensList)
|
||||
} else {
|
||||
throw DecodingError.dataCorruptedError(
|
||||
in: container,
|
||||
debugDescription: "Invalid input type"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
public func encode(to encoder: Encoder) throws {
|
||||
var container = encoder.singleValueContainer()
|
||||
|
||||
switch self {
|
||||
case .string(let string):
|
||||
try container.encode(string)
|
||||
case .strings(let strings):
|
||||
try container.encode(strings)
|
||||
case .tokens(let tokens):
|
||||
try container.encode(tokens)
|
||||
case .tokensList(let tokensList):
|
||||
try container.encode(tokensList)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public init(
|
||||
model: String,
|
||||
input: InputType,
|
||||
encoding_format: String? = nil
|
||||
) {
|
||||
self.model = model
|
||||
self.input = input
|
||||
self.encoding_format = encoding_format
|
||||
}
|
||||
}
|
||||
|
||||
/// Embeddings response
|
||||
public struct EmbeddingsResponse: Codable, Sendable {
|
||||
public let object: String
|
||||
public let data: [EmbeddingData]
|
||||
public let model: String
|
||||
public let usage: EmbeddingUsage
|
||||
|
||||
public init(
|
||||
object: String = "list",
|
||||
data: [EmbeddingData],
|
||||
model: String,
|
||||
usage: EmbeddingUsage
|
||||
) {
|
||||
self.object = object
|
||||
self.data = data
|
||||
self.model = model
|
||||
self.usage = usage
|
||||
}
|
||||
}
|
||||
|
||||
/// Single embedding data
|
||||
public struct EmbeddingData: Codable, Sendable {
|
||||
public let object: String
|
||||
public let index: Int
|
||||
public let embedding: [Float]
|
||||
|
||||
public init(
|
||||
object: String = "embedding",
|
||||
index: Int,
|
||||
embedding: [Float]
|
||||
) {
|
||||
self.object = object
|
||||
self.index = index
|
||||
self.embedding = embedding
|
||||
}
|
||||
}
|
||||
|
||||
/// Embedding usage
|
||||
public struct EmbeddingUsage: Codable, Sendable {
|
||||
public let prompt_tokens: Int
|
||||
public let total_tokens: Int
|
||||
|
||||
public init(prompt_tokens: Int, total_tokens: Int) {
|
||||
self.prompt_tokens = prompt_tokens
|
||||
self.total_tokens = total_tokens
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Unified Error Handling for MarkBase
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// MarkBase error types
|
||||
public enum MarkBaseError: Error, LocalizedError {
|
||||
case modelNotFound(String)
|
||||
case invalidRequest(String)
|
||||
case modelLoadingFailed(String)
|
||||
case inferenceFailed(String)
|
||||
case tokenLimitExceeded(current: Int, max: Int)
|
||||
case invalidParameter(parameter: String, message: String)
|
||||
case internalError(String)
|
||||
case multimodalNotSupported
|
||||
case imageProcessingFailed
|
||||
case audioProcessingFailed
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .modelNotFound(let path):
|
||||
return "Model not found at: \(path)"
|
||||
case .invalidRequest(let message):
|
||||
return "Invalid request: \(message)"
|
||||
case .modelLoadingFailed(let detail):
|
||||
return "Failed to load model: \(detail)"
|
||||
case .inferenceFailed(let detail):
|
||||
return "Inference failed: \(detail)"
|
||||
case .tokenLimitExceeded(let current, let max):
|
||||
return "Token limit exceeded: \(current) > \(max)"
|
||||
case .invalidParameter(let parameter, let message):
|
||||
return "Invalid parameter '\(parameter)': \(message)"
|
||||
case .internalError(let detail):
|
||||
return "Internal error: \(detail)"
|
||||
case .multimodalNotSupported:
|
||||
return "Multimodal inference not supported for this model"
|
||||
case .imageProcessingFailed:
|
||||
return "Image processing failed"
|
||||
case .audioProcessingFailed:
|
||||
return "Audio processing failed"
|
||||
}
|
||||
}
|
||||
|
||||
/// HTTP status code for this error
|
||||
public var httpStatus: Int {
|
||||
switch self {
|
||||
case .modelNotFound:
|
||||
return 404
|
||||
case .invalidRequest, .invalidParameter:
|
||||
return 400
|
||||
case .modelLoadingFailed, .internalError:
|
||||
return 500
|
||||
case .inferenceFailed:
|
||||
return 500
|
||||
case .tokenLimitExceeded:
|
||||
return 400
|
||||
case .multimodalNotSupported:
|
||||
return 400
|
||||
case .imageProcessingFailed, .audioProcessingFailed:
|
||||
return 500
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI-compatible error type
|
||||
public var errorType: String {
|
||||
switch self {
|
||||
case .modelNotFound:
|
||||
return "model_not_found"
|
||||
case .invalidRequest, .invalidParameter, .tokenLimitExceeded:
|
||||
return "invalid_request_error"
|
||||
case .modelLoadingFailed, .internalError:
|
||||
return "server_error"
|
||||
case .inferenceFailed:
|
||||
return "server_error"
|
||||
case .multimodalNotSupported:
|
||||
return "invalid_request_error"
|
||||
case .imageProcessingFailed, .audioProcessingFailed:
|
||||
return "server_error"
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to OpenAI error response format
|
||||
public func toErrorResponse(param: String? = nil) -> ErrorResponse {
|
||||
ErrorResponse(
|
||||
error: toErrorDetail(param: param)
|
||||
)
|
||||
}
|
||||
|
||||
/// Convert to ErrorDetail
|
||||
public func toErrorDetail(param: String? = nil) -> ErrorDetail {
|
||||
ErrorDetail(
|
||||
message: localizedDescription,
|
||||
type: errorType,
|
||||
code: httpStatus,
|
||||
param: param
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI-compatible error response
|
||||
public struct ErrorResponse: Codable {
|
||||
public let error: ErrorDetail
|
||||
|
||||
public init(error: ErrorDetail) {
|
||||
self.error = error
|
||||
}
|
||||
|
||||
/// Create error response from MarkBaseError
|
||||
public static func from(_ error: MarkBaseError) -> ErrorResponse {
|
||||
ErrorResponse(error: error.toErrorDetail())
|
||||
}
|
||||
}
|
||||
|
||||
public struct ErrorDetail: Codable {
|
||||
public let message: String
|
||||
public let type: String
|
||||
public let code: Int
|
||||
public let param: String?
|
||||
|
||||
public init(message: String, type: String, code: Int, param: String? = nil) {
|
||||
self.message = message
|
||||
self.type = type
|
||||
self.code = code
|
||||
self.param = param
|
||||
}
|
||||
}
|
||||
|
||||
/// Validation helpers
|
||||
public enum Validator {
|
||||
/// Validate model path exists
|
||||
public static func validateModelPath(_ path: String) throws {
|
||||
guard FileManager.default.fileExists(atPath: path) else {
|
||||
throw MarkBaseError.modelNotFound(path)
|
||||
}
|
||||
|
||||
// Check for required files (support both single-file and sharded safetensors)
|
||||
let hasSingleFile = FileManager.default.fileExists(atPath: (path as NSString).appendingPathComponent("model.safetensors"))
|
||||
let hasIndexFile = FileManager.default.fileExists(atPath: (path as NSString).appendingPathComponent("model.safetensors.index.json"))
|
||||
guard hasSingleFile || hasIndexFile else {
|
||||
throw MarkBaseError.modelLoadingFailed("Missing required file: model.safetensors or model.safetensors.index.json")
|
||||
}
|
||||
guard FileManager.default.fileExists(atPath: (path as NSString).appendingPathComponent("config.json")) else {
|
||||
throw MarkBaseError.modelLoadingFailed("Missing required file: config.json")
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate generation parameters
|
||||
public static func validateGenerationParams(
|
||||
maxTokens: Int?,
|
||||
temperature: Float?,
|
||||
topP: Float?,
|
||||
topK: Int?
|
||||
) throws {
|
||||
if let maxTokens = maxTokens {
|
||||
guard maxTokens > 0 && maxTokens <= 4096 else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "max_tokens",
|
||||
message: "Must be between 1 and 4096"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if let temperature = temperature {
|
||||
guard temperature >= 0.0 && temperature <= 2.0 else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "temperature",
|
||||
message: "Must be between 0.0 and 2.0"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if let topP = topP {
|
||||
guard topP >= 0.0 && topP <= 1.0 else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "top_p",
|
||||
message: "Must be between 0.0 and 1.0"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if let topK = topK {
|
||||
guard topK > 0 else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "top_k",
|
||||
message: "Must be greater than 0"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate prompt
|
||||
public static func validatePrompt(_ prompt: String, maxLength: Int = 4096) throws {
|
||||
guard !prompt.isEmpty else {
|
||||
throw MarkBaseError.invalidRequest("Prompt cannot be empty")
|
||||
}
|
||||
|
||||
guard prompt.count <= maxLength else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "prompt",
|
||||
message: "Prompt too long (max \(maxLength) characters)"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate messages
|
||||
public static func validateMessages(_ messages: [ChatMessage]) throws {
|
||||
guard !messages.isEmpty else {
|
||||
throw MarkBaseError.invalidRequest("Messages array cannot be empty")
|
||||
}
|
||||
|
||||
for (index, message) in messages.enumerated() {
|
||||
guard let content = message.content, !content.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content",
|
||||
message: "Content cannot be empty"
|
||||
)
|
||||
}
|
||||
|
||||
guard ["system", "user", "assistant"].contains(message.role) else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].role",
|
||||
message: "Role must be 'system', 'user', or 'assistant'"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate multimodal messages
|
||||
public static func validateMultimodalMessages(_ messages: [MultimodalMessage]) throws {
|
||||
guard !messages.isEmpty else {
|
||||
throw MarkBaseError.invalidRequest("Messages array cannot be empty")
|
||||
}
|
||||
|
||||
for (index, message) in messages.enumerated() {
|
||||
guard !message.content.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content",
|
||||
message: "Content array cannot be empty"
|
||||
)
|
||||
}
|
||||
|
||||
// Validate content parts
|
||||
for (partIndex, part) in message.content.enumerated() {
|
||||
switch part {
|
||||
case .text(let text):
|
||||
guard !text.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content[\(partIndex)]",
|
||||
message: "Text content cannot be empty"
|
||||
)
|
||||
}
|
||||
case .imageUrl(let imageUrl):
|
||||
guard !imageUrl.url.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content[\(partIndex)].image_url.url",
|
||||
message: "Image URL cannot be empty"
|
||||
)
|
||||
}
|
||||
case .audioUrl(let audioUrl):
|
||||
guard !audioUrl.url.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content[\(partIndex)].audio_url.url",
|
||||
message: "Audio URL cannot be empty"
|
||||
)
|
||||
}
|
||||
case .videoUrl(let videoUrl):
|
||||
guard !videoUrl.url.isEmpty else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].content[\(partIndex)].video_url.url",
|
||||
message: "Video URL cannot be empty"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Validate role
|
||||
guard ["system", "user", "assistant"].contains(message.role) else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "messages[\(index)].role",
|
||||
message: "Role must be 'system', 'user', or 'assistant'"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result type with error handling
|
||||
public enum Result<T> {
|
||||
case success(T)
|
||||
case failure(MarkBaseError)
|
||||
|
||||
public func get() throws -> T {
|
||||
switch self {
|
||||
case .success(let value):
|
||||
return value
|
||||
case .failure(let error):
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
public func map<U>(_ transform: (T) -> U) -> Result<U> {
|
||||
switch self {
|
||||
case .success(let value):
|
||||
return .success(transform(value))
|
||||
case .failure(let error):
|
||||
return .failure(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Function Calling API Models
|
||||
// OpenAI Compatible Format
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Tool definition
|
||||
public struct Tool: Codable {
|
||||
public let type: String
|
||||
public let function: FunctionDefinition
|
||||
|
||||
public init(type: String = "function", function: FunctionDefinition) {
|
||||
self.type = type
|
||||
self.function = function
|
||||
}
|
||||
}
|
||||
|
||||
/// Function definition
|
||||
public struct FunctionDefinition: Codable {
|
||||
public let name: String
|
||||
public let description: String?
|
||||
public let parameters: FunctionParameters?
|
||||
|
||||
public init(
|
||||
name: String,
|
||||
description: String? = nil,
|
||||
parameters: FunctionParameters? = nil
|
||||
) {
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.parameters = parameters
|
||||
}
|
||||
}
|
||||
|
||||
/// Function parameters (JSON Schema)
|
||||
public struct FunctionParameters: Codable {
|
||||
public let type: String
|
||||
public let properties: [String: PropertySchema]?
|
||||
public let required: [String]?
|
||||
|
||||
public init(
|
||||
type: String = "object",
|
||||
properties: [String: PropertySchema]? = nil,
|
||||
required: [String]? = nil
|
||||
) {
|
||||
self.type = type
|
||||
self.properties = properties
|
||||
self.required = required
|
||||
}
|
||||
}
|
||||
|
||||
/// Property schema
|
||||
public struct PropertySchema: Codable {
|
||||
public let type: String
|
||||
public let description: String?
|
||||
public let `enum`: [String]?
|
||||
|
||||
public init(
|
||||
type: String,
|
||||
description: String? = nil,
|
||||
enum: [String]? = nil
|
||||
) {
|
||||
self.type = type
|
||||
self.description = description
|
||||
self.enum = `enum`
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool call in response
|
||||
public struct ToolCall: Codable, Sendable {
|
||||
public let id: String
|
||||
public let type: String
|
||||
public let function: FunctionCall
|
||||
|
||||
public init(
|
||||
id: String = "call_\(UUID().uuidString.replacingOccurrences(of: "-", with: "").prefix(9))",
|
||||
type: String = "function",
|
||||
function: FunctionCall
|
||||
) {
|
||||
self.id = id
|
||||
self.type = type
|
||||
self.function = function
|
||||
}
|
||||
}
|
||||
|
||||
/// Function call
|
||||
public struct FunctionCall: Codable, Sendable {
|
||||
public let name: String
|
||||
public let arguments: String
|
||||
|
||||
public init(name: String, arguments: String) {
|
||||
self.name = name
|
||||
self.arguments = arguments
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool message
|
||||
public struct ToolMessage: Codable {
|
||||
public let role: String
|
||||
public let content: String?
|
||||
public let tool_call_id: String?
|
||||
public let name: String?
|
||||
|
||||
public init(
|
||||
role: String = "tool",
|
||||
content: String?,
|
||||
tool_call_id: String?,
|
||||
name: String?
|
||||
) {
|
||||
self.role = role
|
||||
self.content = content
|
||||
self.tool_call_id = tool_call_id
|
||||
self.name = name
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// JSON Schema Response API Models
|
||||
// OpenAI Compatible Format
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Response format specification
|
||||
public enum ResponseFormat: Codable {
|
||||
case text
|
||||
case jsonObject
|
||||
case jsonSchema(JSONSchemaDefinition)
|
||||
|
||||
private enum CodingKeys: String, CodingKey {
|
||||
case type
|
||||
case jsonSchema = "json_schema"
|
||||
}
|
||||
|
||||
public init(from decoder: Decoder) throws {
|
||||
let container = try decoder.container(keyedBy: CodingKeys.self)
|
||||
let type = try container.decode(String.self, forKey: .type)
|
||||
|
||||
switch type {
|
||||
case "text":
|
||||
self = .text
|
||||
case "json_object":
|
||||
self = .jsonObject
|
||||
case "json_schema":
|
||||
let schema = try container.decode(JSONSchemaDefinition.self, forKey: .jsonSchema)
|
||||
self = .jsonSchema(schema)
|
||||
default:
|
||||
throw DecodingError.dataCorruptedError(
|
||||
forKey: .type,
|
||||
in: container,
|
||||
debugDescription: "Unknown response format type: \(type)"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
public func encode(to encoder: Encoder) throws {
|
||||
var container = encoder.container(keyedBy: CodingKeys.self)
|
||||
|
||||
switch self {
|
||||
case .text:
|
||||
try container.encode("text", forKey: .type)
|
||||
case .jsonObject:
|
||||
try container.encode("json_object", forKey: .type)
|
||||
case .jsonSchema(let schema):
|
||||
try container.encode("json_schema", forKey: .type)
|
||||
try container.encode(schema, forKey: .jsonSchema)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON Schema definition
|
||||
public struct JSONSchemaDefinition: Codable {
|
||||
public let name: String
|
||||
public let description: String?
|
||||
public let schema: JSONSchema
|
||||
public let strict: Bool?
|
||||
|
||||
public init(
|
||||
name: String,
|
||||
description: String? = nil,
|
||||
schema: JSONSchema,
|
||||
strict: Bool? = nil
|
||||
) {
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.schema = schema
|
||||
self.strict = strict
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON Schema
|
||||
public struct JSONSchema: Codable {
|
||||
public let type: String
|
||||
public let properties: [String: SchemaProperty]?
|
||||
public let required: [String]?
|
||||
public let additionalProperties: Bool?
|
||||
|
||||
public init(
|
||||
type: String,
|
||||
properties: [String: SchemaProperty]? = nil,
|
||||
required: [String]? = nil,
|
||||
additionalProperties: Bool? = false
|
||||
) {
|
||||
self.type = type
|
||||
self.properties = properties
|
||||
self.required = required
|
||||
self.additionalProperties = additionalProperties
|
||||
}
|
||||
}
|
||||
|
||||
/// Schema property
|
||||
public struct SchemaProperty: Codable {
|
||||
public let type: String
|
||||
public let description: String?
|
||||
public let `enum`: [String]?
|
||||
public let properties: [String: SchemaProperty]?
|
||||
public let required: [String]?
|
||||
|
||||
public init(
|
||||
type: String,
|
||||
description: String? = nil,
|
||||
enum: [String]? = nil,
|
||||
properties: [String: SchemaProperty]? = nil,
|
||||
required: [String]? = nil
|
||||
) {
|
||||
self.type = type
|
||||
self.description = description
|
||||
self.enum = `enum`
|
||||
self.properties = properties
|
||||
self.required = required
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Load Balancer for Cross-Device Communication
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Load balancing strategies
|
||||
public enum LoadBalancingStrategy: Sendable {
|
||||
case roundRobin
|
||||
case leastLoaded
|
||||
case random
|
||||
case geographic
|
||||
}
|
||||
|
||||
/// Load balancer
|
||||
public final class LoadBalancer: @unchecked Sendable {
|
||||
private var nodes: [DeviceNode] = []
|
||||
private var currentIndex: Int = 0
|
||||
private let strategy: LoadBalancingStrategy
|
||||
private let lock = NSLock()
|
||||
|
||||
public init(strategy: LoadBalancingStrategy = .roundRobin) {
|
||||
self.strategy = strategy
|
||||
}
|
||||
|
||||
/// Add a node
|
||||
public func addNode(_ node: DeviceNode) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
nodes.append(node)
|
||||
}
|
||||
|
||||
/// Remove a node
|
||||
public func removeNode(id: String) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
nodes.removeAll { $0.id == id }
|
||||
}
|
||||
|
||||
/// Get next node based on strategy
|
||||
public func getNextNode() -> DeviceNode? {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
let healthyNodes = nodes.filter { $0.status == .healthy }
|
||||
guard !healthyNodes.isEmpty else { return nil }
|
||||
|
||||
switch strategy {
|
||||
case .roundRobin:
|
||||
let node = healthyNodes[currentIndex % healthyNodes.count]
|
||||
currentIndex += 1
|
||||
return node
|
||||
|
||||
case .leastLoaded:
|
||||
return healthyNodes.min { $0.load < $1.load }
|
||||
|
||||
case .random:
|
||||
return healthyNodes.randomElement()
|
||||
|
||||
case .geographic:
|
||||
// Simplified: return node with lowest latency
|
||||
return healthyNodes.first
|
||||
}
|
||||
}
|
||||
|
||||
/// Update node status
|
||||
public func updateNodeStatus(id: String, status: DeviceStatus, load: Double? = nil) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
if let index = nodes.firstIndex(where: { $0.id == id }) {
|
||||
var node = nodes[index]
|
||||
node.status = status
|
||||
if let load = load {
|
||||
node.load = load
|
||||
}
|
||||
nodes[index] = node
|
||||
}
|
||||
}
|
||||
|
||||
/// Get all nodes
|
||||
public func getNodes() -> [DeviceNode] {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
return nodes
|
||||
}
|
||||
|
||||
/// Get healthy nodes count
|
||||
public var healthyNodesCount: Int {
|
||||
nodes.filter { $0.status == .healthy }.count
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,176 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model Downloader - Download models from HuggingFace
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Download progress
|
||||
public struct DownloadProgress: Sendable {
|
||||
public let current: Int64
|
||||
public let total: Int64
|
||||
public let percentage: Double
|
||||
public let speed: Double // bytes per second
|
||||
public let eta: TimeInterval // seconds
|
||||
|
||||
public init(current: Int64, total: Int64, speed: Double = 0) {
|
||||
self.current = current
|
||||
self.total = total
|
||||
self.percentage = total > 0 ? Double(current) / Double(total) : 0
|
||||
self.speed = speed
|
||||
self.eta = speed > 0 ? Double(total - current) / speed : 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Model downloader
|
||||
public final class ModelDownloader: @unchecked Sendable {
|
||||
public static let shared = ModelDownloader()
|
||||
|
||||
private let session: URLSession
|
||||
private var progressHandler: ((DownloadProgress) -> Void)?
|
||||
|
||||
private init() {
|
||||
let config = URLSessionConfiguration.default
|
||||
config.timeoutIntervalForRequest = 300
|
||||
config.timeoutIntervalForResource = 3600
|
||||
self.session = URLSession(configuration: config)
|
||||
}
|
||||
|
||||
/// Set progress handler
|
||||
public func onProgress(_ handler: @escaping (DownloadProgress) -> Void) {
|
||||
self.progressHandler = handler
|
||||
}
|
||||
|
||||
/// Download model from HuggingFace
|
||||
public func downloadModel(
|
||||
repoId: String,
|
||||
to destinationDir: String,
|
||||
revision: String = "main"
|
||||
) async throws {
|
||||
let destinationURL = URL(fileURLWithPath: destinationDir)
|
||||
|
||||
// Create destination directory
|
||||
try FileManager.default.createDirectory(
|
||||
at: destinationURL,
|
||||
withIntermediateDirectories: true
|
||||
)
|
||||
|
||||
// List files in repo
|
||||
let files = try await listRepoFiles(repoId: repoId, revision: revision)
|
||||
|
||||
// Download each file
|
||||
for file in files {
|
||||
let fileURL = destinationURL.appendingPathComponent(file)
|
||||
|
||||
// Skip if already exists
|
||||
if FileManager.default.fileExists(atPath: fileURL.path) {
|
||||
print("Skipping \(file) (already exists)")
|
||||
continue
|
||||
}
|
||||
|
||||
print("Downloading \(file)...")
|
||||
try await downloadFile(
|
||||
repoId: repoId,
|
||||
file: file,
|
||||
revision: revision,
|
||||
to: fileURL
|
||||
)
|
||||
}
|
||||
|
||||
print("✓ Model downloaded to \(destinationDir)")
|
||||
}
|
||||
|
||||
/// List files in repo
|
||||
private func listRepoFiles(repoId: String, revision: String) async throws -> [String] {
|
||||
let url = URL(string: "https://huggingface.co/api/models/\(repoId)/tree/\(revision)")!
|
||||
let (data, _) = try await session.data(from: url)
|
||||
|
||||
struct FileInfo: Codable {
|
||||
let path: String
|
||||
let type: String
|
||||
}
|
||||
|
||||
let files = try JSONDecoder().decode([FileInfo].self, from: data)
|
||||
return files.filter { $0.type == "file" }.map { $0.path }
|
||||
}
|
||||
|
||||
/// Download single file
|
||||
private func downloadFile(
|
||||
repoId: String,
|
||||
file: String,
|
||||
revision: String,
|
||||
to destination: URL
|
||||
) async throws {
|
||||
let downloadURL = URL(
|
||||
string: "https://huggingface.co/\(repoId)/resolve/\(revision)/\(file)"
|
||||
)!
|
||||
|
||||
let (tempURL, response) = try await session.download(from: downloadURL)
|
||||
|
||||
// Get total size
|
||||
let total = response.expectedContentLength
|
||||
|
||||
// Move to destination
|
||||
try FileManager.default.moveItem(at: tempURL, to: destination)
|
||||
}
|
||||
|
||||
/// Download with progress
|
||||
private func downloadWithProgress(
|
||||
from url: URL,
|
||||
to destination: URL
|
||||
) async throws {
|
||||
var request = URLRequest(url: url)
|
||||
request.timeoutInterval = 300
|
||||
|
||||
let (bytes, response) = try await session.bytes(for: request)
|
||||
let total = response.expectedContentLength
|
||||
|
||||
var current: Int64 = 0
|
||||
var startTime = Date()
|
||||
|
||||
// Create file
|
||||
let fileHandle = try FileHandle(forWritingTo: destination)
|
||||
defer { try? fileHandle.close() }
|
||||
|
||||
for try await byte in bytes {
|
||||
try fileHandle.write(contentsOf: [byte])
|
||||
current += 1
|
||||
|
||||
// Update progress every 1MB
|
||||
if current % (1024 * 1024) == 0 {
|
||||
let elapsed = Date().timeIntervalSince(startTime)
|
||||
let speed = elapsed > 0 ? Double(current) / elapsed : 0
|
||||
|
||||
let progress = DownloadProgress(
|
||||
current: current,
|
||||
total: total,
|
||||
speed: speed
|
||||
)
|
||||
|
||||
progressHandler?(progress)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Model repository
|
||||
public struct ModelRepository {
|
||||
public let id: String
|
||||
public let name: String
|
||||
public let description: String
|
||||
public let downloads: Int
|
||||
public let likes: Int
|
||||
|
||||
public init(
|
||||
id: String,
|
||||
name: String,
|
||||
description: String,
|
||||
downloads: Int,
|
||||
likes: Int
|
||||
) {
|
||||
self.id = id
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.downloads = downloads
|
||||
self.likes = likes
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model Finder - Search and Select Models
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Model search result
|
||||
public struct ModelSearchResult: Codable, Sendable {
|
||||
public let id: String
|
||||
public let modelId: String
|
||||
public let author: String
|
||||
public let sha: String?
|
||||
public let lastModified: String?
|
||||
public let isPrivate: Bool
|
||||
public let disabled: Bool
|
||||
public let gated: String?
|
||||
public let pipelineTag: String?
|
||||
public let tags: [String]
|
||||
public let downloads: Int
|
||||
public let likes: Int
|
||||
public let libraryName: String?
|
||||
public let createdAt: String?
|
||||
|
||||
public var displayName: String {
|
||||
modelId.split(separator: "/").last.map(String.init) ?? modelId
|
||||
}
|
||||
|
||||
public var isGGUF: Bool {
|
||||
tags.contains { $0.lowercased().contains("gguf") }
|
||||
}
|
||||
|
||||
public var isSafetensors: Bool {
|
||||
tags.contains { $0.lowercased().contains("safetensors") }
|
||||
}
|
||||
}
|
||||
|
||||
/// Model search filters
|
||||
public struct ModelFilters {
|
||||
public var search: String?
|
||||
public var author: String?
|
||||
public var library: String?
|
||||
public var task: String?
|
||||
public var tags: [String]?
|
||||
public var sort: String?
|
||||
public var direction: Int? // -1 for descending, 1 for ascending
|
||||
public var limit: Int?
|
||||
|
||||
public init(
|
||||
search: String? = nil,
|
||||
author: String? = nil,
|
||||
library: String? = nil,
|
||||
task: String? = nil,
|
||||
tags: [String]? = nil,
|
||||
sort: String? = nil,
|
||||
direction: Int? = nil,
|
||||
limit: Int? = nil
|
||||
) {
|
||||
self.search = search
|
||||
self.author = author
|
||||
self.library = library
|
||||
self.task = task
|
||||
self.tags = tags
|
||||
self.sort = sort
|
||||
self.direction = direction
|
||||
self.limit = limit
|
||||
}
|
||||
|
||||
/// Convert to query parameters
|
||||
public func toQueryItems() -> [URLQueryItem] {
|
||||
var items: [URLQueryItem] = []
|
||||
|
||||
if let search = search {
|
||||
items.append(URLQueryItem(name: "search", value: search))
|
||||
}
|
||||
if let author = author {
|
||||
items.append(URLQueryItem(name: "author", value: author))
|
||||
}
|
||||
if let library = library {
|
||||
items.append(URLQueryItem(name: "filter", value: library))
|
||||
}
|
||||
if let task = task {
|
||||
items.append(URLQueryItem(name: "pipeline_tag", value: task))
|
||||
}
|
||||
if let tags = tags {
|
||||
for tag in tags {
|
||||
items.append(URLQueryItem(name: "filter", value: tag))
|
||||
}
|
||||
}
|
||||
if let sort = sort {
|
||||
items.append(URLQueryItem(name: "sort", value: sort))
|
||||
}
|
||||
if let direction = direction {
|
||||
items.append(URLQueryItem(name: "direction", value: "\(direction)"))
|
||||
}
|
||||
if let limit = limit {
|
||||
items.append(URLQueryItem(name: "limit", value: "\(limit)"))
|
||||
}
|
||||
|
||||
return items
|
||||
}
|
||||
}
|
||||
|
||||
/// Model finder
|
||||
public final class ModelFinder {
|
||||
public nonisolated(unsafe) static let shared = ModelFinder()
|
||||
|
||||
private let session: URLSession
|
||||
|
||||
private init() {
|
||||
let config = URLSessionConfiguration.default
|
||||
config.timeoutIntervalForRequest = 30
|
||||
self.session = URLSession(configuration: config)
|
||||
}
|
||||
|
||||
/// Search models on HuggingFace
|
||||
public func searchModels(filters: ModelFilters) async throws -> [ModelSearchResult] {
|
||||
var components = URLComponents(string: "https://huggingface.co/api/models")!
|
||||
components.queryItems = filters.toQueryItems()
|
||||
|
||||
guard let url = components.url else {
|
||||
throw ModelFinderError.invalidURL
|
||||
}
|
||||
|
||||
let (data, _) = try await session.data(from: url)
|
||||
let models = try JSONDecoder().decode([ModelSearchResult].self, from: data)
|
||||
|
||||
return models
|
||||
}
|
||||
|
||||
/// Search by name
|
||||
public func searchByName(_ name: String, limit: Int = 20) async throws -> [ModelSearchResult] {
|
||||
let filters = ModelFilters(
|
||||
search: name,
|
||||
sort: "downloads",
|
||||
direction: -1,
|
||||
limit: limit
|
||||
)
|
||||
return try await searchModels(filters: filters)
|
||||
}
|
||||
|
||||
/// Search GGUF models
|
||||
public func searchGGUF(
|
||||
name: String? = nil,
|
||||
limit: Int = 20,
|
||||
sortBy: String = "downloads"
|
||||
) async throws -> [ModelSearchResult] {
|
||||
let filters = ModelFilters(
|
||||
search: name,
|
||||
tags: ["gguf"],
|
||||
sort: sortBy,
|
||||
direction: -1,
|
||||
limit: limit
|
||||
)
|
||||
return try await searchModels(filters: filters)
|
||||
}
|
||||
|
||||
/// Search Safetensors models
|
||||
public func searchSafetensors(
|
||||
name: String? = nil,
|
||||
limit: Int = 20,
|
||||
sortBy: String = "downloads"
|
||||
) async throws -> [ModelSearchResult] {
|
||||
let filters = ModelFilters(
|
||||
search: name,
|
||||
tags: ["safetensors"],
|
||||
sort: sortBy,
|
||||
direction: -1,
|
||||
limit: limit
|
||||
)
|
||||
return try await searchModels(filters: filters)
|
||||
}
|
||||
|
||||
/// Get model details
|
||||
public func getModelDetails(repoId: String) async throws -> ModelSearchResult {
|
||||
let url = URL(string: "https://huggingface.co/api/models/\(repoId)")!
|
||||
let (data, _) = try await session.data(from: url)
|
||||
return try JSONDecoder().decode(ModelSearchResult.self, from: data)
|
||||
}
|
||||
|
||||
/// List model files
|
||||
public func listModelFiles(repoId: String, revision: String = "main") async throws -> [String] {
|
||||
let url = URL(string: "https://huggingface.co/api/models/\(repoId)/tree/\(revision)")!
|
||||
let (data, _) = try await session.data(from: url)
|
||||
|
||||
struct FileInfo: Codable {
|
||||
let path: String
|
||||
let type: String
|
||||
}
|
||||
|
||||
let files = try JSONDecoder().decode([FileInfo].self, from: data)
|
||||
return files.filter { $0.type == "file" }.map { $0.path }
|
||||
}
|
||||
}
|
||||
|
||||
/// Model finder errors
|
||||
public enum ModelFinderError: Error, LocalizedError {
|
||||
case invalidURL
|
||||
case requestFailed(String)
|
||||
case decodeFailed(String)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .invalidURL:
|
||||
return "Invalid URL"
|
||||
case .requestFailed(let detail):
|
||||
return "Request failed: \(detail)"
|
||||
case .decodeFailed(let detail):
|
||||
return "Failed to decode response: \(detail)"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model Manager - Runtime Model Switching
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Model information
|
||||
public struct ModelInfo: Codable, Sendable {
|
||||
public let id: String
|
||||
public let path: String
|
||||
public let name: String
|
||||
public var loaded: Bool
|
||||
public let parameters: [String: String]?
|
||||
|
||||
public init(
|
||||
id: String,
|
||||
path: String,
|
||||
name: String,
|
||||
loaded: Bool = false,
|
||||
parameters: [String: String]? = nil
|
||||
) {
|
||||
self.id = id
|
||||
self.path = path
|
||||
self.name = name
|
||||
self.loaded = loaded
|
||||
self.parameters = parameters
|
||||
}
|
||||
}
|
||||
|
||||
/// Model errors
|
||||
public enum ModelError: Error, LocalizedError {
|
||||
case modelNotFound(String)
|
||||
case noModelLoaded
|
||||
case loadFailed(String)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .modelNotFound(let id):
|
||||
return "Model not found: \(id)"
|
||||
case .noModelLoaded:
|
||||
return "No model is currently loaded"
|
||||
case .loadFailed(let detail):
|
||||
return "Failed to load model: \(detail)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Model manager for runtime model switching
|
||||
public actor ModelManager {
|
||||
private var models: [String: ModelInfo] = [:]
|
||||
private var currentModelId: String?
|
||||
|
||||
// Active model instances
|
||||
private var engine: MarkBaseEngine?
|
||||
private var model: E4BModel?
|
||||
private var tokenizer: Tokenizer?
|
||||
private var generator: StreamingGenerator?
|
||||
private var sampler: Sampler?
|
||||
|
||||
public init() {}
|
||||
|
||||
/// Register a model
|
||||
public func register(id: String, path: String, name: String) {
|
||||
models[id] = ModelInfo(id: id, path: path, name: name, loaded: false)
|
||||
}
|
||||
|
||||
/// Get list of registered models
|
||||
public func listModels() -> [ModelInfo] {
|
||||
return Array(models.values)
|
||||
}
|
||||
|
||||
/// Get current model info
|
||||
public func getCurrentModel() -> ModelInfo? {
|
||||
guard let id = currentModelId else { return nil }
|
||||
return models[id]
|
||||
}
|
||||
|
||||
/// Load a model
|
||||
public func loadModel(id: String) async throws {
|
||||
guard let modelInfo = models[id] else {
|
||||
throw ModelError.modelNotFound(id)
|
||||
}
|
||||
|
||||
// Validate path
|
||||
try Validator.validateModelPath(modelInfo.path)
|
||||
|
||||
// Create engine
|
||||
let newEngine = try MarkBaseEngine(autoCompile: true)
|
||||
|
||||
// Load model
|
||||
let newModel = try E4BModel(
|
||||
modelDir: modelInfo.path,
|
||||
engine: newEngine,
|
||||
maxContextLength: 512
|
||||
)
|
||||
|
||||
// Load tokenizer
|
||||
let newTokenizer = try TokenizerFactory.load(modelDir: modelInfo.path)
|
||||
|
||||
// Create generator
|
||||
let newGenerator = StreamingGenerator(
|
||||
model: newModel,
|
||||
tokenizer: newTokenizer,
|
||||
engine: newEngine
|
||||
)
|
||||
|
||||
// Update active instances
|
||||
engine = newEngine
|
||||
model = newModel
|
||||
tokenizer = newTokenizer
|
||||
generator = newGenerator
|
||||
sampler = Sampler()
|
||||
currentModelId = id
|
||||
models[id]?.loaded = true
|
||||
}
|
||||
|
||||
/// Unload current model
|
||||
public func unloadModel() {
|
||||
guard let id = currentModelId else { return }
|
||||
|
||||
engine = nil
|
||||
model = nil
|
||||
tokenizer = nil
|
||||
generator = nil
|
||||
sampler = nil
|
||||
currentModelId = nil
|
||||
models[id]?.loaded = false
|
||||
}
|
||||
|
||||
/// Switch to a different model
|
||||
public func switchModel(to id: String) async throws {
|
||||
// Unload current model if loaded
|
||||
if currentModelId != nil {
|
||||
unloadModel()
|
||||
}
|
||||
|
||||
// Load new model
|
||||
try await loadModel(id: id)
|
||||
}
|
||||
|
||||
/// Get active engine
|
||||
public func getEngine() throws -> MarkBaseEngine {
|
||||
guard let engine = engine else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
return engine
|
||||
}
|
||||
|
||||
/// Get active model
|
||||
public func getModel() throws -> E4BModel {
|
||||
guard let model = model else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
/// Get active tokenizer
|
||||
public func getTokenizer() throws -> Tokenizer {
|
||||
guard let tokenizer = tokenizer else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
return tokenizer
|
||||
}
|
||||
|
||||
/// Get active generator
|
||||
public func getGenerator() throws -> StreamingGenerator {
|
||||
guard let generator = generator else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
return generator
|
||||
}
|
||||
|
||||
/// Get active sampler
|
||||
public func getSampler() throws -> Sampler {
|
||||
guard let sampler = sampler else {
|
||||
throw ModelError.noModelLoaded
|
||||
}
|
||||
return sampler
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model API Models
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Chat completion request (for testing)
|
||||
public struct ChatCompletionRequest: Codable {
|
||||
public let model: String
|
||||
public let messages: [ChatMessage]
|
||||
public let max_tokens: Int?
|
||||
public let temperature: Float?
|
||||
public let stream: Bool?
|
||||
|
||||
public func toGenerationConfig() -> GenerationConfig {
|
||||
GenerationConfig(
|
||||
maxTokens: max_tokens ?? 100,
|
||||
temperature: temperature ?? 1.0
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Multimodal chat completion request (for testing)
|
||||
public struct MultimodalChatCompletionRequest: Codable {
|
||||
public let model: String
|
||||
public let messages: [MultimodalMessage]
|
||||
public let max_tokens: Int?
|
||||
public let stream: Bool?
|
||||
|
||||
public func toGenerationConfig() -> GenerationConfig {
|
||||
GenerationConfig(maxTokens: max_tokens ?? 100)
|
||||
}
|
||||
}
|
||||
|
||||
/// Model capabilities
|
||||
public struct ModelCapabilities: Codable, Sendable {
|
||||
public let text: Bool
|
||||
public let vision: Bool
|
||||
public let audio: Bool
|
||||
public let embeddings: Bool
|
||||
public let streaming: Bool
|
||||
|
||||
public init(
|
||||
text: Bool = true,
|
||||
vision: Bool = true,
|
||||
audio: Bool = true,
|
||||
embeddings: Bool = true,
|
||||
streaming: Bool = true
|
||||
) {
|
||||
self.text = text
|
||||
self.vision = vision
|
||||
self.audio = audio
|
||||
self.embeddings = embeddings
|
||||
self.streaming = streaming
|
||||
}
|
||||
}
|
||||
|
||||
/// Model parameters
|
||||
public struct ModelParameters: Codable, Sendable {
|
||||
public let context_length: Int
|
||||
public let num_hidden_layers: Int
|
||||
public let hidden_size: Int
|
||||
public let vocab_size: Int
|
||||
public let num_attention_heads: Int
|
||||
public let num_kv_heads: Int
|
||||
|
||||
public init(
|
||||
context_length: Int,
|
||||
num_hidden_layers: Int,
|
||||
hidden_size: Int,
|
||||
vocab_size: Int,
|
||||
num_attention_heads: Int,
|
||||
num_kv_heads: Int
|
||||
) {
|
||||
self.context_length = context_length
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.hidden_size = hidden_size
|
||||
self.vocab_size = vocab_size
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
}
|
||||
}
|
||||
|
||||
/// Model details response
|
||||
public struct ModelDetails: Codable, Sendable {
|
||||
public let id: String
|
||||
public let object: String
|
||||
public let created: Int
|
||||
public let owned_by: String
|
||||
public let capabilities: ModelCapabilities
|
||||
public let parameters: ModelParameters
|
||||
|
||||
public init(
|
||||
id: String,
|
||||
object: String = "model",
|
||||
created: Int = Int(Date().timeIntervalSince1970),
|
||||
owned_by: String = "markbase",
|
||||
capabilities: ModelCapabilities,
|
||||
parameters: ModelParameters
|
||||
) {
|
||||
self.id = id
|
||||
self.object = object
|
||||
self.created = created
|
||||
self.owned_by = owned_by
|
||||
self.capabilities = capabilities
|
||||
self.parameters = parameters
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Multimodal API Models
|
||||
// OpenAI Compatible Format
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Multimodal message content part
|
||||
public enum ContentPart: Codable {
|
||||
case text(String)
|
||||
case imageUrl(ImageUrl)
|
||||
case audioUrl(AudioUrl)
|
||||
case videoUrl(VideoUrl)
|
||||
|
||||
public struct ImageUrl: Codable {
|
||||
public let url: String
|
||||
|
||||
public init(url: String) {
|
||||
self.url = url
|
||||
}
|
||||
}
|
||||
|
||||
public struct AudioUrl: Codable {
|
||||
public let url: String
|
||||
|
||||
public init(url: String) {
|
||||
self.url = url
|
||||
}
|
||||
}
|
||||
|
||||
public struct VideoUrl: Codable {
|
||||
public let url: String
|
||||
|
||||
public init(url: String) {
|
||||
self.url = url
|
||||
}
|
||||
}
|
||||
|
||||
private enum CodingKeys: String, CodingKey {
|
||||
case type
|
||||
case text
|
||||
case imageUrl = "image_url"
|
||||
case audioUrl = "audio_url"
|
||||
case videoUrl = "video_url"
|
||||
}
|
||||
|
||||
public init(from decoder: Decoder) throws {
|
||||
let container = try decoder.container(keyedBy: CodingKeys.self)
|
||||
let type = try container.decode(String.self, forKey: .type)
|
||||
|
||||
switch type {
|
||||
case "text":
|
||||
let text = try container.decode(String.self, forKey: .text)
|
||||
self = .text(text)
|
||||
case "image_url":
|
||||
let imageUrl = try container.decode(ImageUrl.self, forKey: .imageUrl)
|
||||
self = .imageUrl(imageUrl)
|
||||
case "audio_url":
|
||||
let audioUrl = try container.decode(AudioUrl.self, forKey: .audioUrl)
|
||||
self = .audioUrl(audioUrl)
|
||||
case "video_url":
|
||||
let videoUrl = try container.decode(VideoUrl.self, forKey: .videoUrl)
|
||||
self = .videoUrl(videoUrl)
|
||||
default:
|
||||
throw DecodingError.dataCorruptedError(
|
||||
forKey: .type,
|
||||
in: container,
|
||||
debugDescription: "Unknown content type: \(type)"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
public func encode(to encoder: Encoder) throws {
|
||||
var container = encoder.container(keyedBy: CodingKeys.self)
|
||||
|
||||
switch self {
|
||||
case .text(let text):
|
||||
try container.encode("text", forKey: .type)
|
||||
try container.encode(text, forKey: .text)
|
||||
case .imageUrl(let imageUrl):
|
||||
try container.encode("image_url", forKey: .type)
|
||||
try container.encode(imageUrl, forKey: .imageUrl)
|
||||
case .audioUrl(let audioUrl):
|
||||
try container.encode("audio_url", forKey: .type)
|
||||
try container.encode(audioUrl, forKey: .audioUrl)
|
||||
case .videoUrl(let videoUrl):
|
||||
try container.encode("video_url", forKey: .type)
|
||||
try container.encode(videoUrl, forKey: .videoUrl)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Multimodal message
|
||||
public struct MultimodalMessage: Codable {
|
||||
public let role: String
|
||||
public let content: [ContentPart]
|
||||
|
||||
public init(role: String, content: [ContentPart]) {
|
||||
self.role = role
|
||||
self.content = content
|
||||
}
|
||||
|
||||
/// Extract text content
|
||||
public var textContent: String {
|
||||
content.compactMap { part -> String? in
|
||||
if case .text(let text) = part { return text }
|
||||
return nil
|
||||
}.joined(separator: "\n")
|
||||
}
|
||||
|
||||
/// Extract image URLs
|
||||
public var imageUrls: [ContentPart.ImageUrl] {
|
||||
content.compactMap { part -> ContentPart.ImageUrl? in
|
||||
if case .imageUrl(let url) = part { return url }
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract audio URLs
|
||||
public var audioUrls: [ContentPart.AudioUrl] {
|
||||
content.compactMap { part -> ContentPart.AudioUrl? in
|
||||
if case .audioUrl(let url) = part { return url }
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract video URLs
|
||||
public var videoUrls: [ContentPart.VideoUrl] {
|
||||
content.compactMap { part -> ContentPart.VideoUrl? in
|
||||
if case .videoUrl(let url) = part { return url }
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Multimodal chat completion request
|
||||
public struct MultimodalChatRequest: Codable {
|
||||
public let messages: [MultimodalMessage]
|
||||
public let max_tokens: Int?
|
||||
public let temperature: Float?
|
||||
public let top_p: Float?
|
||||
public let top_k: Int?
|
||||
public let stream: Bool?
|
||||
public let tools: [Tool]?
|
||||
public let response_format: ResponseFormat?
|
||||
|
||||
public init(
|
||||
messages: [MultimodalMessage],
|
||||
max_tokens: Int? = nil,
|
||||
temperature: Float? = nil,
|
||||
top_p: Float? = nil,
|
||||
top_k: Int? = nil,
|
||||
stream: Bool? = nil,
|
||||
tools: [Tool]? = nil,
|
||||
response_format: ResponseFormat? = nil
|
||||
) {
|
||||
self.messages = messages
|
||||
self.max_tokens = max_tokens
|
||||
self.temperature = temperature
|
||||
self.top_p = top_p
|
||||
self.top_k = top_k
|
||||
self.stream = stream
|
||||
self.tools = tools
|
||||
self.response_format = response_format
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Image/Audio Processing Helpers
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public enum MediaProcessor {
|
||||
/// Parse data URI to base64 and mime type
|
||||
public static func parseDataURI(_ uri: String) throws -> (mimeType: String, data: Data) {
|
||||
guard uri.hasPrefix("data:") else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "Expected data URI format"
|
||||
)
|
||||
}
|
||||
|
||||
let parts = uri.split(separator: ",", maxSplits: 1)
|
||||
guard parts.count == 2 else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "Invalid data URI format"
|
||||
)
|
||||
}
|
||||
|
||||
let header = String(parts[0])
|
||||
let base64 = String(parts[1])
|
||||
|
||||
// Extract mime type
|
||||
let mimeType = header.dropFirst(5).split(separator: ";").first.map(String.init) ?? "application/octet-stream"
|
||||
|
||||
// Decode base64
|
||||
guard let data = Data(base64Encoded: base64) else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "Invalid base64 data"
|
||||
)
|
||||
}
|
||||
|
||||
return (mimeType, data)
|
||||
}
|
||||
|
||||
/// Load image from URL or data URI
|
||||
public static func loadImage(from url: String) throws -> Data {
|
||||
if url.hasPrefix("data:") {
|
||||
let (_, data) = try parseDataURI(url)
|
||||
return data
|
||||
} else if url.hasPrefix("http://") || url.hasPrefix("https://") {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "HTTP URLs not yet supported, use base64 data URI"
|
||||
)
|
||||
} else {
|
||||
// Local file path
|
||||
let filePath = url
|
||||
guard FileManager.default.fileExists(atPath: filePath) else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "File not found: \(filePath)"
|
||||
)
|
||||
}
|
||||
return try Data(contentsOf: URL(fileURLWithPath: filePath))
|
||||
}
|
||||
}
|
||||
|
||||
/// Load audio from URL or data URI
|
||||
public static func loadAudio(from url: String) throws -> Data {
|
||||
if url.hasPrefix("data:") {
|
||||
let (_, data) = try parseDataURI(url)
|
||||
return data
|
||||
} else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "Only base64 data URI supported for audio"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Load video from file path or data URI
|
||||
public static func loadVideo(from url: String) throws -> URL {
|
||||
if url.hasPrefix("data:") {
|
||||
let (mimeType, data) = try parseDataURI(url)
|
||||
let ext: String
|
||||
if mimeType.contains("mp4") { ext = "mp4" }
|
||||
else if mimeType.contains("quicktime") { ext = "mov" }
|
||||
else { ext = "mp4" }
|
||||
let tempURL = FileManager.default.temporaryDirectory
|
||||
.appendingPathComponent(UUID().uuidString)
|
||||
.appendingPathExtension(ext)
|
||||
try data.write(to: tempURL)
|
||||
return tempURL
|
||||
} else {
|
||||
let fileURL = URL(fileURLWithPath: url)
|
||||
guard FileManager.default.fileExists(atPath: fileURL.path) else {
|
||||
throw MarkBaseError.invalidParameter(
|
||||
parameter: "url",
|
||||
message: "Video file not found: \(url)"
|
||||
)
|
||||
}
|
||||
return fileURL
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Performance Benchmarking Tool
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public struct PerformanceBenchmark {
|
||||
private let modelDir: String
|
||||
private let modelName: String
|
||||
|
||||
public init(modelDir: String, modelName: String = "markbase") {
|
||||
self.modelDir = modelDir
|
||||
self.modelName = modelName
|
||||
}
|
||||
|
||||
/// Run all benchmarks
|
||||
public mutating func run() async throws {
|
||||
print("""
|
||||
|
||||
╔══════════════════════════════════════╗
|
||||
║ Performance Benchmark ║
|
||||
║ Model: \(modelName) ║
|
||||
╚══════════════════════════════════════╝
|
||||
|
||||
""")
|
||||
|
||||
try await benchmarkModelLoading()
|
||||
try await benchmarkTokenGeneration()
|
||||
try await benchmarkTokenizer()
|
||||
try await benchmarkBufferPool()
|
||||
|
||||
print("\n✅ All benchmarks completed!")
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Model Loading Benchmark
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private mutating func benchmarkModelLoading() async throws {
|
||||
print("\n📊 Model Loading Benchmark")
|
||||
print(String(repeating: "─", count: 40))
|
||||
|
||||
let start = Date()
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let loadEngineTime = Date().timeIntervalSince(start)
|
||||
print(" Engine initialization: \(String(format: "%.3f", loadEngineTime))s")
|
||||
|
||||
let start2 = Date()
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
let loadModelTime = Date().timeIntervalSince(start2)
|
||||
print(" Model loading: \(String(format: "%.3f", loadModelTime))s")
|
||||
|
||||
print(" Total: \(String(format: "%.3f", loadEngineTime + loadModelTime))s")
|
||||
print(" Layers: \(model.numHiddenLayers)")
|
||||
print(" Vocab: \(model.vocabSize)")
|
||||
print(" Hidden: \(model.hiddenSize)")
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Token Generation Benchmark
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private mutating func benchmarkTokenGeneration() async throws {
|
||||
print("\n📊 Token Generation Benchmark")
|
||||
print(String(repeating: "─", count: 40))
|
||||
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
|
||||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||||
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||||
let sampler = Sampler()
|
||||
|
||||
// Test different temperatures
|
||||
print("\nTesting with different temperatures:")
|
||||
|
||||
for temp in [Float(0.0), Float(0.7), Float(1.0)] {
|
||||
let prompt = "Hello, how are you?"
|
||||
let config = GenerationConfig(maxTokens: 20, temperature: temp)
|
||||
|
||||
print("\n Temperature \(temp):")
|
||||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||||
print(" Generated: \"\(response)\"")
|
||||
}
|
||||
|
||||
// Now run benchmark with temperature=0.7
|
||||
let prompt = "Hello, how are you?"
|
||||
let config = GenerationConfig(maxTokens: 20, temperature: Float(0.7))
|
||||
|
||||
// Benchmark
|
||||
let numRuns = 3
|
||||
var totalTokens = 0
|
||||
var totalTime: TimeInterval = 0
|
||||
|
||||
for i in 1...numRuns {
|
||||
let start = Date()
|
||||
let response = try generator.generateComplete(prompt: prompt, config: config)
|
||||
let elapsed = Date().timeIntervalSince(start)
|
||||
|
||||
let tokens = tokenizer.encode(text: response).count
|
||||
totalTokens += tokens
|
||||
totalTime += elapsed
|
||||
|
||||
print(" Run \(i): \(tokens) tokens in \(String(format: "%.3f", elapsed))s (\(String(format: "%.1f", Double(tokens) / elapsed)) tok/s)")
|
||||
if i == 1 {
|
||||
print(" Generated text: \"\(response)\"")
|
||||
}
|
||||
}
|
||||
|
||||
let avgTokens = totalTokens / numRuns
|
||||
let avgTime = totalTime / Double(numRuns)
|
||||
let avgSpeed = Double(avgTokens) / avgTime
|
||||
|
||||
print("\n Average: \(avgTokens) tokens in \(String(format: "%.3f", avgTime))s")
|
||||
print(" Speed: \(String(format: "%.1f", avgSpeed)) tok/s")
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Tokenizer Benchmark
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private func benchmarkTokenizer() throws {
|
||||
print("\n📊 Tokenizer Benchmark")
|
||||
print(String(repeating: "─", count: 40))
|
||||
|
||||
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||||
|
||||
let testTexts = [
|
||||
"Hello, world!",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Swift is a powerful and intuitive programming language for macOS, iOS, watchOS, and tvOS.",
|
||||
"Artificial intelligence is transforming the world in unprecedented ways."
|
||||
]
|
||||
|
||||
// Encode benchmark
|
||||
let numIterations = 100
|
||||
var totalEncodeTime: TimeInterval = 0
|
||||
|
||||
for text in testTexts {
|
||||
let start = Date()
|
||||
for _ in 0..<numIterations {
|
||||
_ = tokenizer.encode(text: text)
|
||||
}
|
||||
let elapsed = Date().timeIntervalSince(start)
|
||||
totalEncodeTime += elapsed
|
||||
|
||||
let tokens = tokenizer.encode(text: text).count
|
||||
print(" Encode: \"\(text.prefix(30))...\" -> \(tokens) tokens in \(String(format: "%.3f", elapsed / Double(numIterations) * 1000))ms")
|
||||
}
|
||||
|
||||
// Decode benchmark
|
||||
var totalDecodeTime: TimeInterval = 0
|
||||
|
||||
for text in testTexts {
|
||||
let tokens = tokenizer.encode(text: text)
|
||||
let start = Date()
|
||||
for _ in 0..<numIterations {
|
||||
_ = tokenizer.decode(tokens: tokens)
|
||||
}
|
||||
let elapsed = Date().timeIntervalSince(start)
|
||||
totalDecodeTime += elapsed
|
||||
|
||||
print(" Decode: \(tokens.count) tokens -> \"\(text.prefix(30))...\" in \(String(format: "%.3f", elapsed / Double(numIterations) * 1000))ms")
|
||||
}
|
||||
|
||||
print("\n Total encode time: \(String(format: "%.3f", totalEncodeTime))s")
|
||||
print(" Total decode time: \(String(format: "%.3f", totalDecodeTime))s")
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Buffer Pool Benchmark
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private func benchmarkBufferPool() throws {
|
||||
print("\n📊 Buffer Pool Benchmark")
|
||||
print(String(repeating: "─", count: 40))
|
||||
|
||||
let engine = try MarkBaseEngine()
|
||||
let pool = engine.bufferPool
|
||||
|
||||
let bufferSize = 1024
|
||||
let numIterations = 1000
|
||||
|
||||
// Without pool
|
||||
let start1 = Date()
|
||||
for _ in 0..<numIterations {
|
||||
let _ = try engine.makeBuffer(length: bufferSize)
|
||||
}
|
||||
let withoutPoolTime = Date().timeIntervalSince(start1)
|
||||
print(" Without pool: \(String(format: "%.3f", withoutPoolTime))s (\(String(format: "%.0f", Double(numIterations) / withoutPoolTime)) allocs/s)")
|
||||
|
||||
// With pool
|
||||
let start2 = Date()
|
||||
for _ in 0..<numIterations {
|
||||
let buf = engine.acquireBuffer(length: bufferSize)
|
||||
engine.releaseBuffer(buf)
|
||||
}
|
||||
let withPoolTime = Date().timeIntervalSince(start2)
|
||||
print(" With pool: \(String(format: "%.3f", withPoolTime))s (\(String(format: "%.0f", Double(numIterations) / withPoolTime)) allocs/s)")
|
||||
|
||||
let speedup = withoutPoolTime / withPoolTime
|
||||
print("\n Speedup: \(String(format: "%.1f", speedup))x")
|
||||
print(" Pool stats:")
|
||||
print(" \(pool.stats)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Performance Metrics API Models
|
||||
// Prometheus Compatible Format
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Performance metrics collector
|
||||
public final class MetricsCollector {
|
||||
public nonisolated(unsafe) static let shared = MetricsCollector()
|
||||
|
||||
private var counters: [String: CounterMetric] = [:]
|
||||
private var histograms: [String: MutableHistogram] = [:]
|
||||
private let lock = NSLock()
|
||||
|
||||
private init() {}
|
||||
|
||||
/// Record a request
|
||||
public func recordRequest(duration: TimeInterval, tokens: Int, model: String) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
// Request count
|
||||
incrementCounter(name: "markbase_requests_total", labels: ["model": model])
|
||||
|
||||
// Request duration
|
||||
observeHistogram(name: "markbase_request_duration_seconds", value: duration, labels: ["model": model])
|
||||
|
||||
// Tokens processed
|
||||
incrementCounter(name: "markbase_tokens_total", labels: ["model": model, "type": "output"], value: Double(tokens))
|
||||
}
|
||||
|
||||
/// Record a token generation
|
||||
public func recordToken(tokens: Int, duration: TimeInterval, model: String) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
let tokensPerSecond = Double(tokens) / duration
|
||||
observeHistogram(name: "markbase_tokens_per_second", value: tokensPerSecond, labels: ["model": model])
|
||||
}
|
||||
|
||||
/// Record an error
|
||||
public func recordError(type: String, model: String) {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
incrementCounter(name: "markbase_errors_total", labels: ["model": model, "type": type])
|
||||
}
|
||||
|
||||
/// Get Prometheus format metrics
|
||||
public func getPrometheusMetrics() -> String {
|
||||
lock.lock()
|
||||
defer { lock.unlock() }
|
||||
|
||||
var output = ""
|
||||
|
||||
for (name, counter) in counters {
|
||||
output += counter.toPrometheus(name: name)
|
||||
}
|
||||
|
||||
for (name, histogram) in histograms {
|
||||
output += histogram.toPrometheus(name: name)
|
||||
}
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Private helpers
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
private func incrementCounter(name: String, labels: [String: String], value: Double = 1.0) {
|
||||
if counters[name] == nil {
|
||||
counters[name] = CounterMetric()
|
||||
}
|
||||
counters[name]?.increment(by: value, labels: labels)
|
||||
}
|
||||
|
||||
private func observeHistogram(name: String, value: Double, labels: [String: String]) {
|
||||
if histograms[name] == nil {
|
||||
histograms[name] = MutableHistogram()
|
||||
}
|
||||
histograms[name]?.observe(value, labels: labels)
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Metric types
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
protocol Metric {
|
||||
func toPrometheus(name: String) -> String
|
||||
}
|
||||
|
||||
struct CounterMetric: Metric {
|
||||
var values: [[String: String]: Double] = [:]
|
||||
|
||||
mutating func increment(by value: Double, labels: [String: String]) {
|
||||
values[labels, default: 0] += value
|
||||
}
|
||||
|
||||
func toPrometheus(name: String) -> String {
|
||||
var output = "# TYPE \(name) counter\n"
|
||||
|
||||
for (labels, value) in values {
|
||||
let labelsStr = labels.map { "\($0.key)=\"\($0.value)\"" }.joined(separator: ",")
|
||||
output += "\(name){\(labelsStr)} \(value)\n"
|
||||
}
|
||||
|
||||
return output + "\n"
|
||||
}
|
||||
}
|
||||
|
||||
struct HistogramMetric: Metric {
|
||||
var counts: [[String: String]: Int] = [:]
|
||||
var sums: [[String: String]: Double] = [:]
|
||||
|
||||
mutating func observe(_ value: Double, labels: [String: String]) {
|
||||
counts[labels, default: 0] += 1
|
||||
sums[labels, default: 0] += value
|
||||
}
|
||||
|
||||
func toPrometheus(name: String) -> String {
|
||||
var output = "# TYPE \(name) histogram\n"
|
||||
|
||||
for (labels, count) in counts {
|
||||
let sum = sums[labels] ?? 0
|
||||
let labelsStr = labels.map { "\($0.key)=\"\($0.value)\"" }.joined(separator: ",")
|
||||
output += "\(name)_count{\(labelsStr)} \(count)\n"
|
||||
output += "\(name)_sum{\(labelsStr)} \(sum)\n"
|
||||
}
|
||||
|
||||
return output + "\n"
|
||||
}
|
||||
}
|
||||
|
||||
// Wrapper for mutable histogram
|
||||
final class MutableHistogram {
|
||||
var histogram: HistogramMetric
|
||||
|
||||
init() {
|
||||
self.histogram = HistogramMetric()
|
||||
}
|
||||
|
||||
func observe(_ value: Double, labels: [String: String]) {
|
||||
histogram.observe(value, labels: labels)
|
||||
}
|
||||
|
||||
func toPrometheus(name: String) -> String {
|
||||
histogram.toPrometheus(name: name)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
import Metal
|
||||
import RDMAKit
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// RDMA Distribution Service — Pipeline Parallelism
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public actor RDMADistributionService {
|
||||
public enum Role: Sendable {
|
||||
case primary(splitLayer: Int)
|
||||
case secondary
|
||||
}
|
||||
|
||||
public struct Config: Sendable {
|
||||
public let role: Role
|
||||
public let peerAddress: String?
|
||||
public let peerPort: UInt16
|
||||
|
||||
public init(role: Role, peerAddress: String? = nil, peerPort: UInt16 = 0) {
|
||||
self.role = role
|
||||
self.peerAddress = peerAddress
|
||||
self.peerPort = peerPort
|
||||
}
|
||||
}
|
||||
|
||||
// RDMA state
|
||||
private var context: RDMAContext?
|
||||
private var pd: RDMAProtectionDomain?
|
||||
private var sendCQ: RDMACompletionQueue?
|
||||
private var recvCQ: RDMACompletionQueue?
|
||||
private var qp: RDMAQueuePair?
|
||||
|
||||
// Registered GPU buffers
|
||||
private var registeredHiddenState: RegisteredMemory?
|
||||
|
||||
public let config: Config
|
||||
public let discovery: RDMADiscovery
|
||||
|
||||
public init(config: Config) {
|
||||
self.config = config
|
||||
self.discovery = RDMADiscovery()
|
||||
}
|
||||
|
||||
/// Initialize RDMA device and resources. Returns false if no RDMA device is available.
|
||||
@discardableResult
|
||||
public func initialize() throws -> Bool {
|
||||
let devices = discovery.listDevices()
|
||||
guard let device = devices.first else {
|
||||
print(" RDMA: No devices found — running in local-only mode")
|
||||
return false
|
||||
}
|
||||
|
||||
let tbInfo = device.isThunderbolt ? " (Thunderbolt)" : ""
|
||||
print(" RDMA: Found device '\(device.name)'\(tbInfo)")
|
||||
context = try discovery.openDevice(device)
|
||||
let attrs = try context!.queryDeviceAttributes()
|
||||
print(" RDMA: Firmware \(attrs.fwVer), max MR size \(attrs.maxMRSize) bytes")
|
||||
|
||||
pd = try RDMAProtectionDomain(context: context!)
|
||||
sendCQ = try RDMACompletionQueue(context: context!, cqe: 128)
|
||||
recvCQ = try RDMACompletionQueue(context: context!, cqe: 128)
|
||||
qp = try RDMAQueuePair(pd: pd!, sendCQ: sendCQ!, recvCQ: recvCQ!)
|
||||
try qp!.modifyToInit()
|
||||
print(" RDMA: QP initialized in RC mode")
|
||||
|
||||
if case .primary = config.role {
|
||||
try qp!.modifyToRTR(destQPN: 0)
|
||||
try qp!.modifyToRTS()
|
||||
}
|
||||
|
||||
print(" RDMA: Ready")
|
||||
return true
|
||||
}
|
||||
|
||||
/// Register a Metal buffer for remote RDMA access
|
||||
public func registerBuffer(_ buffer: MTLBuffer) throws -> RegisteredMemory? {
|
||||
guard let pd else { return nil }
|
||||
let reg = try pd.registerMTLBuffer(
|
||||
UnsafeMutableRawPointer(buffer.contents()),
|
||||
length: buffer.length
|
||||
)
|
||||
print(" RDMA: Registered buffer \(buffer.length) bytes (lkey=\(reg.lkey), rkey=\(reg.rkey))")
|
||||
return reg
|
||||
}
|
||||
|
||||
/// Register the model's hidden state buffer for pipeline-split transfer
|
||||
public func registerHiddenState(_ model: E4BModel) throws {
|
||||
let buf = model.temps.io
|
||||
let reg = try registerBuffer(buf)
|
||||
registeredHiddenState = reg
|
||||
print(" RDMA: Hidden state buffer registered (\(buf.length) bytes)")
|
||||
}
|
||||
|
||||
/// Split point for pipeline parallelism
|
||||
public var splitLayer: Int? {
|
||||
if case .primary(let layer) = config.role { return layer }
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Server-Sent Events (SSE) Support
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// SSE Event structure
|
||||
public struct SSEEvent {
|
||||
public let id: String?
|
||||
public let event: String?
|
||||
public let data: String
|
||||
public let retry: Int?
|
||||
|
||||
public init(id: String? = nil, event: String? = nil, data: String, retry: Int? = nil) {
|
||||
self.id = id
|
||||
self.event = event
|
||||
self.data = data
|
||||
self.retry = retry
|
||||
}
|
||||
|
||||
/// Format as SSE string
|
||||
public func format() -> String {
|
||||
var output = ""
|
||||
|
||||
if let id = id {
|
||||
output += "id: \(id)\n"
|
||||
}
|
||||
|
||||
if let event = event {
|
||||
output += "event: \(event)\n"
|
||||
}
|
||||
|
||||
if let retry = retry {
|
||||
output += "retry: \(retry)\n"
|
||||
}
|
||||
|
||||
// Handle multi-line data
|
||||
let lines = data.split(separator: "\n")
|
||||
for line in lines {
|
||||
output += "data: \(line)\n"
|
||||
}
|
||||
|
||||
output += "\n"
|
||||
return output
|
||||
}
|
||||
}
|
||||
|
||||
/// SSE Stream Generator
|
||||
public final class SSEStream {
|
||||
private var buffer = ""
|
||||
|
||||
public init() {}
|
||||
|
||||
/// Add event to stream
|
||||
public func add(event: SSEEvent) -> String {
|
||||
let formatted = event.format()
|
||||
buffer += formatted
|
||||
return formatted
|
||||
}
|
||||
|
||||
/// Create chat completion chunk
|
||||
public static func chatChunk(
|
||||
id: String,
|
||||
model: String,
|
||||
content: String? = nil,
|
||||
role: String? = nil,
|
||||
finishReason: String? = nil
|
||||
) -> SSEEvent {
|
||||
var delta: [String: String] = [:]
|
||||
if let role = role {
|
||||
delta["role"] = role
|
||||
}
|
||||
if let content = content {
|
||||
delta["content"] = content
|
||||
}
|
||||
|
||||
let chunk: [String: Any] = [
|
||||
"id": id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": Int(Date().timeIntervalSince1970),
|
||||
"model": model,
|
||||
"choices": [
|
||||
[
|
||||
"index": 0,
|
||||
"delta": delta,
|
||||
"finish_reason": finishReason as Any
|
||||
]
|
||||
]
|
||||
]
|
||||
|
||||
guard let jsonData = try? JSONSerialization.data(withJSONObject: chunk),
|
||||
let jsonString = String(data: jsonData, encoding: .utf8) else {
|
||||
return SSEEvent(data: "")
|
||||
}
|
||||
|
||||
return SSEEvent(data: jsonString)
|
||||
}
|
||||
|
||||
/// Create text completion chunk
|
||||
public static func textChunk(
|
||||
id: String,
|
||||
model: String,
|
||||
text: String,
|
||||
finishReason: String? = nil
|
||||
) -> SSEEvent {
|
||||
let chunk: [String: Any] = [
|
||||
"id": id,
|
||||
"object": "text_completion",
|
||||
"created": Int(Date().timeIntervalSince1970),
|
||||
"model": model,
|
||||
"choices": [
|
||||
[
|
||||
"index": 0,
|
||||
"text": text,
|
||||
"finish_reason": finishReason as Any
|
||||
]
|
||||
]
|
||||
]
|
||||
|
||||
guard let jsonData = try? JSONSerialization.data(withJSONObject: chunk),
|
||||
let jsonString = String(data: jsonData, encoding: .utf8) else {
|
||||
return SSEEvent(data: "")
|
||||
}
|
||||
|
||||
return SSEEvent(data: jsonString)
|
||||
}
|
||||
|
||||
/// Create [DONE] event
|
||||
public static func done() -> SSEEvent {
|
||||
SSEEvent(data: "[DONE]")
|
||||
}
|
||||
|
||||
/// Create error event
|
||||
public static func error(message: String) -> SSEEvent {
|
||||
let error: [String: Any] = [
|
||||
"error": [
|
||||
"message": message,
|
||||
"type": "invalid_request_error",
|
||||
"code": nil
|
||||
]
|
||||
]
|
||||
|
||||
guard let jsonData = try? JSONSerialization.data(withJSONObject: error),
|
||||
let jsonString = String(data: jsonData, encoding: .utf8) else {
|
||||
return SSEEvent(data: "")
|
||||
}
|
||||
|
||||
return SSEEvent(event: "error", data: jsonString)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
import Foundation
|
||||
import MarkBase
|
||||
import Hummingbird
|
||||
|
||||
struct SimpleServerApp {
|
||||
static func main() async throws {
|
||||
// 支持命令行參數選擇模型
|
||||
let args = CommandLine.arguments
|
||||
let modelName = args.count > 1 ? args[1] : "E4B-MarkBase"
|
||||
let port = args.count > 2 ? Int(args[2]) ?? 8080 : 8080
|
||||
|
||||
let modelPath = NSString(string: "~/MarkBaseEngine/models/\(modelName)").expandingTildeInPath
|
||||
|
||||
print("═══════════════════════════════════════════════════════════════════")
|
||||
print(" MarkBaseEngine Server")
|
||||
print("═══════════════════════════════════════════════════════════════════")
|
||||
print(" Model: \(modelName)")
|
||||
print(" Port: \(port)")
|
||||
print(" Path: \(modelPath)")
|
||||
print("")
|
||||
|
||||
let engine = try MarkBaseEngine(autoCompile: true)
|
||||
let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
|
||||
let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
|
||||
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
|
||||
|
||||
print("✓ E4B loaded (\(model.numHiddenLayers) layers)")
|
||||
|
||||
let router = Router()
|
||||
|
||||
let layers = model.numHiddenLayers
|
||||
|
||||
@Sendable func helpJSON() -> String {
|
||||
return """
|
||||
{
|
||||
"server": {
|
||||
"name": "MarkBaseEngine",
|
||||
"version": "1.0.0",
|
||||
"model": "E4B (\(layers) layers)",
|
||||
"framework": "Hummingbird 2.x + Metal GPU",
|
||||
"platform": "Apple Silicon",
|
||||
"base_url": "http://localhost:8080"
|
||||
},
|
||||
"endpoints": [
|
||||
{
|
||||
"method": "GET",
|
||||
"path": "/",
|
||||
"summary": "API help and documentation",
|
||||
"content_type": "application/json",
|
||||
"curl": "curl http://localhost:8080/"
|
||||
},
|
||||
{
|
||||
"method": "GET",
|
||||
"path": "/help",
|
||||
"summary": "API help and documentation (alias)",
|
||||
"content_type": "application/json",
|
||||
"curl": "curl http://localhost:8080/help"
|
||||
},
|
||||
{
|
||||
"method": "GET",
|
||||
"path": "/health",
|
||||
"summary": "Health check - returns server status and model information",
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Server is healthy",
|
||||
"body": "OK"
|
||||
}
|
||||
},
|
||||
"curl": "curl http://localhost:8080/health"
|
||||
},
|
||||
{
|
||||
"method": "GET",
|
||||
"path": "/v1/models",
|
||||
"summary": "List available models (OpenAI-compatible)",
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Model list",
|
||||
"body": {
|
||||
"id": "e4b",
|
||||
"object": "model",
|
||||
"owned_by": "markbase"
|
||||
}
|
||||
}
|
||||
},
|
||||
"curl": "curl http://localhost:8080/v1/models"
|
||||
},
|
||||
{
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"summary": "Chat completion (OpenAI-compatible API)",
|
||||
"description": "Generate chat responses using the E4B model. Supports text-only input via messages array.",
|
||||
"request": {
|
||||
"body": {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.95,
|
||||
"top_k": 40,
|
||||
"stream": false
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Successful completion",
|
||||
"body": {
|
||||
"id": "chatcmpl-xxx",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "e4b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I help you?"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"400": {
|
||||
"description": "Invalid request - missing or malformed messages field"
|
||||
}
|
||||
},
|
||||
"parameters": [
|
||||
{"name": "messages", "type": "array", "required": true, "description": "Array of message objects with 'role' (system/user/assistant/tool) and 'content' (string). Supports 'tool_calls' in assistant messages and 'tool_call_id'/'name' in tool messages."},
|
||||
{"name": "max_tokens", "type": "integer", "required": false, "default": 100, "description": "Maximum number of tokens to generate (1-4096)"},
|
||||
{"name": "temperature", "type": "float", "required": false, "default": 0.7, "description": "Sampling temperature (0.0-2.0)"},
|
||||
{"name": "top_p", "type": "float", "required": false, "description": "Nucleus sampling threshold (0.0-1.0)"},
|
||||
{"name": "top_k", "type": "integer", "required": false, "description": "Top-k sampling count"},
|
||||
{"name": "tools", "type": "array", "required": false, "description": "Array of tool definitions (OpenAI format) for function calling. Each tool has 'type' (function) and 'function' with 'name', 'description', 'parameters'."},
|
||||
{"name": "stream", "type": "boolean", "required": false, "default": false, "description": "Enable streaming response (not yet implemented)"}
|
||||
],
|
||||
"curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Hello\"}],\"max_tokens\":100}'",
|
||||
"examples": [
|
||||
{
|
||||
"title": "Text completion",
|
||||
"curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Hello\"}],\"max_tokens\":100}'"
|
||||
},
|
||||
{
|
||||
"title": "Function calling",
|
||||
"curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Search for cats\"}],\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"find_file\",\"description\":\"Search for files\",\"parameters\":{\"type\":\"object\",\"properties\":{\"query\":{\"type\":\"string\"}},\"required\":[\"query\"]}}}],\"max_tokens\":200}'"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"schema": {
|
||||
"content_type": "application/json",
|
||||
"error_format": {
|
||||
"error": {
|
||||
"message": "Error description",
|
||||
"type": "error_type",
|
||||
"code": 400
|
||||
}
|
||||
},
|
||||
"error_codes": [
|
||||
{"code": 400, "type": "invalid_request_error", "description": "Invalid request parameters"},
|
||||
{"code": 404, "type": "not_found_error", "description": "Resource not found"},
|
||||
{"code": 500, "type": "server_error", "description": "Internal server error"}
|
||||
]
|
||||
},
|
||||
"documentation": {
|
||||
"api_spec": "docs/API_SPEC.md",
|
||||
"api_reference": "docs/API.md",
|
||||
"deployment": "docs/DEPLOYMENT.md",
|
||||
"performance": "docs/PERFORMANCE.md"
|
||||
},
|
||||
"notes": [
|
||||
"All responses are in JSON format",
|
||||
"Text generation only (multimodal not yet supported via API)",
|
||||
"E4B model with 42 layers, ~4B parameters",
|
||||
"For multimodal (vision/audio) support, use the MarkBase Swift library directly",
|
||||
"Streaming support is planned but not yet implemented",
|
||||
"Function calling uses native Gemma 4 special tokens",
|
||||
"Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling"
|
||||
]
|
||||
}
|
||||
"""
|
||||
}
|
||||
|
||||
@Sendable func healthResponse() -> String {
|
||||
return "{\"status\":\"healthy\",\"model\":\"e4b\",\"layers\":\(layers)}"
|
||||
}
|
||||
|
||||
router.get("/") { _, _ in
|
||||
return helpJSON()
|
||||
}
|
||||
|
||||
router.get("/help") { _, _ in
|
||||
return helpJSON()
|
||||
}
|
||||
|
||||
router.get("/health") { _, _ in
|
||||
return healthResponse()
|
||||
}
|
||||
|
||||
router.get("/v1/models") { _, _ in
|
||||
return "{\"id\":\"e4b\",\"object\":\"model\",\"owned_by\":\"markbase\"}"
|
||||
}
|
||||
|
||||
router.post("/v1/chat/completions") { request, _ in
|
||||
let buffer = try await request.body.collect(upTo: .max)
|
||||
let data = Data(buffer: buffer)
|
||||
|
||||
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
|
||||
let messages = json["messages"] as? [[String: Any]] else {
|
||||
return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400}"
|
||||
}
|
||||
|
||||
let maxTokens = (json["max_tokens"] as? Int) ?? 100
|
||||
let temperature = Float(json["temperature"] as? Double ?? 0.7)
|
||||
let topK = json["top_k"] as? Int
|
||||
let topP = (json["top_p"] as? Double).map { Float($0) }
|
||||
let tools = json["tools"] as? [[String: Any]]
|
||||
|
||||
let prompt = Gemma4Format.buildChatPrompt(messages: messages, tools: tools)
|
||||
|
||||
let promptTokens = tokenizer.encode(text: prompt)
|
||||
|
||||
let config = GenerationConfig(
|
||||
maxTokens: maxTokens,
|
||||
temperature: temperature,
|
||||
topK: topK,
|
||||
topP: topP
|
||||
)
|
||||
|
||||
let generatedTokens = try generator.generateTokens(promptTokens: promptTokens, config: config)
|
||||
|
||||
let id = UUID().uuidString
|
||||
let ts = Int(Date().timeIntervalSince1970)
|
||||
|
||||
// Check for tool calls (token ID 48 = <|tool_call>)
|
||||
if generatedTokens.contains(48) {
|
||||
var results: [ToolCallResult] = []
|
||||
var i = 0
|
||||
while i < generatedTokens.count {
|
||||
if generatedTokens[i] == 48 {
|
||||
i += 1
|
||||
var callTokens: [Int] = []
|
||||
while i < generatedTokens.count && generatedTokens[i] != 49 {
|
||||
callTokens.append(generatedTokens[i])
|
||||
i += 1
|
||||
}
|
||||
if i < generatedTokens.count && generatedTokens[i] == 49 {
|
||||
i += 1
|
||||
}
|
||||
let callText = tokenizer.decode(tokens: callTokens)
|
||||
if let colonIdx = callText.firstIndex(of: ":"),
|
||||
let braceIdx = callText.firstIndex(of: "{"),
|
||||
let endBrace = callText.lastIndex(of: "}") {
|
||||
let name = String(callText[callText.index(after: colonIdx)..<braceIdx]).trimmingCharacters(in: .whitespaces)
|
||||
let rawArgs = String(callText[callText.index(after: braceIdx)..<endBrace])
|
||||
let jsonArgs = Gemma4Format.gemma4ArgsToJSON(rawArgs)
|
||||
results.append(ToolCallResult(
|
||||
id: "call_\(UUID().uuidString.replacingOccurrences(of: "-", with: "").prefix(16))",
|
||||
function: ToolCallFunction(name: name, arguments: jsonArgs)
|
||||
))
|
||||
}
|
||||
} else {
|
||||
i += 1
|
||||
}
|
||||
}
|
||||
|
||||
let encoder = JSONEncoder()
|
||||
let callsData = try encoder.encode(results)
|
||||
let callsStr = String(data: callsData, encoding: .utf8) ?? "[]"
|
||||
|
||||
return """
|
||||
{"id":"chatcmpl-\(id)","object":"chat.completion","created":\(ts),"model":"e4b","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":\(callsStr)},"finish_reason":"tool_calls"}]}
|
||||
"""
|
||||
} else {
|
||||
let response = tokenizer.decode(tokens: generatedTokens)
|
||||
let trimmed = response.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
let escaped = trimmed
|
||||
.replacingOccurrences(of: "\\", with: "\\\\")
|
||||
.replacingOccurrences(of: "\"", with: "\\\"")
|
||||
.replacingOccurrences(of: "\n", with: "\\n")
|
||||
.replacingOccurrences(of: "\r", with: "\\r")
|
||||
.replacingOccurrences(of: "\t", with: "\\t")
|
||||
|
||||
return """
|
||||
{"id":"chatcmpl-\(id)","object":"chat.completion","created":\(ts),"model":"e4b","choices":[{"index":0,"message":{"role":"assistant","content":"\(escaped)"},"finish_reason":"stop"}]}
|
||||
"""
|
||||
}
|
||||
}
|
||||
|
||||
let app = Application(
|
||||
router: router,
|
||||
configuration: .init(address: .hostname("0.0.0.0", port: port))
|
||||
)
|
||||
|
||||
print("Server starting on port \(port)...")
|
||||
print("Endpoints:")
|
||||
print(" GET / - API help")
|
||||
print(" GET /help - API help")
|
||||
print(" GET /health - Health check")
|
||||
print(" GET /v1/models - Model list")
|
||||
print(" POST /v1/chat/completions - Chat completion")
|
||||
print("")
|
||||
print("Model: \(modelName)")
|
||||
if modelName.contains("E4B") {
|
||||
print(" ⚠️ E4B is multimodal (Vision + Audio + Text)")
|
||||
print(" For text-only, use: 12B-it-MLX-8bit or 31B-it-8bit")
|
||||
} else {
|
||||
print(" ✓ LLM for text generation")
|
||||
}
|
||||
print("")
|
||||
|
||||
try await app.run()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,408 @@
|
||||
import Accelerate
|
||||
import AVFoundation
|
||||
import CoreVideo
|
||||
import Foundation
|
||||
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
// Video Processor — Extract frames + audio from video files
|
||||
// ─────────────────────────────────────────────────────────────
|
||||
|
||||
public struct VideoFrame {
|
||||
public let index: Int
|
||||
public let timestamp: CMTime
|
||||
public let pixelBuffer: CVPixelBuffer
|
||||
}
|
||||
|
||||
public struct VideoData {
|
||||
public let frames: [VideoFrame]
|
||||
public let audioSamples: [Float]
|
||||
public let sampleRate: Int
|
||||
public let duration: CMTime
|
||||
public let naturalSize: CGSize
|
||||
public let estimatedFrameRate: Float
|
||||
}
|
||||
|
||||
public enum VideoProcessor {
|
||||
public struct Config {
|
||||
public let maxFrames: Int
|
||||
public let targetFPS: Float
|
||||
public let audioSampleRate: Int
|
||||
public let sceneThreshold: Double // 0…1; 0 = fixed FPS, >0 = scene detection
|
||||
|
||||
public init(maxFrames: Int = 64, targetFPS: Float = 2.0,
|
||||
audioSampleRate: Int = 16000, sceneThreshold: Double = 0.15) {
|
||||
self.maxFrames = maxFrames
|
||||
self.targetFPS = targetFPS
|
||||
self.audioSampleRate = audioSampleRate
|
||||
self.sceneThreshold = sceneThreshold
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute luminance histogram difference between two pixel buffers (0…1, higher = more different).
|
||||
public static func sceneDiff(_ a: CVPixelBuffer, _ b: CVPixelBuffer) -> Double {
|
||||
CVPixelBufferLockBaseAddress(a, .readOnly)
|
||||
CVPixelBufferLockBaseAddress(b, .readOnly)
|
||||
defer {
|
||||
CVPixelBufferUnlockBaseAddress(a, .readOnly)
|
||||
CVPixelBufferUnlockBaseAddress(b, .readOnly)
|
||||
}
|
||||
|
||||
let w = min(CVPixelBufferGetWidth(a), CVPixelBufferGetWidth(b))
|
||||
let h = min(CVPixelBufferGetHeight(a), CVPixelBufferGetHeight(b))
|
||||
let rowA = CVPixelBufferGetBytesPerRow(a)
|
||||
let rowB = CVPixelBufferGetBytesPerRow(b)
|
||||
guard let baseA = CVPixelBufferGetBaseAddress(a),
|
||||
let baseB = CVPixelBufferGetBaseAddress(b) else { return 0 }
|
||||
|
||||
let ptrA = baseA.assumingMemoryBound(to: UInt8.self)
|
||||
let ptrB = baseB.assumingMemoryBound(to: UInt8.self)
|
||||
|
||||
var histA = [Double](repeating: 0, count: 256)
|
||||
var histB = [Double](repeating: 0, count: 256)
|
||||
let total = w * h
|
||||
|
||||
for y in 0..<h {
|
||||
for x in 0..<w {
|
||||
// Luminance ≈ green channel (BGRA → index 1)
|
||||
histA[Int(ptrA[y * rowA + x * 4 + 1])] += 1
|
||||
histB[Int(ptrB[y * rowB + x * 4 + 1])] += 1
|
||||
}
|
||||
}
|
||||
|
||||
// Normalise and compute intersection
|
||||
var diff: Double = 0
|
||||
for i in 0..<256 {
|
||||
let normA = histA[i] / Double(total)
|
||||
let normB = histB[i] / Double(total)
|
||||
diff += abs(normA - normB)
|
||||
}
|
||||
return diff / 2 // 0..1
|
||||
}
|
||||
|
||||
public static func process(url: URL, config: Config = Config()) async throws -> VideoData {
|
||||
let asset = AVURLAsset(url: url)
|
||||
let duration = try await asset.load(.duration)
|
||||
|
||||
let videoTracks = try await asset.loadTracks(withMediaType: .video)
|
||||
let audioTracks = try await asset.loadTracks(withMediaType: .audio)
|
||||
|
||||
guard let videoTrack = videoTracks.first else {
|
||||
throw MarkBaseError.invalidParameter(parameter: "url", message: "No video track found")
|
||||
}
|
||||
|
||||
let naturalSize = try await videoTrack.load(.naturalSize)
|
||||
let nominalFrameRate = try await videoTrack.load(.nominalFrameRate)
|
||||
|
||||
// Read audio samples concurrently
|
||||
let audioSamples: [Float]
|
||||
if let audioTrack = audioTracks.first {
|
||||
audioSamples = try readAudioTrack(audioTrack, sampleRate: config.audioSampleRate)
|
||||
} else {
|
||||
audioSamples = []
|
||||
}
|
||||
|
||||
// Read video frames (scene detection or fixed FPS)
|
||||
let frames: [VideoFrame]
|
||||
if config.sceneThreshold > 0 {
|
||||
frames = try await readVideoTrackSceneDetect(
|
||||
videoTrack,
|
||||
duration: duration,
|
||||
threshold: config.sceneThreshold,
|
||||
maxFrames: config.maxFrames
|
||||
)
|
||||
print(" Scene detection: \(frames.count) keyframes")
|
||||
} else {
|
||||
frames = try await readVideoTrack(
|
||||
videoTrack,
|
||||
duration: duration,
|
||||
nominalFrameRate: nominalFrameRate,
|
||||
targetFPS: config.targetFPS,
|
||||
maxFrames: config.maxFrames
|
||||
)
|
||||
}
|
||||
|
||||
return VideoData(
|
||||
frames: frames,
|
||||
audioSamples: audioSamples,
|
||||
sampleRate: config.audioSampleRate,
|
||||
duration: duration,
|
||||
naturalSize: naturalSize,
|
||||
estimatedFrameRate: nominalFrameRate
|
||||
)
|
||||
}
|
||||
|
||||
// ── Video reading ─────────────────────────────────
|
||||
|
||||
private static func readVideoTrack(
|
||||
_ track: AVAssetTrack,
|
||||
duration: CMTime,
|
||||
nominalFrameRate: Float,
|
||||
targetFPS: Float,
|
||||
maxFrames: Int
|
||||
) async throws -> [VideoFrame] {
|
||||
let reader = try AVAssetReader(asset: track.asset!)
|
||||
|
||||
let formatDescriptions = try await track.load(.formatDescriptions)
|
||||
let pixelFormat: OSType
|
||||
if let firstDesc = formatDescriptions.first {
|
||||
pixelFormat = CMFormatDescriptionGetMediaSubType(firstDesc)
|
||||
} else {
|
||||
pixelFormat = kCVPixelFormatType_32BGRA
|
||||
}
|
||||
|
||||
let settings: [String: Any] = [
|
||||
kCVPixelBufferPixelFormatTypeKey as String: pixelFormat,
|
||||
kCVPixelBufferMetalCompatibilityKey as String: true,
|
||||
]
|
||||
|
||||
let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
|
||||
reader.add(output)
|
||||
reader.startReading()
|
||||
|
||||
defer {
|
||||
if reader.status == .reading {
|
||||
reader.cancelReading()
|
||||
}
|
||||
}
|
||||
|
||||
let stepSeconds = CMTime(value: CMTimeValue(1.0 / targetFPS), timescale: CMTimeScale(targetFPS * 100))
|
||||
.seconds
|
||||
|
||||
var frames: [VideoFrame] = []
|
||||
var lastSampleTime: Double = -stepSeconds
|
||||
|
||||
while reader.status == .reading, frames.count < maxFrames {
|
||||
guard let sampleBuffer = output.copyNextSampleBuffer() else { break }
|
||||
|
||||
let presentationTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
|
||||
let timeSeconds = presentationTime.seconds
|
||||
|
||||
if timeSeconds - lastSampleTime >= stepSeconds {
|
||||
if let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) {
|
||||
frames.append(VideoFrame(
|
||||
index: frames.count,
|
||||
timestamp: presentationTime,
|
||||
pixelBuffer: pixelBuffer
|
||||
))
|
||||
}
|
||||
lastSampleTime = timeSeconds
|
||||
}
|
||||
}
|
||||
|
||||
return frames
|
||||
}
|
||||
|
||||
// ── Scene-detection reading ─────────────────────────
|
||||
|
||||
private static func readVideoTrackSceneDetect(
|
||||
_ track: AVAssetTrack,
|
||||
duration: CMTime,
|
||||
threshold: Double,
|
||||
maxFrames: Int
|
||||
) async throws -> [VideoFrame] {
|
||||
let reader = try AVAssetReader(asset: track.asset!)
|
||||
let formatDescriptions = try await track.load(.formatDescriptions)
|
||||
let pixelFormat: OSType
|
||||
if let firstDesc = formatDescriptions.first {
|
||||
pixelFormat = CMFormatDescriptionGetMediaSubType(firstDesc)
|
||||
} else {
|
||||
pixelFormat = kCVPixelFormatType_32BGRA
|
||||
}
|
||||
|
||||
let settings: [String: Any] = [
|
||||
kCVPixelBufferPixelFormatTypeKey as String: pixelFormat,
|
||||
kCVPixelBufferMetalCompatibilityKey as String: true,
|
||||
]
|
||||
|
||||
let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
|
||||
reader.add(output)
|
||||
reader.startReading()
|
||||
|
||||
defer {
|
||||
if reader.status == .reading { reader.cancelReading() }
|
||||
}
|
||||
|
||||
var frames: [VideoFrame] = []
|
||||
var prevBuffer: CVPixelBuffer?
|
||||
|
||||
while reader.status == .reading, frames.count < maxFrames {
|
||||
guard let sampleBuffer = output.copyNextSampleBuffer() else { break }
|
||||
guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { continue }
|
||||
|
||||
if let prev = prevBuffer {
|
||||
let diff = sceneDiff(prev, pixelBuffer)
|
||||
if diff >= threshold || frames.isEmpty {
|
||||
frames.append(VideoFrame(
|
||||
index: frames.count,
|
||||
timestamp: CMSampleBufferGetPresentationTimeStamp(sampleBuffer),
|
||||
pixelBuffer: pixelBuffer
|
||||
))
|
||||
}
|
||||
} else {
|
||||
frames.append(VideoFrame(
|
||||
index: frames.count,
|
||||
timestamp: CMSampleBufferGetPresentationTimeStamp(sampleBuffer),
|
||||
pixelBuffer: pixelBuffer
|
||||
))
|
||||
}
|
||||
prevBuffer = pixelBuffer
|
||||
}
|
||||
|
||||
return frames
|
||||
}
|
||||
|
||||
// ── Audio reading ─────────────────────────────────
|
||||
|
||||
private static func readAudioTrack(_ track: AVAssetTrack, sampleRate: Int) throws -> [Float] {
|
||||
let reader = try AVAssetReader(asset: track.asset!)
|
||||
|
||||
let settings: [String: Any] = [
|
||||
AVFormatIDKey: kAudioFormatLinearPCM,
|
||||
AVLinearPCMIsFloatKey: true,
|
||||
AVLinearPCMBitDepthKey: 32,
|
||||
AVNumberOfChannelsKey: 1,
|
||||
AVSampleRateKey: sampleRate,
|
||||
]
|
||||
|
||||
let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
|
||||
reader.add(output)
|
||||
reader.startReading()
|
||||
|
||||
defer {
|
||||
if reader.status == .reading {
|
||||
reader.cancelReading()
|
||||
}
|
||||
}
|
||||
|
||||
var samples: [Float] = []
|
||||
while reader.status == .reading {
|
||||
guard let sampleBuffer = output.copyNextSampleBuffer() else { break }
|
||||
guard let blockBuffer = CMSampleBufferGetDataBuffer(sampleBuffer) else { continue }
|
||||
|
||||
let length = CMBlockBufferGetDataLength(blockBuffer)
|
||||
var data = [Float](repeating: 0, count: length / MemoryLayout<Float>.stride)
|
||||
CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data)
|
||||
samples.append(contentsOf: data)
|
||||
}
|
||||
|
||||
return samples
|
||||
}
|
||||
|
||||
// ── Pixel buffer → float array ────────────────────
|
||||
|
||||
public static func pixelBufferToFloats(_ pixelBuffer: CVPixelBuffer) -> [Float] {
|
||||
CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
|
||||
defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }
|
||||
|
||||
let width = CVPixelBufferGetWidth(pixelBuffer)
|
||||
let height = CVPixelBufferGetHeight(pixelBuffer)
|
||||
let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
|
||||
|
||||
guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else { return [] }
|
||||
|
||||
var floats = [Float](repeating: 0, count: width * height * 3) // RGB
|
||||
let pixelBuffer = baseAddress.assumingMemoryBound(to: UInt8.self)
|
||||
|
||||
for y in 0..<height {
|
||||
for x in 0..<width {
|
||||
let offset = y * bytesPerRow + x * 4
|
||||
let b = Float(pixelBuffer[offset]) / 255.0
|
||||
let g = Float(pixelBuffer[offset + 1]) / 255.0
|
||||
let r = Float(pixelBuffer[offset + 2]) / 255.0
|
||||
let pixelOffset = (y * width + x) * 3
|
||||
floats[pixelOffset] = r
|
||||
floats[pixelOffset + 1] = g
|
||||
floats[pixelOffset + 2] = b
|
||||
}
|
||||
}
|
||||
|
||||
return floats
|
||||
}
|
||||
|
||||
/// Resize pixel buffer to target size using vImage
|
||||
public static func resizePixelBuffer(
|
||||
_ pixelBuffer: CVPixelBuffer,
|
||||
targetWidth: Int,
|
||||
targetHeight: Int
|
||||
) -> CVPixelBuffer? {
|
||||
let srcWidth = CVPixelBufferGetWidth(pixelBuffer)
|
||||
let srcHeight = CVPixelBufferGetHeight(pixelBuffer)
|
||||
|
||||
var srcBuffer = vImage_Buffer()
|
||||
CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
|
||||
srcBuffer.data = CVPixelBufferGetBaseAddress(pixelBuffer)
|
||||
srcBuffer.width = vImagePixelCount(srcWidth)
|
||||
srcBuffer.height = vImagePixelCount(srcHeight)
|
||||
srcBuffer.rowBytes = CVPixelBufferGetBytesPerRow(pixelBuffer)
|
||||
|
||||
var destPixelBuffer: CVPixelBuffer?
|
||||
let attrs: [String: Any] = [
|
||||
kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
|
||||
kCVPixelBufferMetalCompatibilityKey as String: true,
|
||||
]
|
||||
CVPixelBufferCreate(
|
||||
kCFAllocatorDefault,
|
||||
targetWidth, targetHeight,
|
||||
kCVPixelFormatType_32BGRA,
|
||||
attrs as CFDictionary,
|
||||
&destPixelBuffer
|
||||
)
|
||||
|
||||
guard let dest = destPixelBuffer else {
|
||||
CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
|
||||
return nil
|
||||
}
|
||||
|
||||
CVPixelBufferLockBaseAddress(dest, [])
|
||||
var destBuffer = vImage_Buffer()
|
||||
destBuffer.data = CVPixelBufferGetBaseAddress(dest)
|
||||
destBuffer.width = vImagePixelCount(targetWidth)
|
||||
destBuffer.height = vImagePixelCount(targetHeight)
|
||||
destBuffer.rowBytes = CVPixelBufferGetBytesPerRow(dest)
|
||||
|
||||
let scale = vImageScale_ARGB8888(&srcBuffer, &destBuffer, nil, vImage_Flags(kvImageNoFlags))
|
||||
|
||||
CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
|
||||
CVPixelBufferUnlockBaseAddress(dest, [])
|
||||
|
||||
return scale == kvImageNoError ? dest : nil
|
||||
}
|
||||
|
||||
/// Frame to patch embeddings (simple 16×16 grid)
|
||||
public static func frameToPatchEmbeddings(
|
||||
_ pixelBuffer: CVPixelBuffer,
|
||||
patchSize: Int = 16
|
||||
) -> (embeddings: [Float], numPatches: Int, patchDim: Int) {
|
||||
let rgb = pixelBufferToFloats(pixelBuffer)
|
||||
let width = CVPixelBufferGetWidth(pixelBuffer)
|
||||
let height = CVPixelBufferGetHeight(pixelBuffer)
|
||||
|
||||
let numPatchesH = height / patchSize
|
||||
let numPatchesW = width / patchSize
|
||||
let numPatches = numPatchesH * numPatchesW
|
||||
let patchDim = patchSize * patchSize * 3
|
||||
|
||||
var embeddings = [Float](repeating: 0, count: numPatches * patchDim)
|
||||
|
||||
for ph in 0..<numPatchesH {
|
||||
for pw in 0..<numPatchesW {
|
||||
let patchIdx = ph * numPatchesW + pw
|
||||
for py in 0..<patchSize {
|
||||
for px in 0..<patchSize {
|
||||
let srcY = ph * patchSize + py
|
||||
let srcX = pw * patchSize + px
|
||||
let srcIdx = (srcY * width + srcX) * 3
|
||||
let dstIdx = (patchIdx * patchDim) + (py * patchSize + px) * 3
|
||||
if srcIdx + 2 < rgb.count, dstIdx + 2 < embeddings.count {
|
||||
embeddings[dstIdx] = rgb[srcIdx]
|
||||
embeddings[dstIdx + 1] = rgb[srcIdx + 1]
|
||||
embeddings[dstIdx + 2] = rgb[srcIdx + 2]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return (embeddings, numPatches, patchDim)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
import Foundation
|
||||
|
||||
// Entry point — avoids @main conflict with top-level code
|
||||
Task {
|
||||
do {
|
||||
try await SimpleServerApp.main()
|
||||
} catch {
|
||||
print("Server error: \(error)")
|
||||
exit(1)
|
||||
}
|
||||
}
|
||||
dispatchMain()
|
||||
Reference in New Issue
Block a user