serve.py

#!/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)