""" COLLISION — OpenAI-Compatible Local REST Server. Provides a drop-in OpenAI-compatible API server supporting: - POST /v1/chat/completions - POST /v1/completions - GET /v1/models - GET /health Connect immediately with Cursor, Continue.dev, OpenWebUI, LibreChat, AutoGen, and LangChain! Usage: python api_server.py --port 8000 python api_server.py --host 0.0.0.0 --port 8000 --weights model.pt """ import os import sys import json import time import uuid import argparse from http.server import HTTPServer, BaseHTTPRequestHandler from typing import Dict, Any, List CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) if CURRENT_DIR not in sys.path: sys.path.insert(0, CURRENT_DIR) import torch from configuration_collision import CollisionConfig from modeling_collision import CollisionForCausalLM from tokenization_collision import CollisionTokenizer # Global model state GLOBAL_MODEL = None GLOBAL_TOKENIZER = None MODEL_ID = "collision-10M/Collision-1B" def load_global_model(weights_path: str, config_path: str, tokenizer_dir: str): global GLOBAL_MODEL, GLOBAL_TOKENIZER print(f"[*] Initializing COLLISION-1B runtime...") with open(config_path, "r", encoding="utf-8") as f: cfg_dict = json.load(f) cfg = CollisionConfig( vocab_size=cfg_dict.get("vocab_size", 32000), max_seq_len=cfg_dict.get("max_seq_len", 1024), d_model=cfg_dict.get("d_model", 2048), n_layer=cfg_dict.get("n_layer", 24), n_head=cfg_dict.get("n_head", 16), d_ff=cfg_dict.get("d_ff", 5376), dropout=cfg_dict.get("dropout", 0.1), tie_embeddings=cfg_dict.get("tie_embeddings", True) ) model = CollisionForCausalLM(cfg) if os.path.exists(weights_path): checkpoint = torch.load(weights_path, map_location="cpu") state_dict = checkpoint.get("model_state_dict", checkpoint) model.load_state_dict(state_dict, strict=False) print(f"[OK] Checkpoint loaded from {weights_path}") else: print(f"[WARN] Weights not found at {weights_path}, using initialized weights.") model.eval() GLOBAL_MODEL = model vocab_file = os.path.join(tokenizer_dir, "vocab.json") merges_file = os.path.join(tokenizer_dir, "merges.json") GLOBAL_TOKENIZER = CollisionTokenizer(vocab_file=vocab_file, merges_file=merges_file) print(f"[OK] Collision runtime ready for OpenAI requests.\n") def generate_completion(prompt: str, max_tokens: int = 128, temperature: float = 0.7, top_p: float = 0.9) -> str: if GLOBAL_MODEL is None or GLOBAL_TOKENIZER is None: return "Model runtime not loaded." tokens = GLOBAL_TOKENIZER.encode_to_ids(prompt, bos=True) if not tokens: tokens = [GLOBAL_TOKENIZER.special_tokens_map_dict["[BOS]"]] input_ids = torch.tensor([tokens], dtype=torch.long) generated = list(tokens) for _ in range(max_tokens): if input_ids.size(1) >= GLOBAL_MODEL.config.max_seq_len: input_ids = input_ids[:, -GLOBAL_MODEL.config.max_seq_len:] with torch.no_grad(): outputs = GLOBAL_MODEL(input_ids) logits = outputs.logits[0, -1, :] if temperature <= 0.01: next_token = torch.argmax(logits).item() else: logits = logits / temperature probs = torch.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1).item() if next_token == GLOBAL_TOKENIZER.special_tokens_map_dict["[EOS]"]: break generated.append(next_token) input_ids = torch.tensor([generated], dtype=torch.long) completion_ids = generated[len(tokens):] return GLOBAL_TOKENIZER.decode_from_ids(completion_ids, skip_special_tokens=True) class OpenAIApiHandler(BaseHTTPRequestHandler): def _set_headers(self, status: int = 200, content_type: str = "application/json"): self.send_response(status) self.send_header("Content-Type", content_type) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization") self.end_headers() def do_OPTIONS(self): self._set_headers(200) def do_GET(self): if self.path == "/health" or self.path == "/": self._set_headers(200) self.wfile.write(json.dumps({ "status": "healthy", "model": MODEL_ID, "parameters": "999.38M", "uptime": "active" }).encode("utf-8")) elif self.path == "/v1/models": self._set_headers(200) response = { "object": "list", "data": [ { "id": MODEL_ID, "object": "model", "created": int(time.time()), "owned_by": "collision", "permission": [], "root": MODEL_ID, "parent": None }, { "id": "collision", "object": "model", "created": int(time.time()), "owned_by": "collision" } ] } self.wfile.write(json.dumps(response).encode("utf-8")) else: self._set_headers(404) self.wfile.write(json.dumps({"error": "Endpoint not found"}).encode("utf-8")) def do_POST(self): content_length = int(self.headers.get("Content-Length", 0)) post_data = self.rfile.read(content_length) try: req = json.loads(post_data.decode("utf-8")) except Exception: self._set_headers(400) self.wfile.write(json.dumps({"error": "Invalid JSON"}).encode("utf-8")) return if self.path == "/v1/chat/completions": messages = req.get("messages", []) prompt = "" for m in messages: role = m.get("role", "user") content = m.get("content", "") if role == "system": prompt += f"System: {content}\n" elif role == "user": prompt += f"User: {content}\n" elif role == "assistant": prompt += f"Assistant: {content}\n" prompt += "Assistant: " max_tokens = req.get("max_tokens", 128) temperature = float(req.get("temperature", 0.7)) top_p = float(req.get("top_p", 0.9)) t0 = time.time() text = generate_completion(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p) elapsed = time.time() - t0 chat_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" res = { "id": chat_id, "object": "chat.completion", "created": int(time.time()), "model": req.get("model", MODEL_ID), "choices": [ { "index": 0, "message": { "role": "assistant", "content": text }, "finish_reason": "stop" } ], "usage": { "prompt_tokens": len(prompt.split()), "completion_tokens": len(text.split()), "total_tokens": len(prompt.split()) + len(text.split()) }, "latency_seconds": round(elapsed, 4) } self._set_headers(200) self.wfile.write(json.dumps(res).encode("utf-8")) elif self.path == "/v1/completions": prompt = req.get("prompt", "") max_tokens = req.get("max_tokens", 128) temperature = float(req.get("temperature", 0.7)) top_p = float(req.get("top_p", 0.9)) text = generate_completion(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p) cmpl_id = f"cmpl-{uuid.uuid4().hex[:12]}" res = { "id": cmpl_id, "object": "text_completion", "created": int(time.time()), "model": req.get("model", MODEL_ID), "choices": [ { "text": text, "index": 0, "logprobs": None, "finish_reason": "stop" } ], "usage": { "prompt_tokens": len(prompt.split()), "completion_tokens": len(text.split()), "total_tokens": len(prompt.split()) + len(text.split()) } } self._set_headers(200) self.wfile.write(json.dumps(res).encode("utf-8")) else: self._set_headers(404) self.wfile.write(json.dumps({"error": "Endpoint not found"}).encode("utf-8")) def run_server(host: str = "0.0.0.0", port: int = 8000, dry_run: bool = False): weights_path = os.path.join(CURRENT_DIR, "model.pt") config_path = os.path.join(CURRENT_DIR, "config.json") tokenizer_dir = os.path.join(CURRENT_DIR, "tokenizer") if not dry_run: load_global_model(weights_path, config_path, tokenizer_dir) server_address = (host, port) httpd = HTTPServer(server_address, OpenAIApiHandler) print(f"[*] COLLISION OpenAI-Compatible Server running at http://{host}:{port}") print(f"[*] Connect Cursor / Continue.dev / OpenWebUI to: http://localhost:{port}/v1\n") if dry_run: print("[OK] Dry run successful. Exiting.") return try: httpd.serve_forever() except KeyboardInterrupt: print("\nStopping server...") httpd.server_close() if __name__ == "__main__": parser = argparse.ArgumentParser(description="COLLISION OpenAI-Compatible API Server") parser.add_argument("--host", type=str, default="0.0.0.0", help="Host address") parser.add_argument("--port", type=int, default=8000, help="Port to listen on") parser.add_argument("--dry-run", action="store_true", help="Validate server startup and exit") args = parser.parse_args() run_server(host=args.host, port=args.port, dry_run=args.dry_run)