LisaMegaWatts/MonarchSLM
0
1#=2server.jl — OpenAI-compatible inference server for MonarchSLM3 4Serves a Lux.jl trained Monarch Mixer model (sub-quadratic sequence mixing,5RMSNorm, SwiGLU, weight-tied). Downloads artifacts from HuggingFace on first run.6 7Endpoints:8 GET / -> health check / API info9 GET /v1/models -> list available models10 POST /v1/chat/completions -> generate text (OpenAI format, streaming supported)11=#12 13include("checkpoint.jl")14using HTTP15using UUIDs16using Downloads17 18# ═══════════════════════════════════════════════════════════════════19# Download artifacts from HuggingFace20# ═══════════════════════════════════════════════════════════════════21 22const CKPT_DIR = "checkpoints"23const CKPT_PATH = joinpath(CKPT_DIR, "final.jld2")24const CONFIG_PATH = joinpath(CKPT_DIR, "config.toml")25const VOCAB_PATH = joinpath(CKPT_DIR, "vocab.json")26const MERGES_PATH = joinpath(CKPT_DIR, "merges.txt")27const HF_REPO = get(ENV, "HF_REPO", "LisaMegaWatts/MonarchSLM")28const PORT = parse(Int, get(ENV, "PORT", "7860"))29 30function download_from_hf(repo::String, filename::String, local_path::String)31 url = "https://huggingface.co/$repo/resolve/main/$filename"32 println("Downloading $url ...")33 mkpath(dirname(local_path))34 Downloads.download(url, local_path)35 sz = round(filesize(local_path) / 1024^2, digits=1)36 println(" -> $local_path ($sz MB)")37end38 39function ensure_artifacts()40 for (localpath, remote) in [(CKPT_PATH, "final.jld2"),41 (CONFIG_PATH, "config.toml"),42 (VOCAB_PATH, "vocab.json")]43 if !isfile(localpath)44 println("No local $remote found, downloading from $HF_REPO ...")45 try46 download_from_hf(HF_REPO, remote, localpath)47 catch e48 println("Download failed for $remote: $e")49 println("Place $remote at $localpath manually.")50 exit(1)51 end52 end53 end54 if !isfile(MERGES_PATH)55 println("Attempting to download merges.txt (optional, for BPE) ...")56 try57 download_from_hf(HF_REPO, "merges.txt", MERGES_PATH)58 catch e59 println(" merges.txt not found (will use char tokenizer if vocab is array format)")60 end61 end62end63 64# ═══════════════════════════════════════════════════════════════════65# Download and load model66# ═══════════════════════════════════════════════════════════════════67 68ensure_artifacts()69 70println("\nLoading model...")71const INF_MODEL = load_inference_model(CKPT_PATH, CONFIG_PATH, VOCAB_PATH, MERGES_PATH)72const CONFIG = INF_MODEL.config73const PS = INF_MODEL.ps74const TOKENIZER = INF_MODEL.tokenizer75const CACHES = INF_MODEL.caches76const MODEL_CREATED_AT = Int(floor(time()))77 78println("\nModel ready: arch=$(CONFIG.arch), vocab=$(CONFIG.vocab_size), embd=$(CONFIG.embed_dim), " *79 "layers=$(CONFIG.n_layers), monarch_heads=$(CONFIG.n_monarch_heads), ctx=$(CONFIG.context_length)")80 81# ═══════════════════════════════════════════════════════════════════82# HTTP helpers83# ═══════════════════════════════════════════════════════════════════84 85const CORS_HEADERS = [86 "Access-Control-Allow-Origin" => "*",87 "Access-Control-Allow-Methods" => "GET, POST, OPTIONS",88 "Access-Control-Allow-Headers" => "Content-Type, Authorization",89]90 91function json_response(status::Int, body; extra_headers=[])92 json_bytes = JSON3.write(body)93 headers = [94 "Content-Type" => "application/json",95 CORS_HEADERS...,96 extra_headers...97 ]98 return HTTP.Response(status, headers, json_bytes)99end100 101function cors_preflight()102 return HTTP.Response(204, CORS_HEADERS)103end104 105# ═══════════════════════════════════════════════════════════════════106# Extract prompt from OpenAI chat messages107# ═══════════════════════════════════════════════════════════════════108 109function extract_prompt(messages)110 if isempty(messages)111 return ""112 end113 for i in length(messages):-1:1114 role = string(get(messages[i], :role, ""))115 if role == "user"116 return string(get(messages[i], :content, ""))117 end118 end119 return string(get(messages[end], :content, ""))120end121 122# ═══════════════════════════════════════════════════════════════════123# SSE helpers124# ═══════════════════════════════════════════════════════════════════125 126function sse_line(data)127 return "data: $(JSON3.write(data))\n\n"128end129 130# ═══════════════════════════════════════════════════════════════════131# Request handler132# ═══════════════════════════════════════════════════════════════════133 134function handle_request(request::HTTP.Request)135 method = request.method136 target = request.target137 138 if method == "OPTIONS"139 return cors_preflight()140 end141 142 # GET / — health check143 if method == "GET" && target == "/"144 return json_response(200, Dict(145 "name" => "MonarchSLM",146 "version" => "1.0.0",147 "description" => "A Monarch Mixer model trained on classical philosophy texts",148 "architecture" => "Decoder-only (Monarch Mixer, RMSNorm, SwiGLU, weight-tied)",149 "model" => Dict(150 "arch" => CONFIG.arch,151 "vocab_size" => CONFIG.vocab_size,152 "embed_dim" => CONFIG.embed_dim,153 "n_layers" => CONFIG.n_layers,154 "n_monarch_heads" => CONFIG.n_monarch_heads,155 "conv_kernel_size" => CONFIG.conv_kernel_size,156 "context_length" => CONFIG.context_length157 ),158 "endpoints" => ["/v1/models", "/v1/chat/completions"],159 "features" => ["streaming", "OpenAI-compatible", "top-k", "top-p"],160 "compatible_with" => ["OpenAI API", "OpenRouter"]161 ))162 end163 164 # GET /v1/models165 if method == "GET" && target == "/v1/models"166 return json_response(200, Dict(167 "object" => "list",168 "data" => [Dict(169 "id" => "monarchslm-philosophy",170 "object" => "model",171 "created" => MODEL_CREATED_AT,172 "owned_by" => "monarchslm"173 )]174 ))175 end176 177 # POST /v1/chat/completions178 if method == "POST" && target == "/v1/chat/completions"179 local body180 try181 body = JSON3.read(String(request.body))182 catch e183 return json_response(400, Dict("error" => Dict(184 "message" => "Invalid JSON in request body",185 "type" => "invalid_request_error",186 "code" => "invalid_json")))187 end188 189 temperature = Float64(clamp(get(body, :temperature, 0.8), 0.01, 2.0))190 max_tokens = Int(clamp(get(body, :max_tokens, 200), 1, CONFIG.context_length))191 top_k_val = Int(clamp(get(body, :top_k, 40), 0, CONFIG.vocab_size))192 top_p_val = Float64(clamp(get(body, :top_p, 1.0), 0.0, 1.0))193 stream = Bool(get(body, :stream, false))194 195 messages = get(body, :messages, [])196 prompt_text = extract_prompt(messages)197 198 if stream199 completion_id = "chatcmpl-" * string(uuid4())200 created = Int(floor(time()))201 202 buf = IOBuffer()203 204 initial_chunk = Dict(205 "id" => completion_id,206 "object" => "chat.completion.chunk",207 "created" => created,208 "model" => "monarchslm-philosophy",209 "choices" => [Dict(210 "index" => 0,211 "delta" => Dict("role" => "assistant", "content" => ""),212 "finish_reason" => nothing213 )]214 )215 write(buf, sse_line(initial_chunk))216 217 token_count = Ref(0)218 219 generate_streaming(CONFIG, PS, TOKENIZER, prompt_text;220 max_tokens, temperature, top_k=top_k_val, top_p=top_p_val,221 caches=CACHES,222 on_token = function(token_str)223 token_count[] += 1224 chunk = Dict(225 "id" => completion_id,226 "object" => "chat.completion.chunk",227 "created" => created,228 "model" => "monarchslm-philosophy",229 "choices" => [Dict(230 "index" => 0,231 "delta" => Dict("content" => token_str),232 "finish_reason" => nothing233 )]234 )235 write(buf, sse_line(chunk))236 end)237 238 prompt_tokens = length(encode(TOKENIZER, prompt_text))239 finish_chunk = Dict(240 "id" => completion_id,241 "object" => "chat.completion.chunk",242 "created" => created,243 "model" => "monarchslm-philosophy",244 "choices" => [Dict(245 "index" => 0,246 "delta" => Dict(),247 "finish_reason" => token_count[] >= max_tokens ? "length" : "stop"248 )],249 "usage" => Dict(250 "prompt_tokens" => prompt_tokens,251 "completion_tokens" => token_count[],252 "total_tokens" => prompt_tokens + token_count[]253 )254 )255 write(buf, sse_line(finish_chunk))256 write(buf, "data: [DONE]\n\n")257 258 sse_body = take!(buf)259 headers = [260 "Content-Type" => "text/event-stream",261 "Cache-Control" => "no-cache",262 "X-Accel-Buffering" => "no",263 CORS_HEADERS...264 ]265 return HTTP.Response(200, headers, sse_body)266 267 else268 n_completions = Int(clamp(get(body, :n, 1), 1, 4))269 270 choices = []271 total_completion_tokens = 0272 for i in 1:n_completions273 text = generate_streaming(CONFIG, PS, TOKENIZER, prompt_text;274 max_tokens, temperature, top_k=top_k_val, top_p=top_p_val,275 caches=CACHES)276 finish_reason = "length" # generate_streaming always produces exactly max_tokens tokens277 push!(choices, Dict(278 "index" => i - 1,279 "message" => Dict("role" => "assistant", "content" => text),280 "finish_reason" => finish_reason))281 total_completion_tokens += max_tokens # count tokens, not decoded chars282 end283 284 prompt_tokens = length(encode(TOKENIZER, prompt_text))285 return json_response(200, Dict(286 "id" => "chatcmpl-" * string(uuid4()),287 "object" => "chat.completion",288 "created" => Int(floor(time())),289 "model" => "monarchslm-philosophy",290 "choices" => choices,291 "usage" => Dict(292 "prompt_tokens" => prompt_tokens,293 "completion_tokens" => total_completion_tokens,294 "total_tokens" => prompt_tokens + total_completion_tokens),295 "system_fingerprint" => "monarchslm-v1"))296 end297 end298 299 return json_response(404, Dict("error" => Dict(300 "message" => "Not found: $method $target",301 "type" => "invalid_request_error",302 "code" => "not_found")))303end304 305# ═══════════════════════════════════════════════════════════════════306# Start server307# ═══════════════════════════════════════════════════════════════════308 309println("\nMonarchSLM server starting on 0.0.0.0:$PORT ...")310println(" GET http://localhost:$PORT/")311println(" GET http://localhost:$PORT/v1/models")312println(" POST http://localhost:$PORT/v1/chat/completions")313println(" POST http://localhost:$PORT/v1/chat/completions (stream=true)")314println()315 316HTTP.serve(handle_request, "0.0.0.0", PORT)317