Update app.py

app.py

@@ -1,37 +1,596 @@
-# ============================================================
-#  Dockerfile — Qwen3.5-0.8B CPU-Only API for HF Spaces
-#  No GPU required. Port 7860.
-# ============================================================
-
-FROM python:3.11-slim
-
-# ── System deps ──
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends git && \
-    rm -rf /var/lib/apt/lists/*
-
-# ── Python deps (CPU-only torch — no CUDA bloat) ──
-RUN pip install --no-cache-dir \
-    torch --index-url https://download.pytorch.org/whl/cpu
-
-RUN pip install --no-cache-dir \
-    transformers \
-    accelerate \
-    fastapi \
-    uvicorn \
-    pydantic \
-    huggingface_hub
-
-# ── Pre-download model at build time (~1.8 GB baked into image) ──
-ENV HF_HOME=/tmp/hf_cache
-RUN python3 -c "\
-from huggingface_hub import snapshot_download; \
-snapshot_download('Qwen/Qwen3.5-0.8B', cache_dir='/tmp/hf_cache')"
-
-# ── Copy app ──
-WORKDIR /app
-COPY app.py .
-
-EXPOSE 7860
-
-CMD ["python3", "app.py"]
+"""
+=============================================================================
+  Transformers + FastAPI — OpenAI-Compatible Server for Qwen/Qwen3.5-0.8B
+  CPU-ONLY  •  TOOL CALLING  •  STREAMING  •  Port 7860 (HF Spaces)
+=============================================================================
+"""
+
+import json
+import os
+import re
+import time
+import uuid
+from threading import Lock
+from typing import Any, Optional, Union
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━ CONFIG ━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+HOST = "0.0.0.0"
+PORT = 7860
+MAX_NEW_TOKENS = 1024
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+app = FastAPI(
+    title="Qwen3.5-0.8B OpenAI-Compatible API (CPU)",
+    description="Transformers-powered inference with tool calling — runs on CPU",
+    version="2.0.0",
+)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ━━━━━━━━━━━━━━━━━━━━━━ Pydantic Models ━━━━━━━━━━━━━━━━━━━━━━━
+
+
+class FunctionDef(BaseModel):
+    name: str
+    description: Optional[str] = ""
+    parameters: Optional[dict] = None
+
+
+class ToolDef(BaseModel):
+    type: str = "function"
+    function: FunctionDef
+
+
+class FunctionCallModel(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCallObj(BaseModel):
+    id: str
+    type: str = "function"
+    function: FunctionCallModel
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: Optional[str] = None
+    tool_calls: Optional[list[ToolCallObj]] = None
+    tool_call_id: Optional[str] = None
+    name: Optional[str] = None
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str = MODEL_NAME
+    messages: list[ChatMessage]
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 1024
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, list[str]]] = None
+    frequency_penalty: Optional[float] = 0.0
+    presence_penalty: Optional[float] = 0.0
+    repetition_penalty: Optional[float] = 1.0
+    n: Optional[int] = 1
+    tools: Optional[list[ToolDef]] = None
+    tool_choice: Optional[Union[str, dict]] = None
+
+
+class CompletionRequest(BaseModel):
+    model: str = MODEL_NAME
+    prompt: Union[str, list[str]] = ""
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 512
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, list[str]]] = None
+    frequency_penalty: Optional[float] = 0.0
+    presence_penalty: Optional[float] = 0.0
+    repetition_penalty: Optional[float] = 1.0
+    n: Optional[int] = 1
+
+
+# ━━━━━━━━━━━━━━━━━━━ Model Loading (CPU) ━━━━━━━━━━━━━━━━━━━━━━
+
+tokenizer = None
+model = None
+generate_lock = Lock()
+
+
+def load_model():
+    global tokenizer, model
+    if model is not None:
+        return
+
+    print(f"\n🚀  Loading model: {MODEL_NAME} on CPU ...")
+    print(f"    HF_HOME = {os.environ.get('HF_HOME', 'default')}\n")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        use_fast=True,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        device_map="cpu",
+        trust_remote_code=True,
+    )
+    model.eval()
+
+    print("✅  Model loaded on CPU!\n")
+
+
+# ━━━━━━━━━━━━━━━━━━━━ Tool-Prompt Builder (Hermes) ━━━━━━━━━━━━
+
+TOOL_SYSTEM_PROMPT_TEMPLATE = """\
+You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
+
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tool_definitions}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": "<function-name>", "arguments": <args-json-object>}}
+</tool_call>"""
+
+NO_TOOL_SYSTEM_PROMPT = (
+    "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
+)
+
+
+def _serialize_tool_definitions(tools: list[ToolDef]) -> str:
+    lines = []
+    for t in tools:
+        obj: dict[str, Any] = {
+            "type": "function",
+            "function": {
+                "name": t.function.name,
+                "description": t.function.description or "",
+            },
+        }
+        if t.function.parameters:
+            obj["function"]["parameters"] = t.function.parameters
+        lines.append(json.dumps(obj))
+    return "\n".join(lines)
+
+
+def build_chat_prompt(
+    messages: list[ChatMessage],
+    tools: Optional[list[ToolDef]] = None,
+    tool_choice: Optional[Union[str, dict]] = None,
+) -> str:
+    parts: list[str] = []
+    has_system = any(m.role == "system" for m in messages)
+
+    if tools:
+        default_sys = TOOL_SYSTEM_PROMPT_TEMPLATE.format(
+            tool_definitions=_serialize_tool_definitions(tools),
+        )
+    else:
+        default_sys = NO_TOOL_SYSTEM_PROMPT
+
+    if not has_system:
+        parts.append(f"<|im_start|>system\n{default_sys}<|im_end|>\n")
+
+    for msg in messages:
+        role = msg.role
+
+        if role == "system":
+            base = msg.content or ""
+            if tools:
+                tool_block = TOOL_SYSTEM_PROMPT_TEMPLATE.format(
+                    tool_definitions=_serialize_tool_definitions(tools),
+                )
+                merged = f"{base}\n\n{tool_block}" if base else tool_block
+                parts.append(f"<|im_start|>system\n{merged}<|im_end|>\n")
+            else:
+                parts.append(
+                    f"<|im_start|>system\n{base or NO_TOOL_SYSTEM_PROMPT}<|im_end|>\n"
+                )
+
+        elif role == "user":
+            parts.append(f"<|im_start|>user\n{msg.content or ''}<|im_end|>\n")
+
+        elif role == "assistant":
+            if msg.tool_calls:
+                tc_text = ""
+                for tc in msg.tool_calls:
+                    args = tc.function.arguments
+                    if isinstance(args, dict):
+                        args = json.dumps(args)
+                    tc_text += (
+                        f"\n<tool_call>\n"
+                        f'{{"name": "{tc.function.name}", "arguments": {args}}}\n'
+                        f"</tool_call>"
+                    )
+                parts.append(f"<|im_start|>assistant{tc_text}<|im_end|>\n")
+            else:
+                parts.append(
+                    f"<|im_start|>assistant\n{msg.content or ''}<|im_end|>\n"
+                )
+
+        elif role == "tool":
+            parts.append(
+                f"<|im_start|>user\n"
+                f"<tool_response>\n{msg.content or ''}\n</tool_response>"
+                f"<|im_end|>\n"
+            )
+
+    parts.append("<|im_start|>assistant\n")
+    return "".join(parts)
+
+
+# ━━━━━━━━━━━━━━━━━━ Tool-Call Parser ━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+_TOOL_CALL_RE = re.compile(
+    r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
+    re.DOTALL,
+)
+
+
+def parse_tool_calls(text: str) -> tuple[Optional[str], list[dict]]:
+    tool_calls: list[dict] = []
+
+    for raw_json in _TOOL_CALL_RE.findall(text):
+        try:
+            parsed = json.loads(raw_json)
+        except json.JSONDecodeError:
+            continue
+
+        name = parsed.get("name", "")
+        arguments = parsed.get("arguments", {})
+        if isinstance(arguments, dict):
+            arguments = json.dumps(arguments)
+        elif not isinstance(arguments, str):
+            arguments = json.dumps(arguments)
+
+        tool_calls.append({
+            "id": f"call_{uuid.uuid4().hex[:24]}",
+            "type": "function",
+            "function": {
+                "name": name,
+                "arguments": arguments,
+            },
+        })
+
+    content = _TOOL_CALL_RE.sub("", text).strip() or None
+    return content, tool_calls
+
+
+# ━━━━━━━━━━━━━━━━━━ Generation ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def generate_text(prompt: str, req) -> tuple[str, int, int]:
+    """Generate text on CPU. Returns (text, prompt_tokens, completion_tokens)."""
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+    prompt_tokens = input_ids.shape[1]
+
+    max_new = req.max_tokens or MAX_NEW_TOKENS
+
+    # Build generation kwargs
+    gen_kwargs = {
+        "input_ids": input_ids,
+        "attention_mask": inputs.get("attention_mask"),
+        "max_new_tokens": max_new,
+        "do_sample": True,
+        "temperature": max(req.temperature, 0.01),
+        "top_p": req.top_p,
+        "eos_token_id": tokenizer.convert_tokens_to_ids("<|im_end|>"),
+        "pad_token_id": tokenizer.eos_token_id,
+    }
+
+    rep_penalty = getattr(req, "repetition_penalty", 1.0)
+    if rep_penalty and rep_penalty > 1.0:
+        gen_kwargs["repetition_penalty"] = rep_penalty
+
+    with generate_lock:
+        with torch.no_grad():
+            output_ids = model.generate(**gen_kwargs)
+
+    # Slice off the prompt tokens
+    new_ids = output_ids[0][prompt_tokens:]
+    text = tokenizer.decode(new_ids, skip_special_tokens=False)
+
+    # Clean trailing special tokens
+    for tok in ["<|im_end|>", "<|endoftext|>"]:
+        text = text.replace(tok, "")
+
+    completion_tokens = len(new_ids)
+    return text.strip(), prompt_tokens, completion_tokens
+
+
+def generate_text_stream(prompt: str, req):
+    """Generator that yields tokens one-by-one for streaming."""
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+
+    max_new = req.max_tokens or MAX_NEW_TOKENS
+
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=False
+    )
+
+    gen_kwargs = {
+        "input_ids": input_ids,
+        "attention_mask": inputs.get("attention_mask"),
+        "max_new_tokens": max_new,
+        "do_sample": True,
+        "temperature": max(req.temperature, 0.01),
+        "top_p": req.top_p,
+        "eos_token_id": tokenizer.convert_tokens_to_ids("<|im_end|>"),
+        "pad_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+    }
+
+    rep_penalty = getattr(req, "repetition_penalty", 1.0)
+    if rep_penalty and rep_penalty > 1.0:
+        gen_kwargs["repetition_penalty"] = rep_penalty
+
+    thread = Thread(target=_generate_in_thread, args=(gen_kwargs,))
+    thread.start()
+
+    for token_text in streamer:
+        # Stop on special tokens
+        if "<|im_end|>" in token_text or "<|endoftext|>" in token_text:
+            cleaned = token_text.replace("<|im_end|>", "").replace("<|endoftext|>", "")
+            if cleaned:
+                yield cleaned
+            break
+        yield token_text
+
+    thread.join()
+
+
+def _generate_in_thread(gen_kwargs):
+    with generate_lock:
+        with torch.no_grad():
+            model.generate(**gen_kwargs)
+
+
+# ━━━━━━━━━━━━━━━━━━ Response Builders ━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def _uid(prefix: str = "chatcmpl") -> str:
+    return f"{prefix}-{uuid.uuid4().hex[:12]}"
+
+
+def make_chat_response(
+    content: Optional[str],
+    tool_calls: list[dict],
+    model_name: str,
+    prompt_tokens: int,
+    completion_tokens: int,
+) -> dict:
+    message: dict[str, Any] = {"role": "assistant"}
+
+    if tool_calls:
+        message["content"] = content
+        message["tool_calls"] = tool_calls
+        finish_reason = "tool_calls"
+    else:
+        message["content"] = (content or "").strip()
+        finish_reason = "stop"
+
+    return {
+        "id": _uid(),
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model_name,
+        "choices": [{
+            "index": 0,
+            "message": message,
+            "finish_reason": finish_reason,
+        }],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+def make_completion_response(
+    text: str, model_name: str, prompt_tokens: int, completion_tokens: int
+) -> dict:
+    return {
+        "id": _uid("cmpl"),
+        "object": "text_completion",
+        "created": int(time.time()),
+        "model": model_name,
+        "choices": [{"index": 0, "text": text.strip(), "finish_reason": "stop"}],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+# ━━━━━━━━━━━━━━━━━━ Streaming Helpers ━━━━━━━━━━━━━━━━━━━━━━━━
+
+def stream_chat_response(prompt: str, req):
+    """SSE streaming for non-tool-call chat completions."""
+    cid = _uid()
+    created = int(time.time())
+
+    def _chunk(delta: dict, finish: Optional[str] = None) -> str:
+        return "data: " + json.dumps({
+            "id": cid,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": req.model,
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+        }) + "\n\n"
+
+    yield _chunk({"role": "assistant"})
+
+    for token_text in generate_text_stream(prompt, req):
+        if token_text:
+            yield _chunk({"content": token_text})
+
+    yield _chunk({}, finish="stop")
+    yield "data: [DONE]\n\n"
+
+
+def stream_tool_call_chunks(
+    content: Optional[str],
+    tool_calls: list[dict],
+    model_name: str,
+):
+    """SSE streaming for tool-call responses (post-generation)."""
+    cid = _uid()
+    created = int(time.time())
+
+    def _chunk(delta: dict, finish: Optional[str] = None) -> str:
+        return "data: " + json.dumps({
+            "id": cid,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model_name,
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+        }) + "\n\n"
+
+    yield _chunk({"role": "assistant"})
+
+    for idx, tc in enumerate(tool_calls):
+        yield _chunk({
+            "tool_calls": [{
+                "index": idx,
+                "id": tc["id"],
+                "type": "function",
+                "function": {"name": tc["function"]["name"], "arguments": ""},
+            }]
+        })
+        yield _chunk({
+            "tool_calls": [{
+                "index": idx,
+                "function": {"arguments": tc["function"]["arguments"]},
+            }]
+        })
+
+    if content:
+        yield _chunk({"content": content})
+
+    yield _chunk({}, finish="tool_calls" if tool_calls else "stop")
+    yield "data: [DONE]\n\n"
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━ ROUTES ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+@app.get("/")
+async def root():
+    return {
+        "message": "Qwen3.5-0.8B OpenAI-Compatible API (CPU) with Tool Calling",
+        "docs": "/docs",
+        "endpoints": {
+            "models": "/v1/models",
+            "chat": "/v1/chat/completions",
+            "completions": "/v1/completions",
+            "health": "/health",
+        },
+    }
+
+
+@app.get("/v1/models")
+async def list_models():
+    return {
+        "object": "list",
+        "data": [{
+            "id": MODEL_NAME,
+            "object": "model",
+            "created": int(time.time()),
+            "owned_by": "local",
+        }],
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(req: ChatCompletionRequest):
+    try:
+        prompt = build_chat_prompt(req.messages, req.tools, req.tool_choice)
+
+        # ── Tool-calling path (generate fully, then parse) ──
+        if req.tools:
+            text, prompt_tokens, completion_tokens = generate_text(prompt, req)
+            content, tool_calls = parse_tool_calls(text)
+
+            if req.stream:
+                return StreamingResponse(
+                    stream_tool_call_chunks(content, tool_calls, req.model),
+                    media_type="text/event-stream",
+                )
+            return JSONResponse(
+                make_chat_response(
+                    content, tool_calls, req.model, prompt_tokens, completion_tokens
+                )
+            )
+
+        # ── Normal chat (supports true token-by-token streaming) ──
+        if req.stream:
+            return StreamingResponse(
+                stream_chat_response(prompt, req),
+                media_type="text/event-stream",
+            )
+
+        text, prompt_tokens, completion_tokens = generate_text(prompt, req)
+        return JSONResponse(
+            make_chat_response(text, [], req.model, prompt_tokens, completion_tokens)
+        )
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/v1/completions")
+async def completions(req: CompletionRequest):
+    try:
+        prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
+        prompt = prompts[0]
+        text, prompt_tokens, completion_tokens = generate_text(prompt, req)
+
+        return JSONResponse(
+            make_completion_response(text, req.model, prompt_tokens, completion_tokens)
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model": MODEL_NAME, "device": "cpu"}
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━ MAIN ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+if __name__ == "__main__":
+    load_model()
+
+    print(f"\n{'='*60}")
+    print(f"  OpenAI-compatible API with TOOL CALLING (CPU)")
+    print(f"  Model:  {MODEL_NAME}")
+    print(f"  Device: CPU")
+    print(f"  URL:    http://{HOST}:{PORT}/v1")
+    print(f"{'='*60}\n")
+
+    uvicorn.run(app, host=HOST, port=PORT, log_level="info")