UPDATE: support stop tokens

Dockerfile

@@ -15,7 +15,8 @@ RUN ["cargo", "install", "wasm-pack"]
 WORKDIR /app
 ENV PATH=/root/.cargo/bin:$PATH
 RUN npm install -g pnpm
-RUN pnpm install && pnpm run build
+RUN pnpm install
+RUN pnpm run build --mode target-rwkv-hf-space
 
 FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 AS Backend
 

app.py

@@ -3,6 +3,7 @@ from config import CONFIG, ModelConfig
 import os, copy, types, gc, sys, re, time, collections, asyncio
 from huggingface_hub import hf_hub_download
 from loguru import logger
+from rich import print
 
 from snowflake import SnowflakeGenerator
 
@@ -92,6 +93,8 @@ for model_config in CONFIG.MODELS:
         else:
             DEFALUT_MODEL_NAME = model_config.SERVICE_NAME
 
+    print(model_config.DEFAULT_SAMPLER)
+
     MODEL_STORAGE[model_config.SERVICE_NAME] = ModelStorage()
     MODEL_STORAGE[model_config.SERVICE_NAME].MODEL_CONFIG = model_config
     MODEL_STORAGE[model_config.SERVICE_NAME].model = tmp_model
@@ -119,6 +122,7 @@ class ChatCompletionRequest(BaseModel):
     state_name: Optional[str] = Field(default=None)
     include_usage: Optional[bool] = Field(default=False)
     stop: Optional[list[str]] = Field(["\n\n"])
+    stop_tokens: Optional[list[int]] = Field([0])
 
     @model_validator(mode="before")
     @classmethod
@@ -169,7 +173,7 @@ async def runPrefill(
 def generate(
     request: ChatCompletionRequest,
     out,
-    model_tokens,
+    model_tokens: List[int],
     model_state,
     stops=["\n\n"],
     max_tokens=2048,
@@ -184,7 +188,7 @@ def generate(
     )  # stop generation whenever you see any token here
 
     occurrence = {}
-    out_tokens = []
+    out_tokens: List[int] = []
     out_last = 0
 
     output_cache = collections.deque(maxlen=5)
@@ -192,7 +196,7 @@ def generate(
     for i in range(max_tokens):
         for n in occurrence:
             out[n] -= args.alpha_presence + occurrence[n] * args.alpha_frequency
-        out[0] -= 1e10  # disable END_OF_TEXT
+        # out[0] -= 1e10  # disable END_OF_TEXT
 
         token = MODEL_STORAGE[request.model].pipeline.sample_logits(
             out, temperature=args.temperature, top_p=args.top_p
@@ -201,9 +205,21 @@ def generate(
         out, model_state = MODEL_STORAGE[request.model].model.forward(
             [token], model_state
         )
-        model_tokens += [token]
+        model_tokens.append(token)
+
+        out_tokens.append(token)
+
+        if token in request.stop_tokens:
+            yield {
+                "content": "",
+                "tokens": out_tokens[out_last:],
+                "finish_reason": "stop",
+                "state": model_state,
+            }
 
-        out_tokens += [token]
+            del out
+            gc.collect()
+            return
 
         for xxx in occurrence:
             occurrence[xxx] *= request.penalty_decay
@@ -260,6 +276,7 @@ async def chatResponse(
         if request.prompt == None
         else request.prompt.strip()
     )
+    logger.info(f"[REQ] {completionId} - prompt - {prompt}")
 
     out, model_tokens, model_state = await runPrefill(request, prompt, [], model_state)
 
@@ -343,6 +360,8 @@ async def chatResponseStream(
         else request.prompt.strip()
     )
 
+    # logger.info(f"[REQ] {completionId} - prompt - {prompt}")
+
     out, model_tokens, model_state = await runPrefill(request, prompt, [], model_state)
 
     prefillTime = time.time()
@@ -465,7 +484,7 @@ async def chatResponseStream(
                 streamConfig["fullTextCursor"] = len(fullText)
 
             markEnd = fullText.find(">", streamConfig["fullTextCursor"])
-            if streamConfig["isChecking"] and markEnd != -1:
+            if (streamConfig["isChecking"] and markEnd != -1) or finishReason != None:
                 streamConfig["isChecking"] = False
 
                 if (
@@ -626,6 +645,7 @@ async def chat_completions(request: ChatCompletionRequest):
 
     return r
 
+
 app.mount("/", StaticFiles(directory="dist-frontend", html=True), name="static")
 
 if __name__ == "__main__":

config.production.yaml

@@ -18,6 +18,8 @@ MODELS:
       penalty_decay: 0.996
       stop:
         - "\n\n"
+      stop_tokens:
+        - 0
   - SERVICE_NAME: "rwkv7-g1-0.1b-20250307-ctx4096"
     DOWNLOAD_MODEL_FILE_NAME: "rwkv7-g1-0.1b-20250307-ctx4096.pth"
     DOWNLOAD_MODEL_REPO_ID: "BlinkDL/rwkv7-g1"
@@ -32,3 +34,5 @@ MODELS:
       penalty_decay: 0.996
       stop:
         - "\n\n"
+      stop_tokens:
+        - 0

config.py

@@ -23,8 +23,9 @@ class SamplerConfig(BaseModel):
     top_p: float = Field(0.3, description="Top-p sampling threshold.")
     presence_penalty: float = Field(0.5, description="Presence penalty.")
     count_penalty: float = Field(0.5, description="Count penalty.")
-    penalty_decay: float = Field(0.5, description="Penalty decay factor.")
-    stop: List[str] = Field(0.996, description="List of stop sequences.")
+    penalty_decay: float = Field(0.996, description="Penalty decay factor.")
+    stop: List[str] = Field(["\n\n"], description="List of stop sequences.")
+    stop_tokens: List[int] = Field([0], description="List of stop tokens.")
 
 
 class ModelConfig(BaseModel):

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "pydantic>=2.10.6",
     "pydantic-settings>=2.8.1",
     "pynvml>=12.0.0",
+    "rich>=13.9.4",
     "rwkv==0.8.28",
     "setuptools>=75.8.2",
     "snowflake-id>=1.0.2",

uv.lock

@@ -944,6 +944,7 @@ dependencies = [
     { name = "pydantic" },
     { name = "pydantic-settings" },
     { name = "pynvml" },
+    { name = "rich" },
     { name = "rwkv" },
     { name = "setuptools" },
     { name = "snowflake-id" },
@@ -971,6 +972,7 @@ requires-dist = [
     { name = "pydantic", specifier = ">=2.10.6" },
     { name = "pydantic-settings", specifier = ">=2.8.1" },
     { name = "pynvml", specifier = ">=12.0.0" },
+    { name = "rich", specifier = ">=13.9.4" },
     { name = "rwkv", specifier = "==0.8.28" },
     { name = "setuptools", specifier = ">=75.8.2" },
     { name = "snowflake-id", specifier = ">=1.0.2" },