Update app.py

app.py

@@ -19,11 +19,8 @@ from pydantic import BaseModel, Field, ValidationError
 from transformers import AutoTokenizer
 
 # ---------- Configuration ----------
-# 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")
@@ -34,6 +31,12 @@ API_KEY = os.getenv("API_KEY", None)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("uvicorn.error")
 
+# Bonsai architecture constants (from config.json)
+NUM_LAYERS = 28
+NUM_KV_HEADS = 8
+HEAD_DIM = 128
+DTYPE = np.float32
+
 # ---------- Pydantic Models ----------
 class Message(BaseModel):
     role: str = Field(..., pattern="^(system|user|assistant)$")
@@ -151,7 +154,6 @@ 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,
@@ -161,7 +163,7 @@ def _build_chat_prompt(messages: List[Message]) -> str:
         return prompt
     except Exception as e:
         logger.error(f"Chat template error: {e}")
-        # Fallback to a simple concatenation if template fails
+        # Fallback to a simple concatenation
         prompt = ""
         for msg in messages:
             prompt += f"<|{msg.role}|>\n{msg.content}\n"
@@ -197,6 +199,26 @@ def _sample_token(logits: np.ndarray, temperature: float, top_p: float) -> int:
     probs = _softmax(logits)
     return int(np.random.choice(len(probs), p=probs))
 
+def _init_past_key_values(batch_size: int = 1) -> Dict[str, np.ndarray]:
+    """Create empty past_key_values tensors for the first inference step."""
+    past_kv = {}
+    empty_shape = (batch_size, NUM_KV_HEADS, 0, HEAD_DIM)
+    empty_tensor = np.zeros(empty_shape, dtype=DTYPE)
+    for i in range(NUM_LAYERS):
+        past_kv[f"past_key_values.{i}.key"] = empty_tensor.copy()
+        past_kv[f"past_key_values.{i}.value"] = empty_tensor.copy()
+    return past_kv
+
+def _update_past_key_values(outputs: List[np.ndarray], output_names: List[str]) -> Dict[str, np.ndarray]:
+    """Extract present_key_values from ONNX outputs and return as dictionary."""
+    new_past = {}
+    for name, value in zip(output_names, outputs):
+        if name.startswith("present"):
+            # Convert "present_key_values.0.key" -> "past_key_values.0.key"
+            past_name = name.replace("present", "past")
+            new_past[past_name] = value
+    return new_past
+
 def _generate_full(
     prompt: str,
     max_new_tokens: int,
@@ -207,40 +229,61 @@ def _generate_full(
     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),
-    }
-    
+    input_ids = tokenizer.encode(prompt, return_tensors="np").astype(np.int64)
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    position_ids = np.arange(input_ids.shape[1], dtype=np.int64).reshape(1, -1)
+
+    # Initialize KV cache
+    past_kv = _init_past_key_values(batch_size=1)
+
     generated_tokens = []
     stop_sequences = stop_sequences or []
     eos_token_id = tokenizer.eos_token_id
-    
-    for _ in range(max_new_tokens):
+
+    # Prefill step: process full prompt
+    ort_inputs = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "position_ids": position_ids,
+        "num_logits_to_keep": np.array([1], dtype=np.int64),
+        **past_kv,
+    }
+    outputs = ort_session.run(None, ort_inputs)
+    # First output is logits, the rest are present_key_values
+    logits = outputs[0][:, -1, :]
+    next_token = _sample_token(logits[0], temperature, top_p)
+    generated_tokens.append(next_token)
+
+    # Update past_key_values from outputs
+    past_kv = _update_past_key_values(outputs, [out.name for out in ort_session.get_outputs()])
+
+    for step in range(1, max_new_tokens):
+        # Subsequent steps: only the last token
+        last_token = np.array([[next_token]], dtype=np.int64)
+        attention_mask = np.ones((1, past_kv[f"past_key_values.0.key"].shape[2] + 1), dtype=np.int64)
+        position_ids = np.array([[past_kv[f"past_key_values.0.key"].shape[2]]], dtype=np.int64)
+
+        ort_inputs = {
+            "input_ids": last_token,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "num_logits_to_keep": np.array([1], dtype=np.int64),
+            **past_kv,
+        }
         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
-        )
-        
-        # Check stop conditions
+
+        past_kv = _update_past_key_values(outputs, [out.name for out in ort_session.get_outputs()])
+
         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()
 
@@ -255,31 +298,55 @@ async def _generate_stream(
         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),
-    }
-    
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    position_ids = np.arange(input_ids.shape[1], dtype=np.int64).reshape(1, -1)
+
+    past_kv = _init_past_key_values(batch_size=1)
     generated_tokens = []
     stop_sequences = stop_sequences or []
     eos_token_id = tokenizer.eos_token_id
-    
-    for _ in range(max_new_tokens):
+
+    # Prefill
+    ort_inputs = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "position_ids": position_ids,
+        "num_logits_to_keep": np.array([1], dtype=np.int64),
+        **past_kv,
+    }
+    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)
+    past_kv = _update_past_key_values(outputs, [out.name for out in ort_session.get_outputs()])
+
+    new_text = tokenizer.decode([next_token], skip_special_tokens=True)
+    if new_text:
+        yield new_text
+
+    for step in range(1, max_new_tokens):
+        last_token = np.array([[next_token]], dtype=np.int64)
+        attention_mask = np.ones((1, past_kv[f"past_key_values.0.key"].shape[2] + 1), dtype=np.int64)
+        position_ids = np.array([[past_kv[f"past_key_values.0.key"].shape[2]]], dtype=np.int64)
+
+        ort_inputs = {
+            "input_ids": last_token,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "num_logits_to_keep": np.array([1], dtype=np.int64),
+            **past_kv,
+        }
         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
-        )
-        
+
+        past_kv = _update_past_key_values(outputs, [out.name for out in ort_session.get_outputs()])
+
         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:
@@ -375,14 +442,14 @@ def model_info():
 @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
 async def chat_completions(req: ChatCompletionRequest):
     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_ID, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n"
@@ -392,7 +459,7 @@ async def chat_completions(req: ChatCompletionRequest):
             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:
         text = await asyncio.to_thread(
             _generate_full,