Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 =============================================================================
-  Transformers + FastAPI — OpenAI-Compatible Server for Qwen/Qwen3.5-0.8B
+  Transformers + FastAPI — OpenAI-Compatible Server for SmolLM2-360M
   CPU-ONLY  •  TOOL CALLING  •  STREAMING  •  Port 7860 (HF Spaces)
 =============================================================================
 """
@@ -10,7 +10,7 @@ import os
 import re
 import time
 import uuid
-from threading import Lock
+from threading import Lock, Thread
 from typing import Any, Optional, Union
 
 import torch
@@ -20,17 +20,16 @@ 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"
+MODEL_NAME = "HuggingFaceTB/SmolLM2-360M"
 HOST = "0.0.0.0"
 PORT = 7860
 MAX_NEW_TOKENS = 1024
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
 app = FastAPI(
-    title="Qwen3.5-0.8B OpenAI-Compatible API (CPU)",
+    title="SmolLM2-360M OpenAI-Compatible API (CPU)",
     description="Transformers-powered inference with tool calling — runs on CPU",
     version="2.0.0",
 )
@@ -42,8 +41,8 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ━━━━━━━━━━━━━━━━━━━━━━ Pydantic Models ━━━━━━━━━━━━━━━━━━━━━━━
 
+# ━━━━━━━━━━━━━━━━━━━━━━ Pydantic Models ━━━━━━━━━━━━━━━━━━━━━━━
 
 class FunctionDef(BaseModel):
     name: str
@@ -85,7 +84,7 @@ class ChatCompletionRequest(BaseModel):
     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
+    repetition_penalty: Optional[float] = 1.1
     n: Optional[int] = 1
     tools: Optional[list[ToolDef]] = None
     tool_choice: Optional[Union[str, dict]] = None
@@ -101,7 +100,7 @@ class CompletionRequest(BaseModel):
     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
+    repetition_penalty: Optional[float] = 1.1
     n: Optional[int] = 1
 
 
@@ -110,10 +109,12 @@ class CompletionRequest(BaseModel):
 tokenizer = None
 model = None
 generate_lock = Lock()
+# Will hold all token IDs the model should stop on
+stop_token_ids: list[int] = []
 
 
 def load_model():
-    global tokenizer, model
+    global tokenizer, model, stop_token_ids
     if model is not None:
         return
 
@@ -122,8 +123,11 @@ def load_model():
 
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        use_fast=True,
+        trust_remote_code=True,
     )
+    # Ensure pad token exists
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
 
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
@@ -133,13 +137,33 @@ def load_model():
     )
     model.eval()
 
+    # Build stop-token list: eos + any ChatML special tokens the vocab has
+    _stop_ids = set()
+    _stop_ids.add(tokenizer.eos_token_id)
+    for tok_str in ["<|im_end|>", "<|endoftext|>"]:
+        tid = tokenizer.convert_tokens_to_ids(tok_str)
+        # convert_tokens_to_ids returns unk_id when token is missing
+        if tid != tokenizer.unk_token_id and tid is not None:
+            _stop_ids.add(tid)
+    stop_token_ids = list(_stop_ids)
+
+    print(f"    eos_token       = {tokenizer.eos_token!r}")
+    print(f"    stop_token_ids  = {stop_token_ids}")
     print("✅  Model loaded on CPU!\n")
 
 
-# ━━━━━━━━━━━━━━━━━━━━ Tool-Prompt Builder (Hermes) ━━━━━━━━━━━━
+# ━━━━━━━━━━━━━━━━━━━━ Chat-Prompt Builder (ChatML) ━━━━━━━━━━━━
+#
+# SmolLM2 uses the ChatML template:
+#   <|im_start|>system\n...<|im_end|>\n
+#   <|im_start|>user\n...<|im_end|>\n
+#   <|im_start|>assistant\n...<|im_end|>\n
+#
+# For tool calling we inject Hermes-style tool defs into the system prompt.
+#
 
 TOOL_SYSTEM_PROMPT_TEMPLATE = """\
-You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
+You are a helpful assistant.
 
 # Tools
 
@@ -155,9 +179,7 @@ For each function call, return a json object with function name and arguments wi
 {{"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."
-)
+NO_TOOL_SYSTEM_PROMPT = "You are a helpful assistant."
 
 
 def _serialize_tool_definitions(tools: list[ToolDef]) -> str:
@@ -281,24 +303,30 @@ def parse_tool_calls(text: str) -> tuple[Optional[str], list[dict]]:
 
 # ━━━━━━━━━━━━━━━━━━ Generation ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
+def _clean_output(text: str) -> str:
+    """Strip all known special / stop tokens from generated text."""
+    for tok in ["<|im_end|>", "<|im_start|>", "<|endoftext|>"]:
+        text = text.replace(tok, "")
+    return text.strip()
+
+
 def generate_text(prompt: str, req) -> tuple[str, int, int]:
-    """Generate text on CPU. Returns (text, prompt_tokens, completion_tokens)."""
+    """Generate 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 = {
+    gen_kwargs: dict[str, Any] = {
         "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,
+        "eos_token_id": stop_token_ids,
+        "pad_token_id": tokenizer.pad_token_id,
     }
 
     rep_penalty = getattr(req, "repetition_penalty", 1.0)
@@ -309,16 +337,12 @@ def generate_text(prompt: str, req) -> tuple[str, int, int]:
         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, "")
+    text = _clean_output(text)
 
     completion_tokens = len(new_ids)
-    return text.strip(), prompt_tokens, completion_tokens
+    return text, prompt_tokens, completion_tokens
 
 
 def generate_text_stream(prompt: str, req):
@@ -332,15 +356,15 @@ def generate_text_stream(prompt: str, req):
         tokenizer, skip_prompt=True, skip_special_tokens=False
     )
 
-    gen_kwargs = {
+    gen_kwargs: dict[str, Any] = {
         "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,
+        "eos_token_id": stop_token_ids,
+        "pad_token_id": tokenizer.pad_token_id,
         "streamer": streamer,
     }
 
@@ -352,9 +376,8 @@ def generate_text_stream(prompt: str, req):
     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 any(s in token_text for s in ["<|im_end|>", "<|endoftext|>"]):
+            cleaned = _clean_output(token_text)
             if cleaned:
                 yield cleaned
             break
@@ -430,7 +453,6 @@ def make_completion_response(
 # ━━━━━━━━━━━━━━━━━━ Streaming Helpers ━━━━━━━━━━━━━━━━━━━━━━━━
 
 def stream_chat_response(prompt: str, req):
-    """SSE streaming for non-tool-call chat completions."""
     cid = _uid()
     created = int(time.time())
 
@@ -458,7 +480,6 @@ def stream_tool_call_chunks(
     tool_calls: list[dict],
     model_name: str,
 ):
-    """SSE streaming for tool-call responses (post-generation)."""
     cid = _uid()
     created = int(time.time())
 
@@ -501,7 +522,7 @@ def stream_tool_call_chunks(
 @app.get("/")
 async def root():
     return {
-        "message": "Qwen3.5-0.8B OpenAI-Compatible API (CPU) with Tool Calling",
+        "message": "SmolLM2-360M OpenAI-Compatible API (CPU) with Tool Calling",
         "docs": "/docs",
         "endpoints": {
             "models": "/v1/models",
@@ -520,7 +541,7 @@ async def list_models():
             "id": MODEL_NAME,
             "object": "model",
             "created": int(time.time()),
-            "owned_by": "local",
+            "owned_by": "huggingface",
         }],
     }