Update app.py

app.py

@@ -39,52 +39,70 @@ from transformers import (
 # --------------------------------------------------------------------------- #
 # Config + model loading
 # --------------------------------------------------------------------------- #
-MODEL_ID = os.getenv("MODEL_ID", "FlameF0X/TinyGPT")
+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
 
-# 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
+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)
@@ -279,17 +297,16 @@ def messages_to_prompt(
         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"):
+        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.
@@ -303,27 +320,21 @@ def messages_to_prompt(
     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)
+            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 = []
-            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:"
 
@@ -331,17 +342,6 @@ def messages_to_prompt(
     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:
@@ -362,7 +362,7 @@ def messages_to_prompt(
         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.
+        # 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)
@@ -443,17 +443,24 @@ def health():
 def list_models():
     return {
         "object": "list",
-        "data": [{"id": MODEL_ID, "object": "model", "created": 0, "owned_by": "FlameF0X"}],
+        "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):
-    return {"id": MODEL_ID, "object": "model", "created": 0, "owned_by": "FlameF0X"}
+    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.")
@@ -505,6 +512,8 @@ def completions(req: CompletionRequest):
 
 @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()
@@ -574,9 +583,7 @@ DOCS = f"""
 **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.
+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`.
 """
 
@@ -593,12 +600,11 @@ def _chat_ui_params(max_tokens, temperature, top_p, top_k, rep_pen):
     }
 
 
-def chat_playground(message, history, system_prompt, max_tokens, temperature, top_p, top_k, rep_pen):
+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]] = []
-    if system_prompt and str(system_prompt).strip():
-        messages.append({"role": "system", "content": str(system_prompt).strip()})
-
     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"))})
@@ -649,10 +655,15 @@ with gr.Blocks(title=f"{MODEL_ID} API") as demo:
     with gr.Tab("Multi-turn chat"):
         gr.ChatInterface(
             fn=chat_playground,
-            title="Chat with TinyGPT",
+            title="Chat with selected model",
             description="Conversation history is included automatically on each turn.",
             additional_inputs=[
-                gr.Textbox(label="System prompt", value="", placeholder="Optional behavior or role instructions", lines=2),
+                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"),