| |
| """ |
| OpenAI-compatible API server for Ternary Transformer Engine. |
| Drop-in replacement for llama-server. |
| |
| (c) 2026 OpenTransformers Ltd / Scott Bisset |
| """ |
|
|
| import json |
| import time |
| import threading |
| from http.server import HTTPServer, BaseHTTPRequestHandler |
| from inference import TernaryQwen, Tokenizer, load_kernel |
| import os |
|
|
| MODEL_DIR = os.environ.get("TERNARY_MODEL_DIR", "deepseek-r1-1.5b-ternary") |
| TOKENIZER_DIR = os.environ.get("TOKENIZER_DIR", "deepseek-r1-1.5b-hf") |
| HOST = os.environ.get("HOST", "127.0.0.1") |
| PORT = int(os.environ.get("PORT", "8080")) |
|
|
| print("Loading ternary kernel...") |
| kernel = load_kernel(os.path.join(os.path.dirname(__file__), "ternary_kernel.so")) |
|
|
| print(f"Loading model from {MODEL_DIR}...") |
| model = TernaryQwen(MODEL_DIR, kernel) |
|
|
| print(f"Loading tokenizer from {TOKENIZER_DIR}...") |
| tokenizer = Tokenizer(TOKENIZER_DIR) |
|
|
| lock = threading.Lock() |
| print("Ready!") |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def do_POST(self): |
| if self.path == "/v1/chat/completions": |
| length = int(self.headers.get("Content-Length", 0)) |
| body = json.loads(self.rfile.read(length)) |
| |
| messages = body.get("messages", []) |
| max_tokens = body.get("max_tokens", 256) |
| temperature = body.get("temperature", 0.6) |
| top_p = body.get("top_p", 0.95) |
| |
| |
| prompt = tokenizer.apply_chat_template(messages) |
| input_ids = tokenizer.encode(prompt) |
| |
| |
| with lock: |
| gen_ids, stats = model.generate( |
| input_ids, |
| max_new_tokens=max_tokens, |
| temperature=temperature, |
| top_p=top_p |
| ) |
| |
| text = tokenizer.decode(gen_ids) |
| |
| response = { |
| "id": f"chatcmpl-ternary-{int(time.time())}", |
| "object": "chat.completion", |
| "created": int(time.time()), |
| "model": "DeepSeek-R1-Distill-Qwen-1.5B-TERNARY", |
| "choices": [{ |
| "index": 0, |
| "message": {"role": "assistant", "content": text}, |
| "finish_reason": "stop" |
| }], |
| "usage": { |
| "prompt_tokens": len(input_ids), |
| "completion_tokens": stats["tokens_generated"], |
| "total_tokens": len(input_ids) + stats["tokens_generated"] |
| }, |
| "timings": { |
| "prompt_n": stats["prefill_tokens"], |
| "prompt_ms": stats["prefill_ms"], |
| "predicted_n": stats["tokens_generated"], |
| "predicted_ms": stats["decode_ms"], |
| "predicted_per_second": stats["tok_per_sec"], |
| } |
| } |
| |
| self.send_response(200) |
| self.send_header("Content-Type", "application/json") |
| self.end_headers() |
| self.wfile.write(json.dumps(response).encode()) |
| else: |
| self.send_response(404) |
| self.end_headers() |
| |
| def do_GET(self): |
| if self.path == "/health": |
| self.send_response(200) |
| self.send_header("Content-Type", "application/json") |
| self.end_headers() |
| self.wfile.write(b'{"status":"ok","engine":"ternary-avx512"}') |
| else: |
| self.send_response(404) |
| self.end_headers() |
| |
| def log_message(self, format, *args): |
| pass |
|
|
| if __name__ == "__main__": |
| server = HTTPServer((HOST, PORT), Handler) |
| print(f"Ternary engine serving on {HOST}:{PORT}") |
| server.serve_forever() |
|
|
