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))"
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user