Upload app.py

app.py

@@ -1,4 +1,582 @@
-for you. TinyGPT is formatted using its published `Human:` / `Assistant: Response:` format.
+"""
+Gradio Space that hosts FlameF0X/TinyGPT behind an OpenAI-compatible API.
+
+Endpoints (all under /v1):
+    GET  /v1/models
+    POST /v1/completions        (best fit for a base model)
+    POST /v1/chat/completions   (uses the tokenizer chat template if it has one,
+                                 otherwise a simple "User:/Assistant:" fallback)
+
+The Gradio playground is mounted at "/".
+
+Optional environment variables / Space secrets:
+    MODEL_ID         default "FlameF0X/TinyGPT"
+    API_KEY          if set, requests must send "Authorization: Bearer <API_KEY>"
+    MAX_NEW_TOKENS   hard cap on max_tokens per request (default 1024)
+    CONTEXT_WINDOW   optional context window override, never above the model config
+"""
+
+import json
+import os
+import threading
+import time
+import uuid
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+import gradio as gr
+import torch
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel, ConfigDict
+from transformers import (
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer,
+    pipeline,
+    set_seed,
+)
+
+# --------------------------------------------------------------------------- #
+# Config + model loading
+# --------------------------------------------------------------------------- #
+MODEL_ID = os.getenv("MODEL_ID", "FlameF0X/TinyGPT")
+API_KEY = os.getenv("API_KEY")  # optional
+MAX_NEW_TOKENS_LIMIT = int(os.getenv("MAX_NEW_TOKENS", "1024"))
+DEFAULT_MAX_TOKENS = 256
+
+DEVICE = 0 if torch.cuda.is_available() else -1
+
+# Same call as in your snippet (plus a device); we reuse its model + tokenizer
+# for generation so streaming and token accounting work.
+pipe = pipeline(
+    "text-generation",
+    model=MODEL_ID,
+    trust_remote_code=True,
+    device=DEVICE,
+)
+tokenizer = pipe.tokenizer
+model = pipe.model
+model.eval()
+
+# TinyGPT is a GPT-2 model with a 2,048-token context window. Derive the real
+# limit from the loaded config instead of guessing or trusting tokenizer sentinels.
+_config_context_candidates = [
+    getattr(model.config, key, None)
+    for key in ("n_positions", "max_position_embeddings", "n_ctx")
+]
+_config_context_candidates = [
+    int(value) for value in _config_context_candidates
+    if isinstance(value, (int, float)) and 0 < int(value) < 1_000_000
+]
+MODEL_CONTEXT_WINDOW = min(_config_context_candidates) if _config_context_candidates else 2048
+_requested_context_window = os.getenv("CONTEXT_WINDOW")
+if _requested_context_window:
+    try:
+        MODEL_CONTEXT_WINDOW = min(MODEL_CONTEXT_WINDOW, max(128, int(_requested_context_window)))
+    except ValueError as exc:
+        raise RuntimeError("CONTEXT_WINDOW must be an integer.") from exc
+# Leave room for at least a small prompt, even if a client requests too many tokens.
+MAX_OUTPUT_TOKENS = max(1, MODEL_CONTEXT_WINDOW - 64)
+IS_TINYGPT = MODEL_ID.rstrip("/").lower() == "flamef0x/tinygpt"
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# One generation at a time (small model, usually CPU / single GPU Space).
+GEN_LOCK = threading.Lock()
+
+
+# --------------------------------------------------------------------------- #
+# Generation core (shared by the API and the Gradio playground)
+# --------------------------------------------------------------------------- #
+class _Control(StoppingCriteria):
+    """Lets us cancel generation (stop strings, client disconnect) and count tokens."""
+
+    def __init__(self, prompt_len: int):
+        self.prompt_len = prompt_len
+        self.generated = 0
+        self.cancel = threading.Event()
+
+    def __call__(self, input_ids, scores, **kwargs):
+        self.generated = input_ids.shape[1] - self.prompt_len
+        return torch.full(
+            (input_ids.shape[0],),
+            self.cancel.is_set(),
+            dtype=torch.bool,
+            device=input_ids.device,
+        )
+
+
+def _earliest_stop(text: str, stops: List[str]) -> int:
+    hits = [text.find(s) for s in stops]
+    hits = [h for h in hits if h >= 0]
+    return min(hits) if hits else -1
+
+
+def generate(prompt: str, p: Dict[str, Any]) -> Iterator[Tuple]:
+    """
+    Yields:
+        ("text", str)                                   incremental text
+        ("done", finish_reason, prompt_tokens, completion_tokens)   once, at the end
+    """
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs.pop("token_type_ids", None)
+    prompt_len = inputs["input_ids"].shape[1]
+    if prompt_len == 0:
+        raise ValueError("The prompt must not be empty.")
+
+    # Reserve space for the requested continuation. For raw completions, truncate
+    # old tokens from the left; chat prompts are preferably trimmed by whole turns.
+    input_budget = max(1, MODEL_CONTEXT_WINDOW - p["max_tokens"])
+    if prompt_len > input_budget:
+        inputs["input_ids"] = inputs["input_ids"][:, -input_budget:]
+        if "attention_mask" in inputs:
+            inputs["attention_mask"] = inputs["attention_mask"][:, -input_budget:]
+        prompt_len = input_budget
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    ctl = _Control(prompt_len)
+
+    gen_kwargs: Dict[str, Any] = dict(
+        **inputs,
+        streamer=streamer,
+        stopping_criteria=StoppingCriteriaList([ctl]),
+        max_new_tokens=p["max_tokens"],
+        repetition_penalty=p["repetition_penalty"],
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    if p["temperature"] > 0:
+        gen_kwargs.update(
+            do_sample=True,
+            temperature=p["temperature"],
+            top_p=p["top_p"],
+            top_k=p["top_k"],
+        )
+    else:
+        gen_kwargs.update(do_sample=False)
+
+    errors: List[Exception] = []
+
+    def worker():
+        try:
+            with torch.inference_mode():
+                model.generate(**gen_kwargs)
+        except Exception as e:  # make sure the streamer iterator can't hang
+            errors.append(e)
+            streamer.text_queue.put(streamer.stop_signal)
+
+    # TinyGPT sometimes starts generating the next speaker label after its answer.
+    # Always treat User: as a hard stop for this model, even if the client did not
+    # request stop strings. Keeping a partial-match buffer prevents leaking the
+    # marker when it arrives across multiple streamer chunks.
+    stops: List[str] = list(p["stop"])
+    if IS_TINYGPT and "User:" not in stops:
+        stops.append("User:")
+    holdback = max((len(s) for s in stops), default=1) - 1  # partial-match buffer
+    stop_hit = False
+
+    with GEN_LOCK:
+        if p.get("seed") is not None:
+            set_seed(int(p["seed"]))
+        thread = threading.Thread(target=worker, daemon=True)
+        thread.start()
+        buf = ""
+        try:
+            for chunk in streamer:
+                buf += chunk
+                if stops:
+                    idx = _earliest_stop(buf, stops)
+                    if idx >= 0:
+                        if buf[:idx]:
+                            yield ("text", buf[:idx])
+                        buf = ""
+                        stop_hit = True
+                        ctl.cancel.set()
+                        break
+                    safe = len(buf) - holdback
+                    if safe > 0:
+                        yield ("text", buf[:safe])
+                        buf = buf[safe:]
+                else:
+                    yield ("text", buf)
+                    buf = ""
+            if buf and not stop_hit:
+                yield ("text", buf)
+        finally:
+            # Also runs if the client disconnects and the generator is closed.
+            ctl.cancel.set()
+            thread.join()
+
+    if errors:
+        raise errors[0]
+
+    if stop_hit:
+        finish = "stop"
+    elif ctl.generated >= p["max_tokens"]:
+        finish = "length"
+    else:
+        finish = "stop"  # EOS
+    yield ("done", finish, prompt_len, ctl.generated)
+
+
+def run_full(prompt: str, p: Dict[str, Any]) -> Tuple[str, str, int, int]:
+    parts: List[str] = []
+    meta = ("stop", 0, 0)
+    for ev in generate(prompt, p):
+        if ev[0] == "text":
+            parts.append(ev[1])
+        else:
+            meta = ev[1:]
+    return "".join(parts), meta[0], meta[1], meta[2]
+
+
+# --------------------------------------------------------------------------- #
+# Chat formatting
+# --------------------------------------------------------------------------- #
+def _flatten_content(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):  # OpenAI "content parts"
+        return "".join(
+            part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text"
+        )
+    return str(content)
+
+
+def _token_count(text: str) -> int:
+    """Count tokenizer tokens without adding special tokens."""
+    return len(tokenizer.encode(text, add_special_tokens=False))
+
+
+def _render_tinygpt_prompt(system_text: str, dialogue: List[Dict[str, str]]) -> str:
+    """Render TinyGPT conversation history in User: ... Assistant: format."""
+    parts: List[str] = []
+    if system_text:
+        # TinyGPT's published template has no system role. Retain system/developer
+        # instructions as a plain-text preamble rather than silently discarding them.
+        parts.append(f"System: {system_text}\n\n")
+    for message in dialogue:
+        if message["role"] == "user":
+            parts.append(f"User: {message['content']}\n\n")
+        elif message["role"] == "assistant":
+            parts.append(f"Assistant: {message['content']}\n\n")
+    return "".join(parts) + "Assistant: "
+
+
+def messages_to_prompt(
+    messages: List[Dict[str, Any]], max_input_tokens: Optional[int] = None
+) -> str:
+    """Format a full OpenAI messages history and fit it within the model context.
+
+    OpenAI chat endpoints are stateless: callers send prior user and assistant
+    messages on every turn. This function preserves that history and drops oldest
+    complete turns only when the model's context window requires it.
+    """
+    if not isinstance(messages, list) or not messages:
+        raise ValueError("`messages` must be a non-empty list.")
+
+    normalized: List[Dict[str, str]] = []
+    system_parts: List[str] = []
+    for item in messages:
+        if not isinstance(item, dict):
+            raise ValueError("Each message must be an object with `role` and `content`.")
+        role = str(item.get("role", "user")).strip().lower()
+        content = _flatten_content(item.get("content")).strip()
+        if role in ("system", "developer"):
+            if content:
+                system_parts.append(content)
+        elif role in ("user", "assistant"):
+            normalized.append({"role": role, "content": content})
+        elif role == "tool":
+            # Tool execution itself is the client's responsibility. Preserve tool
+            # results in the context so the next assistant turn can use them.
+            tool_name = str(item.get("name", "tool"))
+            normalized.append({"role": "user", "content": f"[Tool result from {tool_name}]\n{content}"})
+        elif role == "function":
+            normalized.append({"role": "user", "content": f"[Function result]\n{content}"})
+        else:
+            raise ValueError(f"Unsupported message role: {role!r}.")
+
+    if not normalized:
+        raise ValueError("`messages` must contain at least one user or assistant message.")
+
+    system_text = "\n".join(system_parts)
+
+    if not IS_TINYGPT and getattr(tokenizer, "chat_template", None):
+        def render() -> str:
+            template_messages: List[Dict[str, str]] = []
+            if system_text:
+                template_messages.append({"role": "system", "content": system_text})
+            template_messages.extend(normalized)
+            return tokenizer.apply_chat_template(
+                template_messages, tokenize=False, add_generation_prompt=True
+            )
+    elif IS_TINYGPT:
+        def render() -> str:
+            return _render_tinygpt_prompt(system_text, normalized)
+    else:
+        # Fallback for models without an installed chat template.
+        labels = {"system": "System", "user": "User", "assistant": "Assistant"}
+        def render() -> str:
+            lines = []
+            if system_text:
+                lines.append(f"System: {system_text}")
+            lines.extend(f"{labels[m['role']]}: {m['content']}" for m in normalized)
+            return "\n".join(lines) + "\nAssistant:"
+
+    prompt = render()
+    if max_input_tokens is None or max_input_tokens < 1:
+        return prompt
+
+    # Shorten the system preamble first, but keep its beginning, where instructions
+    # such as the assistant's role usually live. Reserve most context for dialogue.
+    if system_text and _token_count(prompt) > max_input_tokens:
+        system_budget = min(_token_count(system_text), max(0, max_input_tokens // 4))
+        if system_budget:
+            system_ids = tokenizer.encode(system_text, add_special_tokens=False)[:system_budget]
+            system_text = tokenizer.decode(system_ids, skip_special_tokens=True)
+        else:
+            system_text = ""
+        prompt = render()
+
+    # Drop the oldest user turn and its assistant reply, never the newest request,
+    # until the remaining transcript fits.
+    while _token_count(prompt) > max_input_tokens and len(normalized) > 1:
+        remove_count = 1
+        if normalized[0]["role"] == "user" and len(normalized) > 1 and normalized[1]["role"] == "assistant":
+            remove_count = 2
+        del normalized[:remove_count]
+        prompt = render()
+
+    # A single exceptionally long message can itself exceed the context. Keep the
+    # tail of its content while preserving the role prefix and generation prompt.
+    if _token_count(prompt) > max_input_tokens and normalized:
+        last = normalized[-1]
+        original_content = last["content"]
+        last["content"] = ""
+        empty_prompt = render()
+        available = max(1, max_input_tokens - _token_count(empty_prompt) - 1)
+        content_ids = tokenizer.encode(original_content, add_special_tokens=False)
+        last["content"] = tokenizer.decode(content_ids[-available:], skip_special_tokens=True)
+        prompt = render()
+        # This can only happen with a very long system prefix or tiny custom limit.
+        if _token_count(prompt) > max_input_tokens:
+            prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)[-max_input_tokens:]
+            prompt = tokenizer.decode(prompt_ids, skip_special_tokens=True)
+
+    return prompt
+
+
+# --------------------------------------------------------------------------- #
+# OpenAI-compatible API
+# --------------------------------------------------------------------------- #
+class SamplingParams(BaseModel):
+    model_config = ConfigDict(extra="allow")  # ignore unknown OpenAI params instead of 422
+
+    model: Optional[str] = None
+    max_tokens: Optional[int] = None
+    max_completion_tokens: Optional[int] = None
+    temperature: Optional[float] = 0.8
+    top_p: Optional[float] = 1.0
+    top_k: Optional[int] = 50  # extension
+    repetition_penalty: Optional[float] = 1.0  # extension
+    stop: Optional[Union[str, List[str]]] = None
+    seed: Optional[int] = None
+    stream: Optional[bool] = False
+    stream_options: Optional[Dict[str, Any]] = None
+
+    def to_params(self) -> Dict[str, Any]:
+        stops = [self.stop] if isinstance(self.stop, str) else list(self.stop or [])
+        max_tokens = self.max_completion_tokens or self.max_tokens or DEFAULT_MAX_TOKENS
+        return dict(
+            max_tokens=max(1, min(int(max_tokens), MAX_NEW_TOKENS_LIMIT, MAX_OUTPUT_TOKENS)),
+            temperature=float(self.temperature if self.temperature is not None else 0.8),
+            top_p=float(self.top_p if self.top_p is not None else 1.0),
+            top_k=int(self.top_k if self.top_k is not None else 50),
+            repetition_penalty=float(self.repetition_penalty or 1.0),
+            stop=[s for s in stops if s],
+            seed=self.seed,
+        )
+
+
+class CompletionRequest(SamplingParams):
+    prompt: Union[str, List[str]] = ""
+
+
+class ChatRequest(SamplingParams):
+    messages: List[Dict[str, Any]]
+
+
+def auth(authorization: Optional[str] = Header(default=None)):
+    if API_KEY and authorization != f"Bearer {API_KEY}":
+        raise HTTPException(status_code=401, detail="Invalid or missing API key.")
+
+
+def sse(obj: Any) -> str:
+    return f"data: {json.dumps(obj)}\n\n"
+
+
+def usage(pt: int, ct: int) -> Dict[str, int]:
+    return {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct}
+
+
+api = FastAPI(title="TinyGPT OpenAI-compatible API")
+
+
+@api.exception_handler(HTTPException)
+async def openai_style_errors(_: Request, exc: HTTPException):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"error": {"message": str(exc.detail), "type": "invalid_request_error", "code": exc.status_code}},
+    )
+
+
+@api.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_ID}
+
+
+@api.get("/v1/models", dependencies=[Depends(auth)])
+def list_models():
+    return {
+        "object": "list",
+        "data": [{"id": MODEL_ID, "object": "model", "created": 0, "owned_by": "FlameF0X"}],
+    }
+
+
+@api.get("/v1/models/{model_id:path}", dependencies=[Depends(auth)])
+def get_model(model_id: str):
+    return {"id": MODEL_ID, "object": "model", "created": 0, "owned_by": "FlameF0X"}
+
+
+@api.post("/v1/completions", dependencies=[Depends(auth)])
+def completions(req: CompletionRequest):
+    prompts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt)
+    if not prompts:
+        raise HTTPException(400, "`prompt` is required.")
+    p = req.to_params()
+    cid, created, name = f"cmpl-{uuid.uuid4().hex}", int(time.time()), req.model or MODEL_ID
+
+    if req.stream:
+        if len(prompts) != 1:
+            raise HTTPException(400, "Streaming supports a single prompt only.")
+
+        def stream():
+            try:
+                for ev in generate(prompts[0], p):
+                    if ev[0] == "text":
+                        yield sse({
+                            "id": cid, "object": "text_completion", "created": created, "model": name,
+                            "choices": [{"index": 0, "text": ev[1], "logprobs": None, "finish_reason": None}],
+                        })
+                    else:
+                        _, finish, pt, ct = ev
+                        yield sse({
+                            "id": cid, "object": "text_completion", "created": created, "model": name,
+                            "choices": [{"index": 0, "text": "", "logprobs": None, "finish_reason": finish}],
+                        })
+                        if (req.stream_options or {}).get("include_usage"):
+                            yield sse({
+                                "id": cid, "object": "text_completion", "created": created, "model": name,
+                                "choices": [], "usage": usage(pt, ct),
+                            })
+            except Exception as e:
+                yield sse({"error": {"message": str(e), "type": "server_error"}})
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(stream(), media_type="text/event-stream")
+
+    choices, tot_pt, tot_ct = [], 0, 0
+    try:
+        for i, prompt in enumerate(prompts):
+            text, finish, pt, ct = run_full(prompt, p)
+            choices.append({"index": i, "text": text, "logprobs": None, "finish_reason": finish})
+            tot_pt, tot_ct = tot_pt + pt, tot_ct + ct
+    except ValueError as e:
+        raise HTTPException(400, str(e))
+    return {
+        "id": cid, "object": "text_completion", "created": created, "model": name,
+        "choices": choices, "usage": usage(tot_pt, tot_ct),
+    }
+
+
+@api.post("/v1/chat/completions", dependencies=[Depends(auth)])
+def chat_completions(req: ChatRequest):
+    if not req.messages:
+        raise HTTPException(400, "`messages` is required.")
+    p = req.to_params()
+    try:
+        prompt = messages_to_prompt(
+            req.messages, max_input_tokens=MODEL_CONTEXT_WINDOW - p["max_tokens"]
+        )
+    except ValueError as e:
+        raise HTTPException(400, str(e))
+
+    chat_p = dict(p)
+    cid, created, name = f"chatcmpl-{uuid.uuid4().hex}", int(time.time()), req.model or MODEL_ID
+
+    def chunk(delta: Dict[str, Any], finish: Optional[str] = None, **extra) -> str:
+        return sse({
+            "id": cid, "object": "chat.completion.chunk", "created": created, "model": name,
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], **extra,
+        })
+
+    if req.stream:
+        def stream():
+            try:
+                yield chunk({"role": "assistant", "content": ""})
+                for ev in generate(prompt, chat_p):
+                    if ev[0] == "text":
+                        yield chunk({"content": ev[1]})
+                    else:
+                        _, finish, pt, ct = ev
+                        yield chunk({}, finish)
+                        if (req.stream_options or {}).get("include_usage"):
+                            yield sse({
+                                "id": cid, "object": "chat.completion.chunk", "created": created,
+                                "model": name, "choices": [], "usage": usage(pt, ct),
+                            })
+            except Exception as e:
+                yield sse({"error": {"message": str(e), "type": "server_error"}})
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(stream(), media_type="text/event-stream")
+
+    try:
+        text, finish, pt, ct = run_full(prompt, chat_p)
+    except ValueError as e:
+        raise HTTPException(400, str(e))
+    return {
+        "id": cid, "object": "chat.completion", "created": created, "model": name,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": text.strip()},
+            "finish_reason": finish,
+        }],
+        "usage": usage(pt, ct),
+    }
+
+
+# --------------------------------------------------------------------------- #
+# Gradio playground (mounted at "/")
+# --------------------------------------------------------------------------- #
+space_host = os.getenv("SPACE_HOST")
+BASE_URL = f"https://{space_host}/v1" if space_host else "http://localhost:7860/v1"
+
+DOCS = f"""
+# {MODEL_ID} — OpenAI-compatible API
+
+**Base URL:** `{BASE_URL}`  
+**Endpoints:** `GET /v1/models` · `POST /v1/completions` · `POST /v1/chat/completions` (streaming supported)  
+**Context window:** {MODEL_CONTEXT_WINDOW} tokens  
+**API key:** {"required (Space secret `API_KEY`)" if API_KEY else "not required (set the `API_KEY` secret to enable authentication)"}
+
+The Chat Completions endpoint accepts the full OpenAI `messages` array on every request,
+including previous user and assistant turns. The built-in chat tab keeps conversation history
+for you. TinyGPT chat history is formatted as `User: ...` / `Assistant: ...`. For TinyGPT, generation stops at `User:` and that marker is never streamed to clients.
 Extra sampling parameters supported by the API: `top_k`, `repetition_penalty`, and `seed`.
 """
 
@@ -10,7 +588,7 @@ def _chat_ui_params(max_tokens, temperature, top_p, top_k, rep_pen):
         "top_p": float(top_p),
         "top_k": int(top_k),
         "repetition_penalty": float(rep_pen),
-        "stop": ["\n\nHuman:", "\nHuman:"] if IS_TINYGPT else [],
+        "stop": ["User:"] if IS_TINYGPT else [],
         "seed": None,
     }
 
@@ -86,7 +664,7 @@ with gr.Blocks(title=f"{MODEL_ID} API") as demo:
     with gr.Tab("Text completion"):
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(label="Prompt", lines=6, value="Human: How much is 1 + 1?\n\nAssistant: Response:")
+                prompt = gr.Textbox(label="Prompt", lines=6, value="User: How much is 1 + 1?\n\nAssistant:")
                 with gr.Accordion("Sampling", open=False):
                     max_tokens = gr.Slider(1, min(MAX_OUTPUT_TOKENS, MAX_NEW_TOKENS_LIMIT), value=min(128, MAX_OUTPUT_TOKENS, MAX_NEW_TOKENS_LIMIT), step=1, label="max_tokens")
                     temperature = gr.Slider(0, 2, value=0.8, step=0.05, label="temperature (0 = greedy)")