Update app.py

app.py

@@ -4,32 +4,31 @@ import asyncio
 import json
 import logging
 import os
-import threading
 import time
 import uuid
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
+import numpy as np
+import onnxruntime as ort
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
+from huggingface_hub import snapshot_download
 from pydantic import BaseModel, Field, ValidationError
+from transformers import AutoTokenizer
 
 # ---------- Configuration ----------
-MODEL_REPO = os.getenv("MODEL_REPO", "prism-ml/Bonsai-8B-gguf")
-MODEL_FILE = os.getenv("MODEL_FILE", "Bonsai-8B-Q1_0_g128.gguf")
-MODEL_REVISION = os.getenv("MODEL_REVISION", "main")
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-N_CTX = int(os.getenv("N_CTX", "4096"))
-N_THREADS = int(os.getenv("N_THREADS", "4"))
-N_GPU_LAYERS = int(os.getenv("N_GPU_LAYERS", "0"))  # 0 = CPU only
-MAX_TOKENS_DEFAULT = int(os.getenv("MAX_TOKENS_DEFAULT", "512"))
-TEMPERATURE_DEFAULT = float(os.getenv("TEMPERATURE_DEFAULT", "0.7"))
-TOP_P_DEFAULT = float(os.getenv("TOP_P_DEFAULT", "0.95"))
+# Model Selection: Use "onnx-community/Bonsai-1.7B-ONNX" or "onnx-community/Bonsai-8B-ONNX"
+MODEL_ID = os.getenv("MODEL_ID", "onnx-community/Bonsai-1.7B-ONNX")
+# Quantization: Choose from 'q1', 'q2', 'q4', 'q8' based on the files in the ONNX model repo
+MODEL_QUANTIZATION = os.getenv("MODEL_QUANTIZATION", "q1")
+# Model file name based on quantization
+ONNX_MODEL_FILE = f"model_{MODEL_QUANTIZATION}.onnx"
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/data/bonsai-onnx")
+MAX_NEW_TOKENS_DEFAULT = int(os.getenv("MAX_NEW_TOKENS_DEFAULT", "256"))
 API_KEY = os.getenv("API_KEY", None)
 
 logging.basicConfig(level=logging.INFO)
@@ -37,21 +36,17 @@ logger = logging.getLogger("uvicorn.error")
 
 # ---------- Pydantic Models ----------
 class Message(BaseModel):
-    role: str = Field(..., pattern="^(system|user|assistant|tool)$")
-    content: Optional[Union[str, List[Dict[str, Any]]]] = None
-    tool_calls: Optional[List[Dict[str, Any]]] = None
-    tool_call_id: Optional[str] = None
-    name: Optional[str] = None
+    role: str = Field(..., pattern="^(system|user|assistant)$")
+    content: str
 
 class ChatCompletionRequest(BaseModel):
     messages: List[Message]
-    model: Optional[str] = MODEL_REPO
-    max_tokens: int = Field(default=MAX_TOKENS_DEFAULT, ge=1, le=4096)
-    temperature: float = Field(default=TEMPERATURE_DEFAULT, ge=0.0, le=2.0)
-    top_p: float = Field(default=TOP_P_DEFAULT, gt=0.0, le=1.0)
+    model: Optional[str] = MODEL_ID
+    max_tokens: int = Field(default=MAX_NEW_TOKENS_DEFAULT, ge=1, le=1024)
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    top_p: float = Field(default=0.95, gt=0.0, le=1.0)
     stream: bool = False
     stop: Optional[Union[str, List[str]]] = None
-    user: Optional[str] = None
 
 class ChatCompletionResponseChoice(BaseModel):
     index: int
@@ -71,21 +66,21 @@ class ChatCompletionResponse(BaseModel):
     choices: List[ChatCompletionResponseChoice]
     usage: Usage
 
-class TokenCountRequest(BaseModel):
-    text: str = Field(..., min_length=1)
-
-class TokenCountResponse(BaseModel):
-    text: str
-    token_count: int
+class ModelInfo(BaseModel):
+    model_id: str
+    quantization: str
+    onnx_model_file: str
+    device: str
 
 class ErrorResponse(BaseModel):
     error: str
     detail: Optional[str] = None
 
 # ---------- Global State ----------
-llm: Optional[Llama] = None
-MODEL_LOCK = threading.Lock()
-model_load_error: Optional[str] = None
+tokenizer = None
+ort_session = None
+model_load_error = None
+MODEL_LOCK = asyncio.Lock()
 
 # ---------- Helper Functions ----------
 def _verify_api_key(request: Request) -> None:
@@ -95,124 +90,220 @@ def _verify_api_key(request: Request) -> None:
     if not auth or auth != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid or missing API key")
 
-def _download_model() -> str:
-    """Download the GGUF model from Hugging Face Hub."""
-    os.makedirs("/data", exist_ok=True)
+def _model_device() -> str:
+    return "cuda" if ort.get_device().lower() == "gpu" else "cpu"
+
+def _download_model_snapshot() -> str:
+    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
+    allow_patterns = [
+        "config.json",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "chat_template.jinja",
+        f"onnx/{ONNX_MODEL_FILE}",
+        f"onnx/{ONNX_MODEL_FILE}_data",
+    ]
     try:
-        local_path = hf_hub_download(
-            repo_id=MODEL_REPO,
-            filename=MODEL_FILE,
-            revision=MODEL_REVISION,
+        snapshot_download(
+            repo_id=MODEL_ID,
+            local_dir=LOCAL_MODEL_DIR,
+            local_dir_use_symlinks=False,
+            allow_patterns=allow_patterns,
             token=HF_TOKEN,
-            cache_dir="/data/.cache/huggingface",
         )
-        logger.info(f"Model downloaded to {local_path}")
-        return local_path
     except Exception as e:
         logger.error(f"Model download failed: {e}")
         raise RuntimeError(f"Failed to download model: {str(e)}")
-
-def _ensure_loaded() -> None:
-    global llm, model_load_error
-    if llm is not None:
-        return
-    if model_load_error:
-        raise HTTPException(status_code=503, detail=f"Model failed to load: {model_load_error}")
+    return LOCAL_MODEL_DIR
+
+def _create_ort_session(model_path: str) -> ort.InferenceSession:
+    so = ort.SessionOptions()
+    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    so.intra_op_num_threads = int(os.getenv("ORT_INTRA_OP_THREADS", "2"))
+    so.inter_op_num_threads = 1
+    so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    so.enable_mem_pattern = True
     try:
-        model_path = _download_model()
-        llm = Llama(
-            model_path=model_path,
-            n_ctx=N_CTX,
-            n_threads=N_THREADS,
-            n_gpu_layers=N_GPU_LAYERS,
-            verbose=False,
+        return ort.InferenceSession(model_path, sess_options=so, providers=["CPUExecutionProvider"])
+    except Exception as e:
+        logger.error(f"Failed to load ONNX session from {model_path}: {e}")
+        raise RuntimeError(f"ONNX session creation failed: {str(e)}")
+
+async def _ensure_loaded():
+    global tokenizer, ort_session, model_load_error
+    async with MODEL_LOCK:
+        if tokenizer is not None and ort_session is not None:
+            return
+        if model_load_error:
+            raise HTTPException(status_code=503, detail=f"Model failed to load: {model_load_error}")
+        try:
+            local_dir = _download_model_snapshot()
+            tokenizer = AutoTokenizer.from_pretrained(local_dir, trust_remote_code=True)
+            onnx_path = os.path.join(local_dir, "onnx", ONNX_MODEL_FILE)
+            ort_session = _create_ort_session(onnx_path)
+            logger.info(f"Model loaded successfully: {MODEL_ID} ({MODEL_QUANTIZATION})")
+        except Exception as e:
+            model_load_error = str(e)
+            logger.exception("Model loading failed")
+            raise HTTPException(status_code=503, detail=f"Model unavailable: {model_load_error}")
+
+def _build_chat_prompt(messages: List[Message]) -> str:
+    if tokenizer is None:
+        raise HTTPException(status_code=503, detail="Tokenizer not loaded")
+    try:
+        # Use the tokenizer's chat template to format the conversation
+        formatted_messages = [{"role": msg.role, "content": msg.content} for msg in messages]
+        prompt = tokenizer.apply_chat_template(
+            formatted_messages,
+            tokenize=False,
+            add_generation_prompt=True,
         )
-        logger.info("Model loaded successfully")
+        return prompt
     except Exception as e:
-        model_load_error = str(e)
-        logger.exception("Model loading failed")
-        raise HTTPException(status_code=503, detail=f"Model unavailable: {model_load_error}")
-
-def _format_chat_messages(messages: List[Message]) -> List[Dict[str, str]]:
-    """Convert Pydantic messages to a list of dicts for llama.cpp."""
-    formatted = []
-    for msg in messages:
-        content = msg.content
-        if isinstance(content, list):
-            # For simplicity, we only handle text content here.
-            # Multimodal not needed for Bonsai.
-            text_parts = [p["text"] for p in content if p.get("type") == "text"]
-            content = " ".join(text_parts) if text_parts else ""
-        formatted.append({"role": msg.role, "content": content or ""})
-    return formatted
+        logger.error(f"Chat template error: {e}")
+        # Fallback to a simple concatenation if template fails
+        prompt = ""
+        for msg in messages:
+            prompt += f"<|{msg.role}|>\n{msg.content}\n"
+        prompt += "<|assistant|>\n"
+        return prompt
+
+def _count_tokens(text: str) -> int:
+    if tokenizer is None:
+        return len(text.split())
+    return len(tokenizer.encode(text))
+
+def _softmax(x: np.ndarray) -> np.ndarray:
+    e_x = np.exp(x - np.max(x))
+    return e_x / e_x.sum(axis=-1, keepdims=True)
+
+def _top_p_sampling(logits: np.ndarray, top_p: float) -> int:
+    sorted_indices = np.argsort(logits)[::-1]
+    sorted_logits = logits[sorted_indices]
+    probs = _softmax(sorted_logits)
+    cum_probs = np.cumsum(probs)
+    cutoff_index = np.searchsorted(cum_probs, top_p) + 1
+    top_indices = sorted_indices[:cutoff_index]
+    top_probs = probs[:cutoff_index]
+    top_probs /= top_probs.sum()
+    return int(np.random.choice(top_indices, p=top_probs))
+
+def _sample_token(logits: np.ndarray, temperature: float, top_p: float) -> int:
+    if temperature <= 0:
+        return int(np.argmax(logits))
+    logits = logits / temperature
+    if top_p < 1.0:
+        return _top_p_sampling(logits, top_p)
+    probs = _softmax(logits)
+    return int(np.random.choice(len(probs), p=probs))
 
 def _generate_full(
-    messages: List[Message],
-    max_tokens: int,
+    prompt: str,
+    max_new_tokens: int,
     temperature: float,
     top_p: float,
-    stop: Optional[Union[str, List[str]]],
-) -> tuple[str, Usage]:
-    _ensure_loaded()
-    with MODEL_LOCK:
-        chat_messages = _format_chat_messages(messages)
-        result = llm.create_chat_completion(
-            messages=chat_messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stop=stop,
-            stream=False,
+    stop_sequences: Optional[List[str]] = None,
+) -> str:
+    if ort_session is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    input_ids = tokenizer.encode(prompt, return_tensors="np")
+    input_ids = input_ids.astype(np.int64)
+    
+    # Prepare initial inputs for the ONNX model
+    ort_inputs = {
+        "input_ids": input_ids,
+        "attention_mask": np.ones_like(input_ids, dtype=np.int64),
+    }
+    
+    generated_tokens = []
+    stop_sequences = stop_sequences or []
+    eos_token_id = tokenizer.eos_token_id
+    
+    for _ in range(max_new_tokens):
+        outputs = ort_session.run(None, ort_inputs)
+        logits = outputs[0][:, -1, :]
+        next_token = _sample_token(logits[0], temperature, top_p)
+        generated_tokens.append(next_token)
+        
+        # Update inputs for the next step
+        next_token_id = np.array([[next_token]], dtype=np.int64)
+        ort_inputs["input_ids"] = np.concatenate([input_ids, next_token_id], axis=1)
+        ort_inputs["attention_mask"] = np.concatenate(
+            [ort_inputs["attention_mask"], np.ones((1, 1), dtype=np.int64)], axis=1
         )
-    content = result["choices"][0]["message"]["content"]
-    usage = Usage(
-        prompt_tokens=result["usage"]["prompt_tokens"],
-        completion_tokens=result["usage"]["completion_tokens"],
-        total_tokens=result["usage"]["total_tokens"],
-    )
-    return content, usage
+        
+        # Check stop conditions
+        if next_token == eos_token_id:
+            break
+        partial_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        for stop_seq in stop_sequences:
+            if stop_seq in partial_text:
+                return partial_text.split(stop_seq)[0].strip()
+    
+    full_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+    return full_text.strip()
 
 async def _generate_stream(
-    messages: List[Message],
-    max_tokens: int,
+    prompt: str,
+    max_new_tokens: int,
     temperature: float,
     top_p: float,
-    stop: Optional[Union[str, List[str]]],
-) -> AsyncGenerator[str, None]:
-    _ensure_loaded()
-    chat_messages = _format_chat_messages(messages)
-    with MODEL_LOCK:
-        stream = llm.create_chat_completion(
-            messages=chat_messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stop=stop,
-            stream=True,
+    stop_sequences: Optional[List[str]] = None,
+):
+    if ort_session is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    input_ids = tokenizer.encode(prompt, return_tensors="np").astype(np.int64)
+    ort_inputs = {
+        "input_ids": input_ids,
+        "attention_mask": np.ones_like(input_ids, dtype=np.int64),
+    }
+    
+    generated_tokens = []
+    stop_sequences = stop_sequences or []
+    eos_token_id = tokenizer.eos_token_id
+    
+    for _ in range(max_new_tokens):
+        outputs = ort_session.run(None, ort_inputs)
+        logits = outputs[0][:, -1, :]
+        next_token = _sample_token(logits[0], temperature, top_p)
+        generated_tokens.append(next_token)
+        
+        next_token_id = np.array([[next_token]], dtype=np.int64)
+        ort_inputs["input_ids"] = np.concatenate([input_ids, next_token_id], axis=1)
+        ort_inputs["attention_mask"] = np.concatenate(
+            [ort_inputs["attention_mask"], np.ones((1, 1), dtype=np.int64)], axis=1
         )
-        for chunk in stream:
-            delta = chunk["choices"][0]["delta"]
-            if "content" in delta and delta["content"]:
-                yield delta["content"]
-            if chunk["choices"][0].get("finish_reason") == "stop":
-                break
+        
+        new_text = tokenizer.decode([next_token], skip_special_tokens=True)
+        if new_text:
+            yield new_text
+        
+        full_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        for stop_seq in stop_sequences:
+            if stop_seq in full_text:
+                return
+        if next_token == eos_token_id:
+            break
 
 # ---------- FastAPI App ----------
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     try:
-        _ensure_loaded()
+        await _ensure_loaded()
         logger.info("Model loaded successfully")
     except Exception as e:
         logger.error(f"Startup model load failed: {e}")
     yield
-    global llm
-    llm = None
+    global tokenizer, ort_session
+    tokenizer = None
+    ort_session = None
 
 app = FastAPI(
-    title="Bonsai LLM API (llama.cpp)",
+    title="Bonsai ONNX Inference API",
     version="1.0.0",
-    description="Production‑ready API for PrismML Bonsai models via GGUF.",
+    description="Fast, production-ready inference for 1-bit Bonsai LLMs using ONNX Runtime.",
     docs_url="/docs",
     redoc_url="/redoc",
     lifespan=lifespan,
@@ -229,7 +320,8 @@ app.add_middleware(
 @app.middleware("http")
 async def auth_middleware(request: Request, call_next):
     _verify_api_key(request)
-    return await call_next(request)
+    response = await call_next(request)
+    return response
 
 # ---------- Error Handlers ----------
 @app.exception_handler(HTTPException)
@@ -257,71 +349,66 @@ async def generic_exception_handler(request, exc):
 # ---------- Endpoints ----------
 @app.get("/", summary="Root")
 def root():
-    return {"message": "Bonsai API is running", "docs": "/docs"}
+    return {"message": "Bonsai ONNX API is running", "docs": "/docs"}
 
 @app.get("/health", summary="Health check")
 def health():
-    loaded = llm is not None
+    loaded = tokenizer is not None and ort_session is not None
     return {
         "status": "ok" if loaded else "degraded",
         "model_loaded": loaded,
-        "model_id": MODEL_REPO,
+        "model_id": MODEL_ID,
+        "quantization": MODEL_QUANTIZATION,
+        "device": _model_device(),
         "error": model_load_error if model_load_error else None,
     }
 
-@app.get("/v1/model", summary="Model information")
+@app.get("/v1/model", response_model=ModelInfo, summary="Model information")
 def model_info():
-    return {
-        "model_id": MODEL_REPO,
-        "model_file": MODEL_FILE,
-        "revision": MODEL_REVISION,
-        "context_length": N_CTX,
-        "gpu_layers": N_GPU_LAYERS,
-        "cpu_threads": N_THREADS,
-    }
-
-@app.post("/v1/token/count", response_model=TokenCountResponse)
-def token_count(req: TokenCountRequest):
-    _ensure_loaded()
-    tokens = llm.tokenize(req.text.encode("utf-8"))
-    return TokenCountResponse(text=req.text, token_count=len(tokens))
+    return ModelInfo(
+        model_id=MODEL_ID,
+        quantization=MODEL_QUANTIZATION,
+        onnx_model_file=ONNX_MODEL_FILE,
+        device=_model_device(),
+    )
 
-@app.post("/v1/chat/completions")
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
 async def chat_completions(req: ChatCompletionRequest):
-    _ensure_loaded()
-
+    await _ensure_loaded()
+    
+    try:
+        prompt = _build_chat_prompt(req.messages)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Prompt formatting error: {str(e)}")
+    
+    stop_seq = req.stop if isinstance(req.stop, list) else ([req.stop] if req.stop else None)
+    
     if req.stream:
         async def stream_generator():
-            yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_REPO, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n"
-            async for chunk in _generate_stream(
-                req.messages, req.max_tokens, req.temperature, req.top_p, req.stop
-            ):
-                yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_REPO, 'choices': [{'index': 0, 'delta': {'content': chunk}, 'finish_reason': None}]})}\n\n"
+            yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_ID, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n"
+            async for chunk in _generate_stream(prompt, req.max_tokens, req.temperature, req.top_p, stop_seq):
+                yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_ID, 'choices': [{'index': 0, 'delta': {'content': chunk}, 'finish_reason': None}]})}\n\n"
                 await asyncio.sleep(0)
-            yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_REPO, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
+            yield f"data: {json.dumps({'id': f'chatcmpl-{uuid.uuid4().hex[:12]}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model or MODEL_ID, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
             yield "data: [DONE]\n\n"
         return StreamingResponse(stream_generator(), media_type="text/event-stream")
-
+    
     else:
-        content, usage = await asyncio.to_thread(
+        text = await asyncio.to_thread(
             _generate_full,
-            req.messages,
-            req.max_tokens,
-            req.temperature,
-            req.top_p,
-            req.stop,
+            prompt, req.max_tokens, req.temperature, req.top_p, stop_seq
+        )
+        assistant_msg = Message(role="assistant", content=text)
+        usage = Usage(
+            prompt_tokens=_count_tokens(prompt),
+            completion_tokens=_count_tokens(text),
+            total_tokens=_count_tokens(prompt) + _count_tokens(text),
         )
         return ChatCompletionResponse(
             id=f"chatcmpl-{uuid.uuid4().hex[:12]}",
             created=int(time.time()),
-            model=req.model or MODEL_REPO,
-            choices=[
-                ChatCompletionResponseChoice(
-                    index=0,
-                    message=Message(role="assistant", content=content),
-                    finish_reason="stop",
-                )
-            ],
+            model=req.model or MODEL_ID,
+            choices=[ChatCompletionResponseChoice(index=0, message=assistant_msg, finish_reason="stop")],
             usage=usage,
         )