Update app.py

app.py

@@ -20,16 +20,18 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.gzip import GZipMiddleware
 from huggingface_hub import hf_hub_download
-from loguru import logger
 from snowflake import SnowflakeGenerator
 
 if os.environ.get("MODELSCOPE_ENVIRONMENT") == "studio":
     from modelscope import patch_hub
     patch_hub()
 
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:64"
 os.environ["RWKV_V7_ON"] = "1"
 os.environ["RWKV_JIT_ON"] = "1"
+os.environ["RWKV_CUDA_ON"] = "1"
+
+GPU_LOCK = asyncio.Lock()
 
 class ChatMessage(BaseModel):
     role: str = Field()
@@ -44,10 +46,6 @@ class LogprobsContent(BaseModel):
     content: Optional[List[Logprob]] = None
     refusal: Optional[List[Logprob]] = None
 
-class FunctionCall(BaseModel):
-    name: str
-    arguments: str
-
 class ChatCompletionMessage(BaseModel):
     role: Optional[str] = Field(None)
     content: Optional[str] = Field(None)
@@ -57,11 +55,6 @@ class ChatCompletionMessage(BaseModel):
 class PromptTokensDetails(BaseModel):
     cached_tokens: int
 
-class CompletionTokensDetails(BaseModel):
-    reasoning_tokens: int
-    accepted_prediction_tokens: int
-    rejected_prediction_tokens: int
-
 class Usage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
@@ -75,14 +68,6 @@ class ChatCompletionChoice(BaseModel):
     logprobs: Optional[LogprobsContent] = None
     finish_reason: Optional[str] = Field(...)
 
-class ChatCompletion(BaseModel):
-    id: str = Field(...)
-    object: Literal["chat.completion"] = "chat.completion"
-    created: int = Field(...)
-    model: str
-    choices: List[ChatCompletionChoice]
-    usage: Usage
-
 class ChatCompletionChunk(BaseModel):
     id: str = Field(...)
     object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
@@ -113,20 +98,6 @@ def remove_nested_think_tags_stack(text):
             i += 1
     return result
 
-def parse_think_response(full_response: str):
-    think_start = full_response.find("<think")
-    if think_start == -1:
-        return None, full_response.strip()
-    think_end = full_response.find("</think>")
-    if think_end == -1:
-        reasoning = full_response[think_start:].strip()
-        content = ""
-    else:
-        reasoning = full_response[think_start : think_end + 9].strip()
-        content = full_response[think_end + 9 :].strip()
-    reasoning_content = reasoning.replace("<think", "").replace("</think>", "").strip()
-    return reasoning_content, content
-
 def cleanMessages(messages: List[ChatMessage], removeThinkingContent: bool = False):
     promptStrList = []
     for message in messages:
@@ -138,35 +109,6 @@ def cleanMessages(messages: List[ChatMessage], removeThinkingContent: bool = Fal
         promptStrList.append(f"{role_str}: {content}")
     return "\n\n".join(promptStrList)
 
-def format_bytes(size):
-    power = 2**10
-    n = 0
-    power_labels = {0: "", 1: "K", 2: "M", 3: "G", 4: "T"}
-    while size > power:
-        size /= power
-        n += 1
-    return f"{size:.4f}{power_labels[n]+'B'}"
-
-LOGGER_QUEUE = queue.Queue(5)
-
-def logger_worker():
-    while True:
-        item = LOGGER_QUEUE.get()
-        try:
-            requests.post(
-                os.environ.get("LOG_PORT"),
-                headers={"Content-Type": "application/json"},
-                json=item,
-            )
-        except Exception:
-            pass
-
-if os.environ.get("LOG_PORT"):
-    threading.Thread(target=logger_worker).start()
-
-def log(item):
-    LOGGER_QUEUE.put_nowait(item)
-
 class SamplerConfig(BaseModel):
     max_tokens: int = 4096
     temperature: float = 1.0
@@ -187,6 +129,7 @@ class ModelConfig(BaseModel):
     DEFAULT_REASONING: bool = False
     REASONING: bool = False
     VOCAB: str = "rwkv_vocab_v20230424"
+    CTX_LEN: int = 4096
     DEFAULT_SAMPLER: SamplerConfig = Field(default_factory=SamplerConfig)
 
 class Config(BaseSettings):
@@ -200,19 +143,22 @@ class Config(BaseSettings):
             SERVICE_NAME="rwkv7-g1a4-2.9b-20251118-ctx8192",
             DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a4-2.9b-20251118-ctx8192.pth",
             DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1",
-            REASONING=True
+            REASONING=True,
+            CTX_LEN=8192
         ),
         ModelConfig(
             SERVICE_NAME="rwkv7-g1a3-1.5b-20251015-ctx8192",
             DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a3-1.5b-20251015-ctx8192.pth",
             DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1",
-            REASONING=True
+            REASONING=True,
+            CTX_LEN=8192
         ),
         ModelConfig(
             SERVICE_NAME="rwkv7-g1a-0.4b-20250905-ctx4096",
             DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a-0.4b-20250905-ctx4096.pth",
             DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1",
-            REASONING=True
+            REASONING=True,
+            CTX_LEN=4096
         ),
         ModelConfig(
             SERVICE_NAME="rwkv7-g1a-0.1b-20250728-ctx4096",
@@ -220,7 +166,8 @@ class Config(BaseSettings):
             DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1",
             REASONING=True,
             DEFAULT_CHAT=True,
-            DEFAULT_REASONING=True
+            DEFAULT_REASONING=True,
+            CTX_LEN=4096
         ),
     ]
 
@@ -248,7 +195,6 @@ if "cuda" in CONFIG.STRATEGY.lower() and not torch.cuda.is_available():
 if CONFIG.RWKV_CUDA_ON and "cuda" in CONFIG.STRATEGY.lower():
     from pynvml import *
     nvmlInit()
-    gpu_h = nvmlDeviceGetHandleByIndex(0)
     os.environ["RWKV_CUDA_ON"] = "1"
     torch.backends.cudnn.benchmark = True
     torch.backends.cudnn.allow_tf32 = True
@@ -365,7 +311,7 @@ def needs_verification(msg: str, model: str) -> bool:
     triggers = ["es verdad", "dato", "precio", "cuando", "quien", "noticia", "actualidad", "verify"]
     return any(t in msg.lower() for t in triggers)
 
-app = FastAPI(title="RWKV Zero-Bias Server")
+app = FastAPI(title="RWKV Ultimate Server")
 
 app.add_middleware(
     CORSMiddleware,
@@ -382,6 +328,28 @@ async def privacy_middleware(request: Request, call_next):
         request.scope["client"] = (fake.ipv4(), request.client.port if request.client else 80)
     return await call_next(request)
 
+def prune_context(messages: List[ChatMessage], model_name: str, max_gen_tokens: int):
+    storage = MODEL_STORAGE[model_name]
+    limit = storage.MODEL_CONFIG.CTX_LEN
+    pipeline = storage.pipeline
+    
+    current_text = cleanMessages(messages)
+    tokens = pipeline.encode(current_text)
+    
+    if len(tokens) + max_gen_tokens < limit:
+        return messages
+        
+    system_msgs = [m for m in messages if m.role == "System"]
+    other_msgs = [m for m in messages if m.role != "System"]
+    
+    while len(other_msgs) > 1:
+        candidate_text = cleanMessages(system_msgs + other_msgs)
+        if len(pipeline.encode(candidate_text)) + max_gen_tokens < limit:
+            break
+        other_msgs.pop(0)
+    
+    return system_msgs + other_msgs
+
 async def runPrefill(request: ChatCompletionRequest, ctx: str, model_tokens: List[int], model_state):
     ctx = ctx.replace("\r\n", "\n")
     tokens = MODEL_STORAGE[request.model].pipeline.encode(ctx)
@@ -404,36 +372,63 @@ def generate(request: ChatCompletionRequest, out, model_tokens: List[int], model
     out_tokens = []
     out_last = 0
     cache_word_list = []
+    
+    stop_sequences = request.stop if request.stop else []
+
     for i in range(max_tokens):
         for n in occurrence: out[n] -= args.alpha_presence + occurrence[n] * args.alpha_frequency
         token = MODEL_STORAGE[request.model].pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p)
+        
         if token == 0:
              yield {"content": "".join(cache_word_list), "finish_reason": "stop", "state": model_state}
              del out; gc.collect(); return
+
         out, model_state = MODEL_STORAGE[request.model].model.forward([token], model_state)
         model_tokens.append(token)
         out_tokens.append(token)
+        
         for xxx in occurrence: occurrence[xxx] *= request.penalty_decay
         occurrence[token] = 1 + (occurrence.get(token, 0))
+        
         tmp = MODEL_STORAGE[request.model].pipeline.decode(out_tokens[out_last:])
         if "\ufffd" in tmp: continue
         cache_word_list.append(tmp)
         out_last = i + 1
+        
+        current_buffer = "".join(cache_word_list)
+        for s in stop_sequences:
+            if s in current_buffer:
+                final_content = current_buffer.split(s)[0]
+                yield {"content": final_content, "finish_reason": "stop", "state": model_state}
+                del out; gc.collect(); return
+
         if len(cache_word_list) > 1:
             yield {"content": cache_word_list.pop(0), "finish_reason": None}
+            
     yield {"content": "".join(cache_word_list), "finish_reason": "length"}
 
 async def chatResponseStream(request: ChatCompletionRequest, model_state: any, completionId: str, enableReasoning: bool):
     clean_msg = cleanMessages(request.messages, enableReasoning)
     prompt = f"{clean_msg}\n\nAssistant:{' <think' if enableReasoning else ''}"
-    out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state)
-    yield f"data: {ChatCompletionChunk(id=completionId, created=int(time.time()), model=request.model, choices=[ChatCompletionChoice(index=0, delta=ChatCompletionMessage(role='Assistant', content=''), finish_reason=None)]).model_dump_json()}\n\n"
-    for chunk in generate(request, out, model_tokens, model_state, max_tokens=request.max_tokens or 4096):
-        content = chunk["content"]
-        if content:
-             yield f"data: {ChatCompletionChunk(id=completionId, created=int(time.time()), model=request.model, choices=[ChatCompletionChoice(index=0, delta=ChatCompletionMessage(content=content), finish_reason=None)]).model_dump_json()}\n\n"
-        if chunk.get("finish_reason"): break
-        await asyncio.sleep(0)
+    
+    async with GPU_LOCK:
+        try:
+            out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state)
+            
+            yield f"data: {ChatCompletionChunk(id=completionId, created=int(time.time()), model=request.model, choices=[ChatCompletionChoice(index=0, delta=ChatCompletionMessage(role='Assistant', content=''), finish_reason=None)]).model_dump_json()}\n\n"
+            
+            for chunk in generate(request, out, model_tokens, model_state, max_tokens=request.max_tokens or 4096):
+                content = chunk.get("content", "")
+                finish = chunk.get("finish_reason", None)
+                if content:
+                    yield f"data: {ChatCompletionChunk(id=completionId, created=int(time.time()), model=request.model, choices=[ChatCompletionChoice(index=0, delta=ChatCompletionMessage(content=content), finish_reason=None)]).model_dump_json()}\n\n"
+                if finish:
+                    yield f"data: {ChatCompletionChunk(id=completionId, created=int(time.time()), model=request.model, choices=[ChatCompletionChoice(index=0, delta=ChatCompletionMessage(content=''), finish_reason=finish)]).model_dump_json()}\n\n"
+                    break
+                await asyncio.sleep(0)
+        finally:
+            pass
+            
     yield "data: [DONE]\n\n"
 
 @app.post("/api/v1/chat/completions")
@@ -446,26 +441,34 @@ async def chat_completions(request: ChatCompletionRequest):
     if "rwkv-latest" in model_key:
         if is_reasoning and DEFAULT_REASONING_MODEL_NAME: target_model = DEFAULT_REASONING_MODEL_NAME
         elif DEFALUT_MODEL_NAME: target_model = DEFALUT_MODEL_NAME
+    
     if target_model not in MODEL_STORAGE:
         raise HTTPException(404, f"Model {target_model} not loaded.")
     request.model = target_model
+
     default_sampler = MODEL_STORAGE[target_model].MODEL_CONFIG.DEFAULT_SAMPLER
     req_data = request.model_dump()
     for k, v in default_sampler.model_dump().items():
         if req_data.get(k) is None: req_data[k] = v
     realRequest = ChatCompletionRequest(**req_data)
+
     sys_msg = ChatMessage(role="System", content=TruthProtocol.STRICT_SYSTEM_PROMPT)
     if realRequest.messages:
         if realRequest.messages[0].role == "System":
              realRequest.messages[0].content = f"{TruthProtocol.STRICT_SYSTEM_PROMPT}\n\n{realRequest.messages[0].content}"
         else:
             realRequest.messages.insert(0, sys_msg)
+    
     last_msg = realRequest.messages[-1]
     if last_msg.role == "user" and needs_verification(last_msg.content, raw_model):
         ctx = search_facts(last_msg.content)
         if ctx:
             realRequest.messages.insert(-1, ChatMessage(role="System", content=ctx))
+    
     TruthProtocol.enforce_truth_params(realRequest)
+    
+    realRequest.messages = prune_context(realRequest.messages, target_model, realRequest.max_tokens or 1024)
+
     return StreamingResponse(chatResponseStream(realRequest, None, completionId, is_reasoning), media_type="text/event-stream")
 
 @app.get("/api/v1/models")