Upload app.py

app.py

@@ -6,8 +6,10 @@ Provides streaming chat completion API.
 import os
 import json
 import asyncio
+import time
 from typing import AsyncGenerator, List, Dict, Any, Optional
 from contextlib import asynccontextmanager
+from datetime import datetime
 
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
@@ -17,6 +19,10 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import torch
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
 # Model configuration
 MODEL_NAME = "Nanbeige/Nanbeige4.1-3B"
@@ -36,8 +42,8 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     messages: List[Message]
     stream: bool = True
-    max_tokens: int = 2048
-    temperature: float = 0.6
+    max_tokens: int = 8192  # Increased from 2048 (supports up to 131072)
+    temperature: float = 0.6  # Nanbeige4.1-3B recommended
     tools: Optional[List[Dict]] = None
 
 
@@ -83,59 +89,121 @@ app.add_middleware(
 )
 
 
-def format_messages(messages: List[Message]) -> str:
-    """Format messages into prompt string."""
-    formatted = []
-    for msg in messages:
-        if msg.role == "system":
-            formatted.append(f"System: {msg.content}")
-        elif msg.role == "user":
-            formatted.append(f"User: {msg.content}")
-        elif msg.role == "assistant":
-            formatted.append(f"Assistant: {msg.content}")
+def format_messages_proper(messages: List[Message], tools: Optional[List[Dict]] = None) -> str:
+    """Format messages using the model's proper chat template.
+    
+    Nanbeige4.1-3B uses the HF transformers chat template.
+    This ensures proper formatting for both regular and tool-aware conversations.
+    """
+    global tokenizer
     
-    formatted.append("Assistant:")
-    return "\n\n".join(formatted)
+    # Convert Message objects to dicts for tokenizer
+    message_dicts = [{"role": msg.role, "content": msg.content} for msg in messages]
+    
+    # Use tokenizer's built-in chat template for proper formatting
+    if tools:
+        # Tool-aware formatting (for function calling)
+        prompt = tokenizer.apply_chat_template(
+            message_dicts,
+            tools=tools,
+            add_generation_prompt=True,
+            tokenize=False
+        )
+    else:
+        # Regular chat formatting
+        prompt = tokenizer.apply_chat_template(
+            message_dicts,
+            add_generation_prompt=True,
+            tokenize=False
+        )
+    
+    return prompt
 
 
-async def stream_tokens(prompt: str, max_tokens: int, temperature: float) -> AsyncGenerator[str, None]:
-    """Stream tokens from the model."""
+async def stream_tokens(prompt: str, max_tokens: int, temperature: float, tools: Optional[List[Dict]] = None) -> AsyncGenerator[str, None]:
+    """Stream tokens from the model token-by-token as fast as generated.
+    
+    Uses Nanbeige4.1-3B recommended hyperparameters.
+    """
     global model, tokenizer
     
-    inputs = tokenizer(prompt, return_tensors="pt")
+    start_time = time.time()
+    logger.info(f"Starting token generation for prompt length: {len(prompt)}")
+    
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=2048
+    )
+    
     if torch.cuda.is_available():
         inputs = inputs.to("cuda")
     
+    # Create streamer with timeout to prevent hanging
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_prompt=True,
-        skip_special_tokens=True
+        skip_special_tokens=True,
+        timeout=300.0  # 5 min timeout per token
     )
     
     generation_kwargs = dict(
-        inputs,
+        **inputs,
         streamer=streamer,
-        max_new_tokens=max_tokens,
+        max_new_tokens=min(max_tokens, 131072),  # Support up to model's max (131072)
         temperature=temperature,
+        top_p=0.95,  # Nanbeige4.1-3B recommended
+        repetition_penalty=1.0,  # Nanbeige4.1-3B recommended
         do_sample=temperature > 0,
+        eos_token_id=166101,  # Nanbeige4.1-3B specific EOS token
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # Run generation in separate thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    # Run generation in separate thread (non-blocking)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs, daemon=False)
     thread.start()
     
     generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
-        # Yield each token
-        data = json.dumps({"type": "token", "content": new_text})
-        yield f"data: {data}\n\n"
+    token_count = 0
+    first_token_time = None
+    
+    try:
+        for new_text in streamer:
+            if new_text:  # Skip empty strings
+                generated_text += new_text
+                token_count += 1
+                
+                # Log first token time (time to first byte)
+                if first_token_time is None:
+                    first_token_time = time.time() - start_time
+                    logger.info(f"First token generated in {first_token_time:.2f}s")
+                
+                # preview logging to verify streaming works
+                logger.info(f"streaming token #{token_count}: {repr(new_text)}")
+
+                # Yield SSE event immediately (no buffering)
+                data = json.dumps({"type": "token", "content": new_text})
+                yield f"data: {data}\n\n"
+                logger.debug(f"Token {token_count}: {repr(new_text[:20])}...")
+        
+        # Log generation stats
+        total_time = time.time() - start_time
+        tokens_per_sec = token_count / total_time if total_time > 0 else 0
+        logger.info(f"Generation complete: {token_count} tokens in {total_time:.2f}s ({tokens_per_sec:.2f} tok/s)")
+        
+        # Signal completion
+        yield f"data: {json.dumps({'type': 'done', 'content': ''})}\n\n"
     
-    # Signal completion
-    yield f"data: {json.dumps({'type': 'done', 'content': ''})}\n\n"
+    except Exception as e:
+        logger.error(f"Token generation error: {e}", exc_info=True)
+        yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
     
-    thread.join()
+    finally:
+        # Wait for thread to finish
+        thread.join(timeout=5)
+        if thread.is_alive():
+            logger.warning("Generation thread did not finish within timeout")
 
 
 @app.get("/")
@@ -172,18 +240,21 @@ async def chat(request: ChatRequest):
     if model is None or tokenizer is None:
         raise HTTPException(status_code=503, detail="Model not loaded yet")
     
-    # Format messages into prompt
-    prompt = format_messages(request.messages)
+    # Format messages using the model's proper chat template
+    prompt = format_messages_proper(request.messages, request.tools)
     
     if request.stream:
-        # Return streaming response
+        # Return streaming response with anti-buffering headers
         return StreamingResponse(
-            stream_tokens(prompt, request.max_tokens, request.temperature),
+            stream_tokens(prompt, request.max_tokens, request.temperature, request.tools),
             media_type="text/event-stream",
             headers={
-                "Cache-Control": "no-cache",
+                "Cache-Control": "no-cache, no-store, must-revalidate",
+                "Pragma": "no-cache",
+                "Expires": "0",
                 "Connection": "keep-alive",
-                "X-Accel-Buffering": "no"
+                "X-Accel-Buffering": "no",
+                "Transfer-Encoding": "chunked"
             }
         )
     else:
@@ -194,9 +265,12 @@ async def chat(request: ChatRequest):
         
         outputs = model.generate(
             **inputs,
-            max_new_tokens=request.max_tokens,
+            max_new_tokens=min(request.max_tokens, 131072),  # Support up to model's max
             temperature=request.temperature,
+            top_p=0.95,  # Nanbeige4.1-3B recommended
+            repetition_penalty=1.0,  # Nanbeige4.1-3B recommended
             do_sample=request.temperature > 0,
+            eos_token_id=166101,  # Model-specific EOS token
             pad_token_id=tokenizer.eos_token_id
         )
         
@@ -215,4 +289,4 @@ async def chat(request: ChatRequest):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)