Blame for app.py

"""
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
# --------------------------------------------------------------------------- #
DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "FlameF0X/TinyGPT")
AVAILABLE_MODELS = [
    "FlameF0X/TinyGPT",
    "FWKV/Myosotis-1.5-base-Preview",
]
MODEL_ID = DEFAULT_MODEL_ID if DEFAULT_MODEL_ID in AVAILABLE_MODELS else "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

def _context_window_for(loaded_model) -> int:
    candidates = [
        getattr(loaded_model.config, key, None)
        for key in ("n_positions", "max_position_embeddings", "n_ctx", "max_sequence_length")
    ]
    candidates = [
        int(value) for value in candidates
        if isinstance(value, (int, float)) and 0 < int(value) < 1_000_000
    ]
    context_window = min(candidates) if candidates else 2048
    requested = os.getenv("CONTEXT_WINDOW")
    if requested:
        try:
            context_window = min(context_window, max(128, int(requested)))
        except ValueError as exc:
            raise RuntimeError("CONTEXT_WINDOW must be an integer.") from exc
    return context_window

def _load_model(model_id: str):
    global pipe, tokenizer, model, MODEL_ID, MODEL_CONTEXT_WINDOW
    global MAX_OUTPUT_TOKENS, IS_TINYGPT
    if model_id not in AVAILABLE_MODELS:
        raise ValueError(f"Unsupported model: {model_id}")
    new_pipe = pipeline(
        "text-generation",
        model=model_id,
        trust_remote_code=True,
        device=DEVICE,
    )
    new_tokenizer = new_pipe.tokenizer
    new_model = new_pipe.model
    new_model.eval()
    if new_tokenizer.pad_token_id is None:
        new_tokenizer.pad_token = new_tokenizer.eos_token
    pipe, tokenizer, model = new_pipe, new_tokenizer, new_model
    MODEL_ID = model_id
    MODEL_CONTEXT_WINDOW = _context_window_for(model)
    MAX_OUTPUT_TOKENS = max(1, MODEL_CONTEXT_WINDOW - 64)
    IS_TINYGPT = MODEL_ID.rstrip("/").lower() == "flamef0x/tinygpt"

_load_model(MODEL_ID)

# One generation at a time (small model, usually CPU / single GPU Space).
GEN_LOCK = threading.Lock()

def select_model(model_id: str):
    # Avoid swapping tokenizer/model while a generation is using them.
    with GEN_LOCK:
        _load_model(model_id)
    return f"Active model: {MODEL_ID}"



# --------------------------------------------------------------------------- #
# 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]] = []
    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 ("user", "assistant"):
            normalized.append({"role": role, "content": content})
        elif role in ("system", "developer"):
            # System/developer prompts are intentionally ignored.
            continue
        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.")

    if not IS_TINYGPT and getattr(tokenizer, "chat_template", None):
        def render() -> str:
            template_messages: List[Dict[str, str]] = []
            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("", normalized)
    else:
        # Fallback for models without an installed chat template.
        labels = {"system": "System", "user": "User", "assistant": "Assistant"}
        def render() -> str:
            lines = []
            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

    # 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 tiny custom context 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": model_id.split("/")[0]}
            for model_id in AVAILABLE_MODELS
        ],
    }


@api.get("/v1/models/{model_id:path}", dependencies=[Depends(auth)])
def get_model(model_id: str):
    if model_id not in AVAILABLE_MODELS:
        raise HTTPException(404, "Model not found.")
    return {"id": model_id, "object": "model", "created": 0, "owned_by": model_id.split("/")[0]}


@api.post("/v1/completions", dependencies=[Depends(auth)])
def completions(req: CompletionRequest):
    if req.model in AVAILABLE_MODELS and req.model != MODEL_ID:
        select_model(req.model)
    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 req.model in AVAILABLE_MODELS and req.model != MODEL_ID:
        select_model(req.model)
    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)"}

Available models: {", ".join(AVAILABLE_MODELS)}. The Chat Completions endpoint accepts the OpenAI `messages` array on every request. System and developer messages are ignored. The built-in chat tab keeps conversation history for you. TinyGPT chat history is formatted as `User: ...` / `Assistant: ...`; generation stops at `User:` for TinyGPT.
Extra sampling parameters supported by the API: `top_k`, `repetition_penalty`, and `seed`.
"""


def _chat_ui_params(max_tokens, temperature, top_p, top_k, rep_pen):
    return {
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": int(top_k),
        "repetition_penalty": float(rep_pen),
        "stop": ["User:"] if IS_TINYGPT else [],
        "seed": None,
    }


def chat_playground(message, history, model_choice, max_tokens, temperature, top_p, top_k, rep_pen):
    """Gradio's ChatInterface callback; `history` contains all earlier turns."""
    if model_choice and model_choice != MODEL_ID:
        select_model(model_choice)
    messages: List[Dict[str, Any]] = []
    for item in history or []:
        if isinstance(item, dict) and item.get("role") in ("user", "assistant"):
            messages.append({"role": item["role"], "content": _flatten_content(item.get("content"))})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            # Backward compatibility if Gradio supplies legacy (user, assistant) pairs.
            if item[0] is not None:
                messages.append({"role": "user", "content": str(item[0])})
            if item[1] is not None:
                messages.append({"role": "assistant", "content": str(item[1])})

    messages.append({"role": "user", "content": str(message or "")})
    params = _chat_ui_params(max_tokens, temperature, top_p, top_k, rep_pen)
    prompt = messages_to_prompt(messages, max_input_tokens=MODEL_CONTEXT_WINDOW - params["max_tokens"])
    response = ""
    try:
        for event in generate(prompt, params):
            if event[0] == "text":
                response += event[1]
                yield response
    except ValueError as exc:
        raise gr.Error(str(exc))
    yield response


def playground(prompt, max_tokens, temperature, top_p, top_k, rep_pen, stop_text):
    params = {
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": int(top_k),
        "repetition_penalty": float(rep_pen),
        "stop": [s for s in (stop_text or "").split("|") if s],
        "seed": None,
    }
    out = ""
    try:
        for event in generate(prompt, params):
            if event[0] == "text":
                out += event[1]
                yield out
    except ValueError as exc:
        raise gr.Error(str(exc))
    yield out


with gr.Blocks(title=f"{MODEL_ID} API") as demo:
    gr.Markdown(DOCS)
    with gr.Tab("Multi-turn chat"):
        gr.ChatInterface(
            fn=chat_playground,
            title="Chat with selected model",
            description="Conversation history is included automatically on each turn.",
            additional_inputs=[
                gr.Dropdown(
                    choices=AVAILABLE_MODELS,
                    value=MODEL_ID,
                    label="Model",
                    interactive=True,
                ),
                gr.Slider(1, min(MAX_OUTPUT_TOKENS, MAX_NEW_TOKENS_LIMIT), value=min(256, MAX_OUTPUT_TOKENS, MAX_NEW_TOKENS_LIMIT), step=1, label="max_tokens"),
                gr.Slider(0, 2, value=0.8, step=0.05, label="temperature (0 = greedy)"),
                gr.Slider(0.05, 1, value=0.9, step=0.05, label="top_p"),
                gr.Slider(1, 200, value=50, step=1, label="top_k"),
                gr.Slider(1.0, 2.0, value=1.05, step=0.05, label="repetition_penalty"),
            ],
            additional_inputs_accordion="Generation settings",
        )
    with gr.Tab("Text completion"):
        with gr.Row():
            with gr.Column():
                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)")
                    top_p = gr.Slider(0.05, 1, value=0.9, step=0.05, label="top_p")
                    top_k = gr.Slider(1, 200, value=50, step=1, label="top_k")
                    rep_pen = gr.Slider(1.0, 2.0, value=1.05, step=0.05, label="repetition_penalty")
                    stop_text = gr.Textbox(label="Stop strings (separate with |)", value="")
                btn = gr.Button("Generate", variant="primary")
            with gr.Column():
                output = gr.Textbox(label="Completion", lines=14)
        btn.click(
            playground,
            [prompt, max_tokens, temperature, top_p, top_k, rep_pen, stop_text],
            output,
            api_name=False,
        )

def attach_api_to_gradio():
    """Copy the /v1 routes onto Gradio's own FastAPI app (same server, same port)."""
    gradio_app = demo.app
    api_routes = [r for r in api.routes if getattr(r, "path", "").startswith(("/v1", "/health"))]
    # Insert at the front so no Gradio catch-all route can shadow them.
    for route in reversed(api_routes):
        gradio_app.router.routes.insert(0, route)
    gradio_app.add_exception_handler(HTTPException, openai_style_errors)


if __name__ == "__main__":
    # Let Gradio own the port (Spaces expects it), then add the OpenAI routes.
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        prevent_thread_lock=True,
        ssr_mode=False,
    )
    attach_api_to_gradio()
    demo.block_thread()