Update app.py

app.py

@@ -1,596 +1,37 @@
-"""
-=============================================================================
-  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/Qwen3.5-0.8B"
-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,
-        trust_remote_code=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")
+# ============================================================
+#  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"]