INIT

.gitignore

@@ -0,0 +1,16 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+.cache
+
+*pth
+*.pt
+*.st

.python-version

@@ -0,0 +1 @@
+3.10

Dockerfile

@@ -0,0 +1,27 @@
+ARG CUDA_IMAGE="12.1.1-devel-ubuntu22.04"
+FROM nvidia/cuda:${CUDA_IMAGE}
+
+RUN apt-get update && apt-get install --no-install-recommends -y \
+    build-essential \
+    git \
+    ffmpeg &&
+    apt-get clean && rm -rf /var/lib/apt/lists/*
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+COPY . .
+
+RUN uv sync --frozen
+
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+CMD ["uv", "app.py","--strategy","cuda fp16","--model_title","RWKV-x070-World-0.1B-v2.8-20241210-ctx4096","--download_repo_id","BlinkDL/rwkv-7-world"]

README.md

@@ -0,0 +1,21 @@
+# Simple RWKV OpenAI-Compatible API
+
+## Usage
+
+`RWKV-x070-World-0.1B-v2.8-20241210-ctx4096`
+
+```shell
+python app.py --strategy "cuda fp16" --model_title "RWKV-x070-World-0.1B-v2.8-20241210-ctx4096" --download_repo_id "BlinkDL/rwkv-7-world" --download_model_dir ./
+```
+
+`RWKV7-G1-0.1B-68%trained-20250303-ctx4k`
+
+```shell
+python app.py --strategy "cuda fp16" --model_title "RWKV7-G1-0.1B-68%trained-20250303-ctx4k" --download_repo_id "BlinkDL/temp-latest-training-models" --download_model_dir ./
+```
+
+`RWKV7-G1-0.1B-68%trained-20250303-ctx4k`
+
+```shell
+python app.py --strategy "cuda fp16" --model_title "RWKV7-G1-0.4B-32%trained-20250304-ctx4k" --download_repo_id "BlinkDL/temp-latest-training-models" --download_model_dir ./
+```

api_types.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union, Dict, Any, Literal
+from pydantic import BaseModel, Field
+
+
+class ChatMessage(BaseModel):
+    role: str = Field()
+    content: str = Field()
+
+
+class Logprob(BaseModel):
+    token: str
+    logprob: float
+    top_logprobs: Optional[List[Dict[str, Any]]] = None
+
+
+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, description="The role of the author of this message"
+    )
+    content: Optional[str] = Field(None, description="The contents of the message")
+    reasoning_content: Optional[str] = Field(
+        None, description="The reasoning contents of the message"
+    )
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        None, description="Tool calls generated by the model"
+    )
+
+
+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
+    total_tokens: int
+    prompt_tokens_details: Optional[PromptTokensDetails]
+    # completion_tokens_details: CompletionTokensDetails
+
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: Optional[ChatCompletionMessage] = None
+    delta: Optional[ChatCompletionMessage] = None
+    logprobs: Optional[LogprobsContent] = None
+    finish_reason: Optional[str] = Field(
+        ..., description="Reason for stopping: stop, length, content_filter, tool_calls"
+    )
+
+
+class ChatCompletion(BaseModel):
+    id: str = Field(..., description="Unique identifier for the chat completion")
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int = Field(..., description="Unix timestamp of creation")
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Usage
+
+
+class ChatCompletionChunk(BaseModel):
+    id: str = Field(..., description="Unique identifier for the chat completion")
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
+    created: int = Field(..., description="Unix timestamp of creation")
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Optional[Usage]

app.py

@@ -0,0 +1,555 @@
+import os, copy, types, gc, sys, re, time, collections, asyncio
+from huggingface_hub import hf_hub_download
+from loguru import logger
+
+from snowflake import SnowflakeGenerator
+
+CompletionIdGenerator = SnowflakeGenerator(42, timestamp=1741101491595)
+
+from pynvml import *
+
+nvmlInit()
+gpu_h = nvmlDeviceGetHandleByIndex(0)
+
+from typing import List, Optional, Union
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings
+
+
+class Config(BaseSettings, cli_parse_args=True, cli_use_class_docs_for_groups=True):
+    HOST: str = Field("127.0.0.1", description="Host")
+    PORT: int = Field(8000, description="Port")
+    DEBUG: bool = Field(False, description="Debug mode")
+    STRATEGY: str = Field("cpu", description="Stratergy")
+    MODEL_TITLE: str = Field("RWKV-x070-World-0.1B-v2.8-20241210-ctx4096")
+    DOWNLOAD_REPO_ID: str = Field("BlinkDL/rwkv-7-world")
+    DOWNLOAD_MODEL_DIR: Union[str, None] = Field(None, description="Model Download Dir")
+    MODEL_FILE_PATH: Union[str, None] = Field(None, description="Model Path")
+    GEN_penalty_decay: float = Field(0.996, description="Default penalty decay")
+    CHUNK_LEN: int = Field(
+        256,
+        description="split input into chunks to save VRAM (shorter -> slower, but saves VRAM)",
+    )
+    VOCAB: str = Field("rwkv_vocab_v20230424", description="Vocab Name")
+
+
+CONFIG = Config()
+
+
+import numpy as np
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+os.environ["RWKV_V7_ON"] = "1"  # enable this for rwkv-7 models
+os.environ["RWKV_JIT_ON"] = "1"
+os.environ["RWKV_CUDA_ON"] = (
+    "0"  # !!! '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries !!!
+)
+
+from rwkv.model import RWKV
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
+
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+
+from api_types import (
+    ChatMessage,
+    ChatCompletion,
+    ChatCompletionChunk,
+    Usage,
+    PromptTokensDetails,
+    ChatCompletionChoice,
+    ChatCompletionMessage,
+)
+from utils import cleanMessages, parse_think_response
+
+
+logger.info(f"STRATEGY - {CONFIG.STRATEGY}")
+if CONFIG.MODEL_FILE_PATH == None:
+    CONFIG.MODEL_FILE_PATH = hf_hub_download(
+        repo_id=CONFIG.DOWNLOAD_REPO_ID,
+        filename=f"{CONFIG.MODEL_TITLE}.pth",
+        local_dir=CONFIG.DOWNLOAD_MODEL_DIR,
+    )
+
+logger.info(f"Load Model - {CONFIG.MODEL_FILE_PATH}")
+model = RWKV(model=CONFIG.MODEL_FILE_PATH.replace(".pth", ""), strategy=CONFIG.STRATEGY)
+pipeline = PIPELINE(model, CONFIG.VOCAB)
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str = Field(
+        default="rwkv-latest",
+        description="Add `:thinking` suffix to the model name to enable reasoning. Example: `rwkv-latest:thinking`",
+    )
+    messages: List[ChatMessage]
+    prompt: Union[str, None] = Field(default=None)
+    max_tokens: int = Field(default=512)
+    temperature: float = Field(default=1.0)
+    top_p: float = Field(default=0.3)
+    presencePenalty: float = Field(default=0.5)
+    countPenalty: float = Field(default=0.5)
+    stream: bool = Field(default=False)
+    state_name: str = Field(default=None)
+    include_usage: bool = Field(default=False)
+
+
+app = FastAPI(title="RWKV OpenAI-Compatible API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def runPrefill(ctx: str, model_tokens: List[int], model_state):
+    ctx = ctx.replace("\r\n", "\n")
+
+    tokens = pipeline.encode(ctx)
+    tokens = [int(x) for x in tokens]
+    model_tokens += tokens
+
+    while len(tokens) > 0:
+        out, model_state = model.forward(tokens[: CONFIG.CHUNK_LEN], model_state)
+        tokens = tokens[CONFIG.CHUNK_LEN :]
+
+    return out, model_tokens, model_state
+
+
+def generate(
+    request: ChatCompletionRequest,
+    out,
+    model_tokens,
+    model_state,
+    stops=["\n\n"],
+    max_tokens=2048,
+):
+    args = PIPELINE_ARGS(
+        temperature=max(0.2, request.temperature),
+        top_p=request.top_p,
+        alpha_frequency=request.countPenalty,
+        alpha_presence=request.presencePenalty,
+        token_ban=[],  # ban the generation of some tokens
+        token_stop=[0],
+    )  # stop generation whenever you see any token here
+
+    occurrence = {}
+    out_tokens = []
+    out_last = 0
+
+    output_cache = collections.deque(maxlen=5)
+
+    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
+
+        token = pipeline.sample_logits(
+            out, temperature=args.temperature, top_p=args.top_p
+        )
+
+        out, model_state = model.forward([token], model_state)
+        model_tokens += [token]
+
+        out_tokens += [token]
+
+        for xxx in occurrence:
+            occurrence[xxx] *= CONFIG.GEN_penalty_decay
+        occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0)
+
+        tmp: str = pipeline.decode(out_tokens[out_last:])
+
+        if "\ufffd" in tmp:
+            continue
+
+        output_cache.append(tmp)
+        output_cache_str = "".join(output_cache)
+
+        for stop_words in stops:
+            if stop_words in output_cache_str:
+
+                yield {
+                    "content": tmp.replace(stop_words, ""),
+                    "tokens": out_tokens[out_last:],
+                    "finish_reason": "stop",
+                    "state": model_state,
+                }
+
+                del out
+                gc.collect()
+                return
+
+        yield {
+            "content": tmp,
+            "tokens": out_tokens[out_last:],
+            "finish_reason": None,
+        }
+
+        out_last = i + 1
+
+    else:
+        yield {
+            "content": "",
+            "tokens": [],
+            "finish_reason": "length",
+        }
+
+
+async def chatResponse(
+    request: ChatCompletionRequest, model_state: any, completionId: str
+) -> ChatCompletion:
+    createTimestamp = time.time()
+
+    enableReasoning = request.model.endswith(":thinking")
+
+    prompt = (
+        f"{cleanMessages(request.messages)}\n\nAssistant:{' <think' if enableReasoning else ''}"
+        if request.prompt == None
+        else request.prompt.strip()
+    )
+
+    out, model_tokens, model_state = runPrefill(prompt, [], model_state)
+
+    prefillTime = time.time()
+    promptTokenCount = len(model_tokens)
+
+    fullResponse = " <think" if enableReasoning else ""
+    completionTokenCount = 0
+    finishReason = None
+
+    for chunk in generate(
+        request,
+        out,
+        model_tokens,
+        model_state,
+        max_tokens=(
+            64000
+            if "max_tokens" not in request.model_fields_set and enableReasoning
+            else request.max_tokens
+        ),
+    ):
+        fullResponse += chunk["content"]
+        completionTokenCount += 1
+
+        if chunk["finish_reason"]:
+            finishReason = chunk["finish_reason"]
+        await asyncio.sleep(0)
+
+    genenrateTime = time.time()
+
+    responseLog = {
+        "content": fullResponse,
+        "finish": finishReason,
+        "prefill_len": promptTokenCount,
+        "prefill_tps": round(promptTokenCount / (prefillTime - createTimestamp), 2),
+        "gen_len": completionTokenCount,
+        "gen_tps": round(completionTokenCount / (genenrateTime - prefillTime), 2),
+    }
+    logger.info(f"[RES] {completionId} - {responseLog}")
+
+    reasoning_content, content = parse_think_response(fullResponse)
+
+    response = ChatCompletion(
+        id=completionId,
+        created=int(createTimestamp),
+        model=request.model,
+        usage=Usage(
+            prompt_tokens=promptTokenCount,
+            completion_tokens=completionTokenCount,
+            total_tokens=promptTokenCount + completionTokenCount,
+            prompt_tokens_details={"cached_tokens": 0},
+        ),
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                message=ChatCompletionMessage(
+                    role="Assistant",
+                    content=content,
+                    reasoning_content=reasoning_content if reasoning_content else None,
+                ),
+                logprobs=None,
+                finish_reason=finishReason,
+            )
+        ],
+    )
+
+    return response
+
+
+async def chatResponseStream(
+    request: ChatCompletionRequest, model_state: any, completionId: str
+):
+    createTimestamp = int(time.time())
+
+    enableReasoning = request.model.endswith(":thinking")
+
+    prompt = (
+        f"{cleanMessages(request.messages)}\n\nAssistant:{' <think' if enableReasoning else ''}"
+        if request.prompt == None
+        else request.prompt.strip()
+    )
+
+    out, model_tokens, model_state = runPrefill(prompt, [], model_state)
+
+    prefillTime = time.time()
+    promptTokenCount = len(model_tokens)
+
+    completionTokenCount = 0
+    finishReason = None
+
+    response = ChatCompletionChunk(
+        id=completionId,
+        created=createTimestamp,
+        model=request.model,
+        usage=(
+            Usage(
+                prompt_tokens=promptTokenCount,
+                completion_tokens=completionTokenCount,
+                total_tokens=promptTokenCount + completionTokenCount,
+                prompt_tokens_details={"cached_tokens": 0},
+            )
+            if request.include_usage
+            else None
+        ),
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                delta=ChatCompletionMessage(
+                    role="Assistant",
+                    content="",
+                    reasoning_content="" if enableReasoning else None,
+                ),
+                logprobs=None,
+                finish_reason=finishReason,
+            )
+        ],
+    )
+    yield f"data: {response.model_dump_json()}\n\n"
+
+    buffer = []
+
+    if enableReasoning:
+        buffer.append(" <think")
+
+        streamConfig = {
+            "isChecking": False,
+            "fullTextCursor": 0,
+            "in_think": False,
+            "cacheStr": "",
+        }
+
+        for chunk in generate(
+            request,
+            out,
+            model_tokens,
+            model_state,
+            max_tokens=(
+                64000
+                if "max_tokens" not in request.model_fields_set and enableReasoning
+                else request.max_tokens
+            ),
+        ):
+            completionTokenCount += 1
+
+            chunkContent: str = chunk["content"]
+            buffer.append(chunkContent)
+
+            fullText = "".join(buffer)
+
+            if chunk["finish_reason"]:
+                finishReason = chunk["finish_reason"]
+
+            response = ChatCompletionChunk(
+                id=completionId,
+                created=createTimestamp,
+                model=request.model,
+                usage=(
+                    Usage(
+                        prompt_tokens=promptTokenCount,
+                        completion_tokens=completionTokenCount,
+                        total_tokens=promptTokenCount + completionTokenCount,
+                        prompt_tokens_details={"cached_tokens": 0},
+                    )
+                    if request.include_usage
+                    else None
+                ),
+                choices=[
+                    ChatCompletionChoice(
+                        index=0,
+                        delta=ChatCompletionMessage(
+                            content=None, reasoning_content=None
+                        ),
+                        logprobs=None,
+                        finish_reason=finishReason,
+                    )
+                ],
+            )
+
+            markStart = fullText.find("<", streamConfig["fullTextCursor"])
+            if not streamConfig["isChecking"] and markStart != -1:
+                streamConfig["isChecking"] = True
+
+                if streamConfig["in_think"]:
+                    response.choices[0].delta.reasoning_content = fullText[
+                        streamConfig["fullTextCursor"] : markStart
+                    ]
+                else:
+                    response.choices[0].delta.content = fullText[
+                        streamConfig["fullTextCursor"] : markStart
+                    ]
+
+                streamConfig["cacheStr"] = ""
+                streamConfig["fullTextCursor"] = markStart
+
+            if streamConfig["isChecking"]:
+                streamConfig["cacheStr"] = fullText[streamConfig["fullTextCursor"] :]
+            else:
+                if streamConfig["in_think"]:
+                    response.choices[0].delta.reasoning_content = chunkContent
+                else:
+                    response.choices[0].delta.content = chunkContent
+                streamConfig["fullTextCursor"] = len(fullText)
+
+            markEnd = fullText.find(">", streamConfig["fullTextCursor"])
+            if streamConfig["isChecking"] and markEnd != -1:
+                streamConfig["isChecking"] = False
+
+                if (
+                    not streamConfig["in_think"]
+                    and streamConfig["cacheStr"].find("<think>") != -1
+                ):
+                    streamConfig["in_think"] = True
+
+                    response.choices[0].delta.reasoning_content = (
+                        response.choices[0].delta.reasoning_content
+                        if response.choices[0].delta.reasoning_content != None
+                        else "" + streamConfig["cacheStr"].replace("<think>", "")
+                    )
+
+                elif (
+                    streamConfig["in_think"]
+                    and streamConfig["cacheStr"].find("</think>") != -1
+                ):
+                    streamConfig["in_think"] = False
+
+                    response.choices[0].delta.content = (
+                        response.choices[0].delta.content
+                        if response.choices[0].delta.content != None
+                        else "" + streamConfig["cacheStr"].replace("</think>", "")
+                    )
+                else:
+                    if streamConfig["in_think"]:
+                        response.choices[0].delta.reasoning_content = (
+                            response.choices[0].delta.reasoning_content
+                            if response.choices[0].delta.reasoning_content != None
+                            else "" + streamConfig["cacheStr"]
+                        )
+                    else:
+                        response.choices[0].delta.content = (
+                            response.choices[0].delta.content
+                            if response.choices[0].delta.content != None
+                            else "" + streamConfig["cacheStr"]
+                        )
+                streamConfig["fullTextCursor"] = len(fullText)
+
+            if (
+                response.choices[0].delta.content != None
+                or response.choices[0].delta.reasoning_content != None
+            ):
+                yield f"data: {response.model_dump_json()}\n\n"
+
+            await asyncio.sleep(0)
+
+        del streamConfig
+    else:
+        for chunk in generate(request, out, model_tokens, model_state):
+            completionTokenCount += 1
+            buffer.append(chunk["content"])
+
+            if chunk["finish_reason"]:
+                finishReason = chunk["finish_reason"]
+
+            response = ChatCompletionChunk(
+                id=completionId,
+                created=createTimestamp,
+                model=request.model,
+                usage=(
+                    Usage(
+                        prompt_tokens=promptTokenCount,
+                        completion_tokens=completionTokenCount,
+                        total_tokens=promptTokenCount + completionTokenCount,
+                        prompt_tokens_details={"cached_tokens": 0},
+                    )
+                    if request.include_usage
+                    else None
+                ),
+                choices=[
+                    ChatCompletionChoice(
+                        index=0,
+                        delta=ChatCompletionMessage(content=chunk["content"]),
+                        logprobs=None,
+                        finish_reason=finishReason,
+                    )
+                ],
+            )
+
+            yield f"data: {response.model_dump_json()}\n\n"
+            await asyncio.sleep(0)
+
+    genenrateTime = time.time()
+
+    responseLog = {
+        "content": "".join(buffer),
+        "finish": finishReason,
+        "prefill_len": promptTokenCount,
+        "prefill_tps": round(promptTokenCount / (prefillTime - createTimestamp), 2),
+        "gen_len": completionTokenCount,
+        "gen_tps": round(completionTokenCount / (genenrateTime - prefillTime), 2),
+    }
+    logger.info(f"[RES] {completionId} - {responseLog}")
+
+    del buffer
+
+    yield "data: [DONE]\n\n"
+
+
+
+
+
+@app.post("/api/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest):
+    completionId = str(next(CompletionIdGenerator))
+    logger.info(f"[REQ] {completionId} - {request.model_dump()}")
+
+    def chatResponseStreamDisconnect():  
+        gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+        logger.info(
+        f"[STATUS] vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}"
+    )
+
+    model_state = None
+
+    if request.stream:
+        r = StreamingResponse(
+            chatResponseStream(request, model_state, completionId),
+            media_type="text/event-stream",
+            background=chatResponseStreamDisconnect,
+        )
+    else:
+        r = await chatResponse(request, model_state, completionId)
+
+
+    return r
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host=CONFIG.HOST, port=CONFIG.PORT)

openai_test.py

@@ -0,0 +1,78 @@
+"""
+uv pip install openai
+"""
+
+import os
+
+import logging
+
+# logging.basicConfig(
+#     level=logging.DEBUG,
+# )
+
+os.environ["NO_PROXY"] = "127.0.0.1"
+
+from openai import OpenAI
+
+client = OpenAI(base_url="http://127.0.0.1:8000/api/v1", api_key="sk-test")
+
+
+def completionStreamTest():
+    print("[*] Stream completion: ")
+
+    completion = client.chat.completions.create(
+        model="rwkv-latest",
+        messages=[
+            {
+                "role": "User",
+                "content": "请讲个关于一只灰猫和一个小女孩之间的简短故事。",
+            },
+        ],
+        stream=True,
+        max_tokens=2048,
+    )
+
+    isReasoning = False
+
+    for chunk in completion:
+        if chunk.choices[0].delta.reasoning_content and not isReasoning:
+            print("<- Reasoning ->")
+            isReasoning = True
+        elif chunk.choices[0].delta.content and isReasoning:
+            isReasoning = False
+            print("<- Stop Reasoning ->")
+
+        if chunk.choices[0].delta.reasoning_content:
+            print(chunk.choices[0].delta.reasoning_content, end="", flush=True)
+        if chunk.choices[0].delta.content:
+            print(chunk.choices[0].delta.content, end="", flush=True)
+
+    print("")
+
+
+def completionTest():
+    completion = client.chat.completions.create(
+        model="rwkv-latest:thinking",
+        messages=[
+            {
+                "role": "User",
+                "content": "How many planets are there in our solar system?",
+            },
+        ],
+        max_tokens=2048,
+    )
+
+    print("[*] Completion: ", completion)
+
+
+if __name__ == "__main__":
+    try:
+        # completionTest()
+
+        testRounds = input("Test rounds (Default: 10) :")
+
+        for i in range(int(testRounds) if testRounds != "" else 10):
+            print("\n", "=" * 10, i + 1, "/", testRounds, "=" * 10)
+            completionStreamTest()
+    except KeyboardInterrupt:
+        pass

pyproject.toml

@@ -0,0 +1,47 @@
+[project]
+name = "rwkv-hf-space"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "fastapi[standard]>=0.115.11",
+    "huggingface-hub>=0.29.1",
+    "loguru>=0.7.3",
+    "numpy>=2.2.3",
+    "pydantic>=2.10.6",
+    "pydantic-settings>=2.8.1",
+    "pynvml>=12.0.0",
+    "rwkv==0.8.28",
+    "snowflake-id>=1.0.2",
+]
+
+[project.optional-dependencies]
+cpu = ["torch>=2.6.0"]
+cu124 = ["torch>=2.6.0"]
+cu113 = ["torch"]
+
+[tool.uv]
+conflicts = [[{ extra = "cpu" }, { extra = "cu124" }, { extra = "cu113" }]]
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cpu", extra = "cpu" },
+    { index = "pytorch-cu124", extra = "cu124" },
+    { index = "pytorch-cu113", extra = "cu113" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu113"
+url = "https://download.pytorch.org/whl/cu113"
+explicit = true

utils.py

@@ -0,0 +1,35 @@
+import re
+from typing import List, Optional, Union
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings
+
+from api_types import ChatMessage
+
+
+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()  # +9包含完整标签
+        content = full_response[think_end + 9 :].strip()
+
+    # 清理标签保留内容
+    reasoning_content = reasoning.replace("<think", "").replace("</think>", "").strip()
+    return reasoning_content, content
+
+
+def cleanMessages(messages: List[ChatMessage]):
+    promptStrList = []
+
+    for message in messages:
+        content = message.content.strip()
+        content = re.sub(r"\n+", "\n", content)
+        promptStrList.append(f"{message.role.strip()}: {content}")
+
+    return "\n\n".join(promptStrList)