Update app.py

app.py

@@ -1,14 +1,17 @@
 """
 =============================================================================
-  SmolLM2-360M — Ultra FAST OpenAI-Compatible API (CPU)
-  NO TOOL CALLING  •  LOW LATENCY  •  STREAMING ENABLED
+  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, Thread
-from typing import Optional, Union
+from threading import Lock
+from typing import Any, Optional, Union
 
 import torch
 import uvicorn
@@ -17,16 +20,20 @@ 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 = "HuggingFaceTB/SmolLM2-360M"
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━ CONFIG ━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+MODEL_NAME = "Qwen/Qwen3.5-0.8B"
 HOST = "0.0.0.0"
 PORT = 7860
-MAX_NEW_TOKENS = 128   # 🔥 FAST
-# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-
-app = FastAPI(title="SmolLM2-360M Fast API", version="1.0")
+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=["*"],
@@ -35,180 +42,555 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ━━━━━━━━━━━━━━━━━━━ MODELS ━━━━━━━━━━━━━━━━━━━
+# ━━━━━━━━━━━━━━━━━━━━━━ 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.3
-    max_tokens: Optional[int] = 128
+    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]]
-    max_tokens: Optional[int] = 128
+    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
 
 
-# ━━━━━━━━━━━━━━━━━━━ LOAD MODEL ━━━━━━━━━━━━━━━━━━━
+# ━━━━━━━━━━━━━━━━━━━ Model Loading (CPU) ━━━━━━━━━━━━━━━━━━━━━━
 
 tokenizer = None
 model = None
-lock = Lock()
+generate_lock = Lock()
 
 
 def load_model():
     global tokenizer, model
+    if model is not None:
+        return
 
-    print(f"🚀 Loading {MODEL_NAME} ...")
+    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,
+        trust_remote_code=True,
     )
 
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         torch_dtype=torch.float32,
         device_map="cpu",
-        low_cpu_mem_usage=True,
+        trust_remote_code=True,
     )
-
     model.eval()
-    print("✅ Model loaded!")
-
 
-# ━━━━━━━━━━━━━━━━━━━ PROMPT BUILDER ━━━━━━━━━━━━━━━━━━━
+    print("✅  Model loaded on CPU!\n")
 
-def build_prompt(messages):
-    prompt = ""
 
-    for m in messages:
-        if m.role == "user":
-            prompt += f"User: {m.content}\n"
-        elif m.role == "assistant":
-            prompt += f"Assistant: {m.content}\n"
+# ━━━━━━━━━━━━━━━━━━━━ Tool-Prompt Builder (Hermes) ━━━━━━━━━━━━
 
-    prompt += "Assistant: "
-    return prompt
+TOOL_SYSTEM_PROMPT_TEMPLATE = """\
+You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
 
+# Tools
 
-# ━━━━━━━━━━━━━━━━━━━ GENERATION ━━━━━━━━━━━━━━━━━━━
+You may call one or more functions to assist with the user query.
 
-def generate(prompt, req):
-    inputs = tokenizer(prompt, return_tensors="pt")
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tool_definitions}
+</tools>
 
-    with lock:
-        with torch.no_grad():
-            output = model.generate(
-                **inputs,
-                max_new_tokens=req.max_tokens or MAX_NEW_TOKENS,
-                do_sample=False,   # 🔥 stable + fast
-                pad_token_id=tokenizer.eos_token_id,
-            )
+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>"""
 
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return text[len(prompt):].strip()
+NO_TOOL_SYSTEM_PROMPT = (
+    "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
+)
 
 
-def generate_stream(prompt, req):
-    inputs = tokenizer(prompt, return_tensors="pt")
+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"
+                )
 
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-
-    def run():
-        with lock:
-            with torch.no_grad():
-                model.generate(
-                    **inputs,
-                    max_new_tokens=req.max_tokens or MAX_NEW_TOKENS,
-                    do_sample=False,
-                    streamer=streamer,
-                    pad_token_id=tokenizer.eos_token_id,
+        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"
                 )
 
-    Thread(target=run).start()
+        elif role == "tool":
+            parts.append(
+                f"<|im_start|>user\n"
+                f"<tool_response>\n{msg.content or ''}\n</tool_response>"
+                f"<|im_end|>\n"
+            )
 
-    for token in streamer:
-        yield token
+    parts.append("<|im_start|>assistant\n")
+    return "".join(parts)
 
 
-# ━━━━━━━━━━━━━━━━━━━ HELPERS ━━━━━━━━━━━━━━━━━━━
+# ━━━━━━━━━━━━━━━━━━ Tool-Call Parser ━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-def uid():
-    return f"chatcmpl-{uuid.uuid4().hex[:10]}"
+_TOOL_CALL_RE = re.compile(
+    r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
+    re.DOTALL,
+)
 
 
-# ━━━━━━━━━━━━━━━━━━━ ROUTES ━━━━━━━━━━━━━━━━━━━
+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,
+            },
+        })
 
-@app.get("/")
-def root():
-    return {"message": "SmolLM2 Fast API running 🚀"}
+    content = _TOOL_CALL_RE.sub("", text).strip() or None
+    return content, tool_calls
 
 
-@app.get("/health")
-def health():
-    return {"status": "ok", "model": MODEL_NAME}
+# ━━━━━━━━━━━━━━━━━━ 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)
 
-@app.post("/v1/chat/completions")
-async def chat(req: ChatCompletionRequest):
-    try:
-        prompt = build_prompt(req.messages)
+    # Slice off the prompt tokens
+    new_ids = output_ids[0][prompt_tokens:]
+    text = tokenizer.decode(new_ids, skip_special_tokens=False)
 
-        if req.stream:
-            def stream():
-                yield f'data: {{"id":"{uid()}","choices":[{{"delta":{{"role":"assistant"}}}}]}}\n\n'
+    # Clean trailing special tokens
+    for tok in ["<|im_end|>", "<|endoftext|>"]:
+        text = text.replace(tok, "")
 
-                for token in generate_stream(prompt, req):
-                    yield f'data: {{"choices":[{{"delta":{{"content":"{token}"}}}}]}}\n\n'
+    completion_tokens = len(new_ids)
+    return text.strip(), prompt_tokens, completion_tokens
 
-                yield "data: [DONE]\n\n"
 
-            return StreamingResponse(stream(), media_type="text/event-stream")
+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
 
-        text = generate(prompt, req)
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=False
+    )
 
-        return JSONResponse({
-            "id": uid(),
-            "object": "chat.completion",
-            "choices": [{
-                "index": 0,
-                "message": {"role": "assistant", "content": text},
-                "finish_reason": "stop"
+    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(500, str(e))
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 @app.post("/v1/completions")
-async def completion(req: CompletionRequest):
+async def completions(req: CompletionRequest):
     try:
-        prompt = req.prompt if isinstance(req.prompt, str) else req.prompt[0]
-        text = generate(prompt, req)
-
-        return {
-            "id": uid(),
-            "object": "text_completion",
-            "choices": [{"text": text}]
-        }
+        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(500, str(e))
+        raise HTTPException(status_code=500, detail=str(e))
 
 
-# ━━━━━━━━━━━━━━━━━━━ MAIN ━━━━━━━━━━━━━━━━━━━
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model": MODEL_NAME, "device": "cpu"}
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━ MAIN ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
 if __name__ == "__main__":
     load_model()
 
-    print(f"🔥 Running on http://{HOST}:{PORT}")
-    uvicorn.run(app, host=HOST, port=PORT)
+    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")