#!/usr/bin/env python3 """ OpenAI-Compatible Local Inference Server for Q-TensorFormer. Provides: - GET /v1/models - POST /v1/chat/completions (with streaming SSE `stream=True` and non-streaming) - GET /v1/telemetry (real-time active TT-rank, energy usage, and KV memory) Compatible with Open-WebUI, Cursor, Continue.dev, and standard OpenAI SDK clients: from openai import OpenAI client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed") """ import argparse import asyncio import json import time import uuid from typing import AsyncGenerator, Dict, List, Optional, Union import torch from transformers import AutoTokenizer from configuration_qtensorformer import QTensorFormerConfig from modeling_qtensorformer import QTensorFormerForCausalLM try: from fastapi import FastAPI, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse import uvicorn HAS_FASTAPI = True except ImportError: HAS_FASTAPI = False app = FastAPI(title="Q-TensorFormer Inference Engine", version="4.0.0") if HAS_FASTAPI else None if app is not None: app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global model state MODEL = None CONFIG = None TOKENIZER = None DEVICE = "cpu" @app.get("/v1/models") if app else lambda: None async def list_models(): return { "object": "list", "data": [ { "id": "qtensorformer", "object": "model", "created": int(time.time()), "owned_by": "Premchan369", "permission": [], "root": "qtensorformer", "parent": None, } ], } @app.get("/v1/telemetry") if app else lambda: None async def get_telemetry(): if MODEL is None: raise HTTPException(status_code=503, detail="Model not initialized") pid_ctrl = MODEL.model.pid_controller return { "active_tt_rank": int(pid_ctrl.rank_levels[int(pid_ctrl.active_rank_idx.item())]), "pid_error": float(pid_ctrl.prev_err.item()), "integral_accumulated": float(pid_ctrl.integral_err.item()), "target_budget": float(pid_ctrl.target_budget), "device": str(DEVICE), "vocab_size": int(CONFIG.vocab_size), "num_layers": int(CONFIG.num_hidden_layers), } async def generate_stream_tokens( prompt_text: str, max_tokens: int = 128, temperature: float = 0.7, ) -> AsyncGenerator[str, None]: req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created_ts = int(time.time()) # Tokenize input if TOKENIZER: input_ids = TOKENIZER(prompt_text, return_tensors="pt").input_ids.to(DEVICE) else: # Fallback character tokenization for standalone test mode tokens = [ord(c) % CONFIG.vocab_size for c in prompt_text] input_ids = torch.tensor([tokens], dtype=torch.long, device=DEVICE) curr_ids = input_ids past_key_values = None for _ in range(max_tokens): with torch.no_grad(): outputs = MODEL(curr_ids[:, -1:], past_key_values=past_key_values, use_cache=True) logits = outputs.logits[:, -1, :] past_key_values = outputs.past_key_values if temperature > 0: probs = torch.softmax(logits / max(1e-4, temperature), dim=-1) next_tok = torch.multinomial(probs, num_samples=1) else: next_tok = logits.argmax(dim=-1, keepdim=True) curr_ids = torch.cat([curr_ids, next_tok], dim=1) tok_id = int(next_tok[0, 0].item()) if TOKENIZER: tok_str = TOKENIZER.decode([tok_id]) else: tok_str = chr(tok_id) if 32 <= tok_id < 127 else " " chunk = { "id": req_id, "object": "chat.completion.chunk", "created": created_ts, "model": "qtensorformer", "choices": [ { "index": 0, "delta": {"content": tok_str}, "finish_reason": None, } ], } yield f"data: {json.dumps(chunk)}\n\n" await asyncio.sleep(0.001) # Final stopping chunk done_chunk = { "id": req_id, "object": "chat.completion.chunk", "created": created_ts, "model": "qtensorformer", "choices": [ { "index": 0, "delta": {}, "finish_reason": "stop", } ], } yield f"data: {json.dumps(done_chunk)}\n\n" yield "data: [DONE]\n\n" @app.post("/v1/chat/completions") if app else lambda: None async def chat_completions(request: Request): if MODEL is None: raise HTTPException(status_code=503, detail="Model is still loading...") body = await request.json() messages = body.get("messages", []) max_tokens = body.get("max_tokens", 128) temperature = body.get("temperature", 0.7) stream = body.get("stream", False) # Format user prompt from messages prompt = "" for m in messages: role = m.get("role", "user") content = m.get("content", "") prompt += f"{role.capitalize()}: {content}\n" prompt += "Assistant: " if stream: return StreamingResponse( generate_stream_tokens(prompt, max_tokens=max_tokens, temperature=temperature), media_type="text/event-stream", ) # Non-streaming full response if TOKENIZER: input_ids = TOKENIZER(prompt, return_tensors="pt").input_ids.to(DEVICE) else: tokens = [ord(c) % CONFIG.vocab_size for c in prompt] input_ids = torch.tensor([tokens], dtype=torch.long, device=DEVICE) t0 = time.perf_counter() with torch.no_grad(): out_ids = MODEL.generate( input_ids, max_new_tokens=max_tokens, do_sample=(temperature > 0), temperature=max(1e-4, temperature), ) elapsed = time.perf_counter() - t0 new_tokens = out_ids[0, input_ids.shape[1] :] if TOKENIZER: reply_text = TOKENIZER.decode(new_tokens, skip_special_tokens=True) else: reply_text = "".join(chr(t) if 32 <= t < 127 else " " for t in new_tokens.tolist()) pid_ctrl = MODEL.model.pid_controller active_rank = int(pid_ctrl.rank_levels[int(pid_ctrl.active_rank_idx.item())]) return { "id": f"chatcmpl-{uuid.uuid4().hex[:12]}", "object": "chat.completion", "created": int(time.time()), "model": "qtensorformer", "choices": [ { "index": 0, "message": {"role": "assistant", "content": reply_text}, "finish_reason": "stop", } ], "usage": { "prompt_tokens": int(input_ids.shape[1]), "completion_tokens": int(len(new_tokens)), "total_tokens": int(out_ids.shape[1]), }, "telemetry": { "latency_seconds": round(elapsed, 4), "tokens_per_second": round(len(new_tokens) / max(1e-6, elapsed), 2), "active_tt_rank": active_rank, "target_budget": float(pid_ctrl.target_budget), }, } def start_server(model_path: str = ".", host: str = "0.0.0.0", port: int = 8000, device_str: str = "auto"): global MODEL, CONFIG, TOKENIZER, DEVICE if not HAS_FASTAPI: print("[Error] fastapi and uvicorn are required to run serve.py: pip install fastapi uvicorn") sys.exit(1) DEVICE = torch.device("cuda" if torch.cuda.is_available() and device_str != "cpu" else "cpu") print(f"\n[Server] Loading Q-TensorFormer from {model_path} onto {DEVICE}...") CONFIG = QTensorFormerConfig.from_pretrained(model_path) MODEL = QTensorFormerForCausalLM(CONFIG).to(DEVICE) MODEL.eval() try: TOKENIZER = AutoTokenizer.from_pretrained(model_path) except Exception: print("[Server] Note: No tokenizer found on path, using character fallback.") print(f"[Server] Model loaded successfully ({sum(p.numel() for p in MODEL.parameters()):,} parameters)") print(f"[Server] Starting OpenAI API server on http://{host}:{port}/v1/chat/completions\n") uvicorn.run(app, host=host, port=port) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Q-TensorFormer OpenAI-Compatible Server") parser.add_argument("--model-path", type=str, default=".", help="Model checkpoint path") parser.add_argument("--host", type=str, default="0.0.0.0", help="Binding host") parser.add_argument("--port", type=int, default=8000, help="Binding port") parser.add_argument("--device", type=str, default="auto", help="Compute device ('cuda', 'cpu')") args = parser.parse_args() start_server(model_path=args.model_path, host=args.host, port=args.port, device_str=args.device)