time check

main_gguf.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 from llama_cpp import Llama, llama_cpp
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import json
@@ -95,6 +95,11 @@ def generate_base_model_stream(messages: List[dict], max_length: int = 256, temp
     """Generate streaming response from base model."""
     print(f"Using: base model", flush=True)
     print(f"Messages: {messages}", flush=True)
+
+    start = time.time()
+    first_token_time = None
+    token_count = 0
+
     for chunk in base_model.create_chat_completion(
         messages=messages,
         max_tokens=max_length,
@@ -103,8 +108,17 @@ def generate_base_model_stream(messages: List[dict], max_length: int = 256, temp
     ):
         delta = chunk["choices"][0]["delta"]
         if "content" in delta:
+            if first_token_time is None:
+                first_token_time = time.time()
+                print(f"\n[Prefill: {first_token_time - start:.2f}s]", flush=True)
+            token_count += 1
+            print(delta["content"], end="", flush=True)
             yield delta["content"]
 
+    total_time = time.time() - start
+    gen_time = total_time - (first_token_time - start) if first_token_time else 0
+    print(f"\n[Total: {total_time:.2f}s | Tokens: {token_count} | Speed: {token_count/gen_time:.1f} tok/s]", flush=True)
+
 
 
 def generate_text_stream(messages: List[dict], language: str, max_length: int = 256, temperature: float = 0.7):
@@ -193,20 +207,14 @@ async def generate_stream_api(request: GenerateRequest):
 
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest, raw_request: Request):
-    # Log raw body to debug message loss
-    raw_body = await raw_request.body()
-
+async def chat_completions(request: ChatCompletionRequest):
     messages_dicts = [{"role": msg.role, "content": msg.content} for msg in request.messages]
     language = request.model if request.model in GGUF_MODELS else "English"
-
-    print(f"Received messages_dicts: {messages_dicts}")
     
     chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
     created = int(time.time())
 
     def event_generator():
-        print(f"[EVENT_GENERATOR] Starting with messages: {messages_dicts}", flush=True)
         try:
             for token in generate_base_model_stream(
                 messages_dicts,