Push all files including app.py, requirements.txt, and other config files

app.py

@@ -1,258 +1,872 @@
 import os
 import time
 import gc
-import torch
+import sys
+import threading
+from itertools import islice
+from datetime import datetime
+import re  # for parsing <think> blocks
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torch
+from transformers import pipeline, TextIteratorStreamer, StoppingCriteria
+from transformers import AutoTokenizer
+from ddgs import DDGS
 import spaces  # Import spaces early to enable ZeroGPU support
+from torch.utils._pytree import tree_map
 
-# Configuration
-MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "600"))
-DEFAULT_TEMPERATURE = float(os.environ.get("DEFAULT_TEMPERATURE", "0.2"))
-DEFAULT_TOP_P = float(os.environ.get("DEFAULT_TOP_P", "0.9"))
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# Global event to signal cancellation from the UI thread to the generation thread
+cancel_event = threading.Event()
 
-MODEL_ID = "Alovestocode/router-gemma3-merged"
+access_token=os.environ['HF_TOKEN']
 
-# Global model cache
-_MODEL = None
-_TOKENIZER = None
-ACTIVE_STRATEGY = None
+# Optional: Disable GPU visibility if you wish to force CPU usage
+# os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
-# Detect ZeroGPU environment
-IS_ZEROGPU = os.environ.get("SPACE_RUNTIME_STATELESS", "0") == "1"
-if os.environ.get("SPACES_ZERO_GPU") is not None:
-    IS_ZEROGPU = True
-
-def load_model():
-    """Load the model on CPU. GPU movement happens inside @spaces.GPU decorated function."""
-    global _MODEL, _TOKENIZER, ACTIVE_STRATEGY
-    
-    if _MODEL is None:
-        print(f"Loading model {MODEL_ID}...")
-        _TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, token=HF_TOKEN)
-        
-        if IS_ZEROGPU:
-            # ZeroGPU: load on CPU with device_map=None
-            try:
-                kwargs = {
-                    "device_map": None,  # Stay on CPU for ZeroGPU
-                    "quantization_config": BitsAndBytesConfig(load_in_8bit=True),
-                    "trust_remote_code": True,
-                    "low_cpu_mem_usage": True,
-                    "token": HF_TOKEN,
-                }
-                _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
-                ACTIVE_STRATEGY = "8bit"
-            except Exception:
-                # Fallback to bf16 on CPU
-                kwargs = {
-                    "device_map": None,
-                    "torch_dtype": torch.bfloat16,
-                    "trust_remote_code": True,
-                    "low_cpu_mem_usage": True,
-                    "token": HF_TOKEN,
-                }
-                _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
-                ACTIVE_STRATEGY = "bf16"
-        else:
-            # Local environment: use GPU if available
-            if torch.cuda.is_available():
-                try:
-                    kwargs = {
-                        "device_map": "auto",
-                        "quantization_config": BitsAndBytesConfig(load_in_8bit=True),
-                        "trust_remote_code": True,
-                        "low_cpu_mem_usage": True,
-                        "token": HF_TOKEN,
-                    }
-                    _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
-                    ACTIVE_STRATEGY = "8bit"
-                except Exception:
-                    kwargs = {
-                        "device_map": "auto",
-                        "torch_dtype": torch.bfloat16,
-                        "trust_remote_code": True,
-                        "low_cpu_mem_usage": True,
-                        "token": HF_TOKEN,
-                    }
-                    _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
-                    ACTIVE_STRATEGY = "bf16"
-            else:
-                kwargs = {
-                    "device_map": "cpu",
-                    "torch_dtype": torch.float32,
-                    "trust_remote_code": True,
-                    "low_cpu_mem_usage": True,
-                    "token": HF_TOKEN,
-                }
-                _MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)
-                ACTIVE_STRATEGY = "cpu"
-        
-        _MODEL = _MODEL.eval()
-        print(f"Loaded {MODEL_ID} with strategy='{ACTIVE_STRATEGY}' (ZeroGPU={IS_ZEROGPU})")
+# ------------------------------
+# Torch-Compatible Model Definitions with Adjusted Descriptions
+# ------------------------------
+MODELS = {
+    # "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8": {
+    #     "repo_id": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8",
+    #     "description": "Sparse Mixture-of-Experts (MoE) causal language model with 80B total parameters and approximately 3B activated per inference step. Features include native 32,768-token context (extendable to 131,072 via YaRN), 16 query heads and 2 KV heads, head dimension of 256, and FP8 quantization for efficiency. Optimized for fast, stable instruction-following dialogue without 'thinking' traces, making it ideal for general chat and low-latency applications [[2]][[3]][[5]][[8]].",
+    #     "params_b": 80.0
+    # },
+    # "Qwen/Qwen3-Next-80B-A3B-Thinking-FP8": {
+    #     "repo_id": "Qwen/Qwen3-Next-80B-A3B-Thinking-FP8",
+    #     "description": "Sparse Mixture-of-Experts (MoE) causal language model with 80B total parameters and approximately 3B activated per inference step. Features include native 32,768-token context (extendable to 131,072 via YaRN), 16 query heads and 2 KV heads, head dimension of 256, and FP8 quantization. Specialized for complex reasoning, math, and coding tasks, this model outputs structured 'thinking' traces by default and is designed to be used with a reasoning parser [[10]][[11]][[14]][[18]].",
+    #     "params_b": 80.0
+    # },
+    "Qwen3-32B-FP8": {
+        "repo_id": "Qwen/Qwen3-32B-FP8",
+        "description": "Dense causal language model with 32.8B total parameters (31.2B non-embedding), 64 layers, 64 query heads & 8 KV heads, native 32,768-token context (extendable to 131,072 via YaRN). Features seamless switching between thinking mode (for complex reasoning, math, coding) and non-thinking mode (for efficient dialogue), strong multilingual support (100+ languages), and leading open-source agent capabilities.",
+        "params_b": 32.8
+    },
+    # ~30.5B total parameters (MoE: 3.3B activated)
+    # "Qwen3-30B-A3B-Instruct-2507": {
+    #     "repo_id": "Qwen/Qwen3-30B-A3B-Instruct-2507",
+    #     "description": "non-thinking-mode MoE model based on Qwen3-30B-A3B-Instruct-2507. Features 30.5B total parameters (3.3B activated), 128 experts (8 activated), 48 layers, and native 262,144-token context. Excels in instruction following, logical reasoning, multilingualism, coding, and long-context understanding. Supports only non-thinking mode (no <think> blocks). Quantized using AWQ (W4A16) with lm_head and gating layers preserved in higher precision.",
+    #     "params_b": 30.5
+    # },
+    # "Qwen3-30B-A3B-Thinking-2507": {
+    #     "repo_id": "Qwen/Qwen3-30B-A3B-Thinking-2507",
+    #     "description": "thinking-mode MoE model based on Qwen3-30B-A3B-Thinking-2507. Contains 30.5B total parameters (3.3B activated), 128 experts (8 activated), 48 layers, and 262,144-token native context. Optimized for deep reasoning in mathematics, science, coding, and agent tasks. Outputs include automatic reasoning delimiters (<think>...</think>). Quantized with AWQ (W4A16), preserving lm_head and expert gating layers.",
+    #     "params_b": 30.5
+    # },
+    "gpt-oss-20b-BF16": {
+        "repo_id": "unsloth/gpt-oss-20b-BF16",
+        "description": "A 20B-parameter open-source GPT-style language model quantized to INT4 using AutoRound, with FP8 key-value cache for efficient inference. Optimized for performance and memory efficiency on Intel hardware while maintaining strong language generation capabilities.",
+        "params_b": 20.0
+    },
+    "Qwen3-4B-Instruct-2507": {
+        "repo_id": "Qwen/Qwen3-4B-Instruct-2507",
+        "description": "Updated non-thinking instruct variant of Qwen3-4B with 4.0B parameters, featuring significant improvements in instruction following, logical reasoning, multilingualism, and 256K long-context understanding. Strong performance across knowledge, coding, alignment, and agent benchmarks.",
+        "params_b": 4.0
+    },    
+    "Apriel-1.5-15b-Thinker": {
+        "repo_id": "ServiceNow-AI/Apriel-1.5-15b-Thinker",
+        "description": "Multimodal reasoning model with 15B parameters, trained via extensive mid-training on text and image data, and fine-tuned only on text (no image SFT). Achieves competitive performance on reasoning benchmarks like Artificial Analysis (score: 52), Tau2 Bench Telecom (68), and IFBench (62). Supports both text and image understanding, fits on a single GPU, and includes structured reasoning output with tool and function calling capabilities.",
+        "params_b": 15.0
+    },
     
-    return _MODEL, _TOKENIZER
+    # 14.8B total parameters
+    "Qwen3-14B": {
+        "repo_id": "Qwen/Qwen3-14B",
+        "description": "Dense causal language model with 14.8 B total parameters (13.2 B non-embedding), 40 layers, 40 query heads & 8 KV heads, 32 768-token context (131 072 via YaRN), enhanced human preference alignment & advanced agent integration.",
+        "params_b": 14.8
+    },
+    "Qwen/Qwen3-14B-FP8": {
+        "repo_id": "Qwen/Qwen3-14B-FP8",
+        "description": "FP8-quantized version of Qwen3-14B for efficient inference.",
+        "params_b": 14.8
+    },
 
-def get_duration(prompt, max_new_tokens, temperature, top_p):
-    """Estimate generation duration for ZeroGPU scheduling."""
-    # Base time + token generation time
-    base_duration = 20
-    token_duration = max_new_tokens * 0.005  # ~200 tokens/second
-    return base_duration + token_duration
+    # ~15B (commented out in original, but larger than 14B)
+    # "Apriel-1.5-15b-Thinker": { ... },
 
-@spaces.GPU(duration=get_duration)
-def generate_response(prompt, max_new_tokens=MAX_NEW_TOKENS, temperature=DEFAULT_TEMPERATURE, top_p=DEFAULT_TOP_P):
+    # 5B
+    # "Apriel-5B-Instruct": {
+    #     "repo_id": "ServiceNow-AI/Apriel-5B-Instruct",
+    #     "description": "A 5B-parameter instruction-tuned model from ServiceNow’s Apriel series, optimized for enterprise tasks and general-purpose instruction following."
+    # },
+
+    # 4.3B
+    "Phi-4-mini-Reasoning": {
+        "repo_id": "microsoft/Phi-4-mini-reasoning",
+        "description": "Phi-4-mini-Reasoning (4.3B parameters)",
+        "params_b": 4.3
+    },
+    "Phi-4-mini-Instruct": {
+        "repo_id": "microsoft/Phi-4-mini-instruct",
+        "description": "Phi-4-mini-Instruct (4.3B parameters)",
+        "params_b": 4.3
+    },
+
+    # 4.0B
+    "Qwen3-4B": {
+        "repo_id": "Qwen/Qwen3-4B",
+        "description": "Dense causal language model with 4.0 B total parameters (3.6 B non-embedding), 36 layers, 32 query heads & 8 KV heads, native 32 768-token context (extendable to 131 072 via YaRN), balanced mid-range capacity & long-context reasoning.",
+        "params_b": 4.0
+    },
+
+    "Gemma-3-4B-IT": {
+        "repo_id": "unsloth/gemma-3-4b-it",
+        "description": "Gemma-3-4B-IT",
+        "params_b": 4.0
+    },
+    "MiniCPM3-4B": {
+        "repo_id": "openbmb/MiniCPM3-4B",
+        "description": "MiniCPM3-4B",
+        "params_b": 4.0
+    },
+    "Gemma-3n-E4B": {
+        "repo_id": "google/gemma-3n-E4B",
+        "description": "Gemma 3n base model with effective 4 B parameters (≈3 GB VRAM)",
+        "params_b": 4.0
+    },
+    "SmallThinker-4BA0.6B-Instruct": {
+        "repo_id": "PowerInfer/SmallThinker-4BA0.6B-Instruct",
+        "description": "SmallThinker 4 B backbone with 0.6 B activated parameters, instruction‑tuned",
+        "params_b": 4.0
+    },
+
+    # ~3B
+    # "AI21-Jamba-Reasoning-3B": {
+    #     "repo_id": "ai21labs/AI21-Jamba-Reasoning-3B",
+    #     "description": "A compact 3B hybrid Transformer–Mamba reasoning model with 256K context length, strong intelligence benchmark scores (61% MMLU-Pro, 52% IFBench), and efficient inference suitable for edge and datacenter use. Outperforms Gemma-3 4B and Llama-3.2 3B despite smaller size."
+    # },
+    "Qwen2.5-Taiwan-3B-Reason-GRPO": {
+        "repo_id": "benchang1110/Qwen2.5-Taiwan-3B-Reason-GRPO",
+        "description": "Qwen2.5-Taiwan model with 3 B parameters, Reason-GRPO fine-tuned",
+        "params_b": 3.0
+    },
+    "Llama-3.2-Taiwan-3B-Instruct": {
+        "repo_id": "lianghsun/Llama-3.2-Taiwan-3B-Instruct",
+        "description": "Llama-3.2-Taiwan-3B-Instruct",
+        "params_b": 3.0
+    },
+    "Qwen2.5-3B-Instruct": {
+        "repo_id": "Qwen/Qwen2.5-3B-Instruct",
+        "description": "Qwen2.5-3B-Instruct",
+        "params_b": 3.0
+    },
+    "Qwen2.5-Omni-3B": {
+        "repo_id": "Qwen/Qwen2.5-Omni-3B",
+        "description": "Qwen2.5-Omni-3B",
+        "params_b": 3.0
+    },
+    "Granite-4.0-Micro": {
+        "repo_id": "ibm-granite/granite-4.0-micro",
+        "description": "A 3B-parameter long-context instruct model from IBM, finetuned for enhanced instruction following and tool-calling. Supports 12 languages including English, Chinese, Arabic, and Japanese. Built on a dense Transformer with GQA, RoPE, SwiGLU, and 128K context length. Trained using SFT, RL alignment, and model merging techniques for enterprise applications.",
+        "params_b": 3.0
+    },
+
+    # 2.6B
+    "LFM2-2.6B": {
+        "repo_id": "LiquidAI/LFM2-2.6B",
+        "description": "The 2.6B parameter model in the LFM2 series, it outperforms models in the 3B+ class and features a hybrid architecture for faster inference.",
+        "params_b": 2.6
+    },
+
+    # 1.7B
+    "Qwen3-1.7B": {
+        "repo_id": "Qwen/Qwen3-1.7B",
+        "description": "Dense causal language model with 1.7 B total parameters (1.4 B non-embedding), 28 layers, 16 query heads & 8 KV heads, 32 768-token context, stronger reasoning vs. 0.6 B variant, dual-mode inference, instruction following across 100+ languages.",
+        "params_b": 1.7
+    },
+
+    # ~2B (effective)
+    "Gemma-3n-E2B": {
+        "repo_id": "google/gemma-3n-E2B",
+        "description": "Gemma 3n base model with effective 2 B parameters (≈2 GB VRAM)",
+        "params_b": 2.0
+    },
+
+    # 1.5B
+    "Nemotron-Research-Reasoning-Qwen-1.5B": {
+        "repo_id": "nvidia/Nemotron-Research-Reasoning-Qwen-1.5B",
+        "description": "Nemotron-Research-Reasoning-Qwen-1.5B",
+        "params_b": 1.5
+    },
+    "Falcon-H1-1.5B-Instruct": {
+        "repo_id": "tiiuae/Falcon-H1-1.5B-Instruct",
+        "description": "Falcon‑H1 model with 1.5 B parameters, instruction‑tuned",
+        "params_b": 1.5
+    },
+    "Qwen2.5-Taiwan-1.5B-Instruct": {
+        "repo_id": "benchang1110/Qwen2.5-Taiwan-1.5B-Instruct",
+        "description": "Qwen2.5-Taiwan-1.5B-Instruct",
+        "params_b": 1.5
+    },
+
+    # 1.2B
+    "LFM2-1.2B": {
+        "repo_id": "LiquidAI/LFM2-1.2B",
+        "description": "A 1.2B parameter hybrid language model from Liquid AI, designed for efficient on-device and edge AI deployment, outperforming larger models like Llama-2-7b-hf in specific tasks.",
+        "params_b": 1.2
+    },
+
+    # 1.1B
+    "Taiwan-ELM-1_1B-Instruct": {
+        "repo_id": "liswei/Taiwan-ELM-1_1B-Instruct",
+        "description": "Taiwan-ELM-1_1B-Instruct",
+        "params_b": 1.1
+    },
+
+    # 1B
+    "Llama-3.2-Taiwan-1B": {
+        "repo_id": "lianghsun/Llama-3.2-Taiwan-1B",
+        "description": "Llama-3.2-Taiwan base model with 1 B parameters",
+        "params_b": 1.0
+    },
+
+    # 700M
+    "LFM2-700M": {
+        "repo_id": "LiquidAI/LFM2-700M",
+        "description": "A 700M parameter model from the LFM2 family, designed for high efficiency on edge devices with a hybrid architecture of multiplicative gates and short convolutions.",
+        "params_b": 0.7
+    },
+
+    # 600M
+    "Qwen3-0.6B": {
+        "repo_id": "Qwen/Qwen3-0.6B",
+        "description": "Dense causal language model with 0.6 B total parameters (0.44 B non-embedding), 28 transformer layers, 16 query heads & 8 KV heads, native 32 768-token context window, dual-mode generation, full multilingual & agentic capabilities.",
+        "params_b": 0.6
+    },
+    "Qwen3-0.6B-Taiwan": {
+        "repo_id": "ShengweiPeng/Qwen3-0.6B-Taiwan",
+        "description": "Qwen3-Taiwan model with 0.6 B parameters",
+        "params_b": 0.6
+    },
+
+    # 500M
+    "Qwen2.5-0.5B-Taiwan-Instruct": {
+        "repo_id": "ShengweiPeng/Qwen2.5-0.5B-Taiwan-Instruct",
+        "description": "Qwen2.5-Taiwan model with 0.5 B parameters, instruction-tuned",
+        "params_b": 0.5
+    },
+
+    # 360M
+    "SmolLM2-360M-Instruct": {
+        "repo_id": "HuggingFaceTB/SmolLM2-360M-Instruct",
+        "description": "Original SmolLM2‑360M Instruct",
+        "params_b": 0.36
+    },
+    "SmolLM2-360M-Instruct-TaiwanChat": {
+        "repo_id": "Luigi/SmolLM2-360M-Instruct-TaiwanChat",
+        "description": "SmolLM2‑360M Instruct fine-tuned on TaiwanChat",
+        "params_b": 0.36
+    },
+
+    # 350M
+    "LFM2-350M": {
+        "repo_id": "LiquidAI/LFM2-350M",
+        "description": "A compact 350M parameter hybrid model optimized for edge and on-device applications, offering significantly faster training and inference speeds compared to models like Qwen3.",
+        "params_b": 0.35
+    },
+
+    # 270M
+    "parser_model_ner_gemma_v0.1": {
+        "repo_id": "myfi/parser_model_ner_gemma_v0.1",
+        "description": "A lightweight named‑entity‑like (NER) parser fine‑tuned from Google’s **Gemma‑3‑270M** model. The base Gemma‑3‑270M is a 270 M‑parameter, hyper‑efficient LLM designed for on‑device inference, supporting >140 languages, a 128 k‑token context window, and instruction‑following capabilities [2][7]. This variant is further trained on standard NER corpora (e.g., CoNLL‑2003, OntoNotes) to extract PERSON, ORG, LOC, and MISC entities with high precision while keeping the memory footprint low (≈240 MB VRAM in BF16 quantized form) [1]. It is released under the Apache‑2.0 license and can be used for fast, cost‑effective entity extraction in low‑resource environments.",
+        "params_b": 0.27
+    },
+    "Gemma-3-Taiwan-270M-it": {
+        "repo_id": "lianghsun/Gemma-3-Taiwan-270M-it",
+        "description": "google/gemma-3-270m-it fintuned on Taiwan Chinese dataset",
+        "params_b": 0.27
+    },
+    "gemma-3-270m-it": {
+        "repo_id": "google/gemma-3-270m-it",
+        "description": "Gemma‑3‑270M‑IT is a compact, 270‑million‑parameter language model fine‑tuned for Italian, offering fast and efficient on‑device text generation and comprehension in the Italian language.",
+        "params_b": 0.27
+    },
+    "Taiwan-ELM-270M-Instruct": {
+        "repo_id": "liswei/Taiwan-ELM-270M-Instruct",
+        "description": "Taiwan-ELM-270M-Instruct",
+        "params_b": 0.27
+    },
+
+    # 135M
+    "SmolLM2-135M-multilingual-base": {
+        "repo_id": "agentlans/SmolLM2-135M-multilingual-base",
+        "description": "SmolLM2-135M-multilingual-base",
+        "params_b": 0.135
+    },
+    "SmolLM-135M-Taiwan-Instruct-v1.0": {
+        "repo_id": "benchang1110/SmolLM-135M-Taiwan-Instruct-v1.0",
+        "description": "135-million-parameter F32 safetensors instruction-finetuned variant of SmolLM-135M-Taiwan, trained on the 416 k-example ChatTaiwan dataset for Traditional Chinese conversational and instruction-following tasks",
+        "params_b": 0.135
+    },
+    "SmolLM2_135M_Grpo_Gsm8k": {
+        "repo_id": "prithivMLmods/SmolLM2_135M_Grpo_Gsm8k",
+        "description": "SmolLM2_135M_Grpo_Gsm8k",
+        "params_b": 0.135
+    },
+    "SmolLM2-135M-Instruct": {
+        "repo_id": "HuggingFaceTB/SmolLM2-135M-Instruct",
+        "description": "Original SmolLM2‑135M Instruct",
+        "params_b": 0.135
+    },
+    "SmolLM2-135M-Instruct-TaiwanChat": {
+        "repo_id": "Luigi/SmolLM2-135M-Instruct-TaiwanChat",
+        "description": "SmolLM2‑135M Instruct fine-tuned on TaiwanChat",
+        "params_b": 0.135
+    },
+}
+
+# Global cache for pipelines to avoid re-loading.
+PIPELINES = {}
+
+def load_pipeline(model_name):
+    """
+    Load and cache a transformers pipeline for text generation.
+    Tries bfloat16, falls back to float16 or float32 if unsupported.
+    """
+    global PIPELINES
+    if model_name in PIPELINES:
+        return PIPELINES[model_name]
+    repo = MODELS[model_name]["repo_id"]
+    tokenizer = AutoTokenizer.from_pretrained(repo,
+                token=access_token)
+    for dtype in (torch.bfloat16, torch.float16, torch.float32):
+        try:
+            pipe = pipeline(
+                task="text-generation",
+                model=repo,
+                tokenizer=tokenizer,
+                trust_remote_code=True,
+                dtype=dtype, # Use `dtype` instead of deprecated `torch_dtype`
+                device_map="auto",
+                use_cache=True,      # Enable past-key-value caching
+                token=access_token)
+            PIPELINES[model_name] = pipe
+            return pipe
+        except Exception:
+            continue
+    # Final fallback
+    pipe = pipeline(
+        task="text-generation",
+        model=repo,
+        tokenizer=tokenizer,
+        trust_remote_code=True,
+        device_map="auto",
+        use_cache=True
+    )
+    PIPELINES[model_name] = pipe
+    return pipe
+
+
+def retrieve_context(query, max_results=6, max_chars=50):
     """
-    Generate response using the router model.
-    In ZeroGPU mode: model is loaded on CPU, moved to GPU here, then back to CPU after.
+    Retrieve search snippets from DuckDuckGo (runs in background).
+    Returns a list of result strings.
     """
-    if not prompt.strip():
-        return "ERROR: Prompt must not be empty."
+    try:
+        with DDGS() as ddgs:
+            return [f"{i+1}. {r.get('title','No Title')} - {r.get('body','')[:max_chars]}"
+                    for i, r in enumerate(islice(ddgs.text(query, region="wt-wt", safesearch="off", timelimit="y"), max_results))]
+    except Exception:
+        return []
+
+def format_conversation(history, system_prompt, tokenizer):
+    if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
+        messages = [{"role": "system", "content": system_prompt.strip()}] + history
+        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+    else:
+        # Fallback for base LMs without chat template
+        prompt = system_prompt.strip() + "\n"
+        for msg in history:
+            if msg['role'] == 'user':
+                prompt += "User: " + msg['content'].strip() + "\n"
+            elif msg['role'] == 'assistant':
+                prompt += "Assistant: " + msg['content'].strip() + "\n"
+        if not prompt.strip().endswith("Assistant:"):
+            prompt += "Assistant: "
+        return prompt
+
+def get_duration(user_msg, chat_history, system_prompt, enable_search, max_results, max_chars, model_name, max_tokens, temperature, top_k, top_p, repeat_penalty, search_timeout):
+    # Get model size from the MODELS dict (more reliable than string parsing)
+    model_size = MODELS[model_name].get("params_b", 4.0)  # Default to 4B if not found
     
-    global _MODEL
-    model, tokenizer = load_model()
+    # Only use AOT for models >= 2B parameters
+    use_aot = model_size >= 2
     
-    # In ZeroGPU: move model to GPU inside this @spaces.GPU function
-    current_device = torch.device("cpu")
-    if IS_ZEROGPU and torch.cuda.is_available():
-        current_device = torch.device("cuda")
-        model = model.to(current_device)
-    elif torch.cuda.is_available() and not IS_ZEROGPU:
-        current_device = torch.device("cuda")
+    # Adjusted for H200 performance: faster inference, quicker compilation
+    base_duration = 20 if not use_aot else 40  # Reduced base times
+    token_duration = max_tokens * 0.005  # ~200 tokens/second average on H200
+    search_duration = 10 if enable_search else 0  # Reduced search time
+    aot_compilation_buffer = 20 if use_aot else 0  # Faster compilation on H200
     
-    inputs = tokenizer(prompt, return_tensors="pt").to(current_device)
-    eos = tokenizer.eos_token_id
+    return base_duration + token_duration + search_duration + aot_compilation_buffer
+
+@spaces.GPU(duration=get_duration)
+def chat_response(user_msg, chat_history, system_prompt,
+                  enable_search, max_results, max_chars,
+                  model_name, max_tokens, temperature,
+                  top_k, top_p, repeat_penalty, search_timeout):
+    """
+    Generates streaming chat responses, optionally with background web search.
+    This version includes cancellation support.
+    """
+    # Clear the cancellation event at the start of a new generation
+    cancel_event.clear()
     
-    try:
-        with torch.inference_mode():
-            output_ids = model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=True,
-                eos_token_id=eos,
-                pad_token_id=eos,
+    history = list(chat_history or [])
+    history.append({'role': 'user', 'content': user_msg})
+
+    # Launch web search if enabled
+    debug = ''
+    search_results = []
+    if enable_search:
+        debug = 'Search task started.'
+        thread_search = threading.Thread(
+            target=lambda: search_results.extend(
+                retrieve_context(user_msg, int(max_results), int(max_chars))
             )
-        
-        text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        result = text[len(prompt):].strip() or text.strip()
+        )
+        thread_search.daemon = True
+        thread_search.start()
+    else:
+        debug = 'Web search disabled.'
+
+    try:
+        cur_date = datetime.now().strftime('%Y-%m-%d')
+        # merge any fetched search results into the system prompt
+        if search_results:
+            
+            enriched = system_prompt.strip() + \
+            f'''\n# The following contents are the search results related to the user's message:
+            {search_results}
+            In the search results I provide to you, each result is formatted as [webpage X begin]...[webpage X end], where X represents the numerical index of each article. Please cite the context at the end of the relevant sentence when appropriate. Use the citation format [citation:X] in the corresponding part of your answer. If a sentence is derived from multiple contexts, list all relevant citation numbers, such as [citation:3][citation:5]. Be sure not to cluster all citations at the end; instead, include them in the corresponding parts of the answer.
+            When responding, please keep the following points in mind:
+            - Today is {cur_date}.
+            - Not all content in the search results is closely related to the user's question. You need to evaluate and filter the search results based on the question.
+            - For listing-type questions (e.g., listing all flight information), try to limit the answer to 10 key points and inform the user that they can refer to the search sources for complete information. Prioritize providing the most complete and relevant items in the list. Avoid mentioning content not provided in the search results unless necessary.
+            - For creative tasks (e.g., writing an essay), ensure that references are cited within the body of the text, such as [citation:3][citation:5], rather than only at the end of the text. You need to interpret and summarize the user's requirements, choose an appropriate format, fully utilize the search results, extract key information, and generate an answer that is insightful, creative, and professional. Extend the length of your response as much as possible, addressing each point in detail and from multiple perspectives, ensuring the content is rich and thorough.
+            - If the response is lengthy, structure it well and summarize it in paragraphs. If a point-by-point format is needed, try to limit it to 5 points and merge related content.
+            - For objective Q&A, if the answer is very brief, you may add one or two related sentences to enrich the content.
+            - Choose an appropriate and visually appealing format for your response based on the user's requirements and the content of the answer, ensuring strong readability.
+            - Your answer should synthesize information from multiple relevant webpages and avoid repeatedly citing the same webpage.
+            - Unless the user requests otherwise, your response should be in the same language as the user's question.
+            # The user's message is:
+            '''
+        else:
+            enriched = system_prompt
+
+        # wait up to 1s for snippets, then replace debug with them
+        if enable_search:
+            thread_search.join(timeout=float(search_timeout))
+            if search_results:
+                debug = "### Search results merged into prompt\n\n" + "\n".join(
+                    f"- {r}" for r in search_results
+                )
+            else:
+                debug = "*No web search results found.*"
+
+        # merge fetched snippets into the system prompt
+        if search_results:
+            enriched = system_prompt.strip() + \
+            f'''\n# The following contents are the search results related to the user's message:
+            {search_results}
+            In the search results I provide to you, each result is formatted as [webpage X begin]...[webpage X end], where X represents the numerical index of each article. Please cite the context at the end of the relevant sentence when appropriate. Use the citation format [citation:X] in the corresponding part of your answer. If a sentence is derived from multiple contexts, list all relevant citation numbers, such as [citation:3][citation:5]. Be sure not to cluster all citations at the end; instead, include them in the corresponding parts of the answer.
+            When responding, please keep the following points in mind:
+            - Today is {cur_date}.
+            - Not all content in the search results is closely related to the user's question. You need to evaluate and filter the search results based on the question.
+            - For listing-type questions (e.g., listing all flight information), try to limit the answer to 10 key points and inform the user that they can refer to the search sources for complete information. Prioritize providing the most complete and relevant items in the list. Avoid mentioning content not provided in the search results unless necessary.
+            - For creative tasks (e.g., writing an essay), ensure that references are cited within the body of the text, such as [citation:3][citation:5], rather than only at the end of the text. You need to interpret and summarize the user's requirements, choose an appropriate format, fully utilize the search results, extract key information, and generate an answer that is insightful, creative, and professional. Extend the length of your response as much as possible, addressing each point in detail and from multiple perspectives, ensuring the content is rich and thorough.
+            - If the response is lengthy, structure it well and summarize it in paragraphs. If a point-by-point format is needed, try to limit it to 5 points and merge related content.
+            - For objective Q&A, if the answer is very brief, you may add one or two related sentences to enrich the content.
+            - Choose an appropriate and visually appealing format for your response based on the user's requirements and the content of the answer, ensuring strong readability.
+            - Your answer should synthesize information from multiple relevant webpages and avoid repeatedly citing the same webpage.
+            - Unless the user requests otherwise, your response should be in the same language as the user's question.
+            # The user's message is:
+            '''
+        else:
+            enriched = system_prompt
+
+        pipe = load_pipeline(model_name)
+
+        prompt = format_conversation(history, enriched, pipe.tokenizer)
+        prompt_debug = f"\n\n--- Prompt Preview ---\n```\n{prompt}\n```"
+        streamer = TextIteratorStreamer(pipe.tokenizer,
+                                        skip_prompt=True,
+                                        skip_special_tokens=True)
+        gen_thread = threading.Thread(
+            target=pipe,
+            args=(prompt,),
+            kwargs={
+                'max_new_tokens': max_tokens,
+                'temperature': temperature,
+                'top_k': top_k,
+                'top_p': top_p,
+                'repetition_penalty': repeat_penalty,
+                'streamer': streamer,
+                'return_full_text': False,
+            }
+        )
+        gen_thread.start()
+
+        # Buffers for thought vs answer
+        thought_buf = ''
+        answer_buf = ''
+        in_thought = False
+        assistant_message_started = False
+
+        # First yield contains the user message
+        yield history, debug
+
+        # Stream tokens
+        for chunk in streamer:
+            # Check for cancellation signal
+            if cancel_event.is_set():
+                if assistant_message_started and history and history[-1]['role'] == 'assistant':
+                    history[-1]['content'] += " [Generation Canceled]"
+                yield history, debug
+                break
+            
+            text = chunk
+
+            # Detect start of thinking
+            if not in_thought and '<think>' in text:
+                in_thought = True
+                history.append({'role': 'assistant', 'content': '', 'metadata': {'title': '💭 Thought'}})
+                assistant_message_started = True
+                after = text.split('<think>', 1)[1]
+                thought_buf += after
+                if '</think>' in thought_buf:
+                    before, after2 = thought_buf.split('</think>', 1)
+                    history[-1]['content'] = before.strip()
+                    in_thought = False
+                    answer_buf = after2
+                    history.append({'role': 'assistant', 'content': answer_buf})
+                else:
+                    history[-1]['content'] = thought_buf
+                yield history, debug
+                continue
+
+            if in_thought:
+                thought_buf += text
+                if '</think>' in thought_buf:
+                    before, after2 = thought_buf.split('</think>', 1)
+                    history[-1]['content'] = before.strip()
+                    in_thought = False
+                    answer_buf = after2
+                    history.append({'role': 'assistant', 'content': answer_buf})
+                else:
+                    history[-1]['content'] = thought_buf
+                yield history, debug
+                continue
+
+            # Stream answer
+            if not assistant_message_started:
+                history.append({'role': 'assistant', 'content': ''})
+                assistant_message_started = True
+
+            answer_buf += text
+            history[-1]['content'] = answer_buf.strip()
+            yield history, debug
+
+        gen_thread.join()
+        yield history, debug + prompt_debug
+    except GeneratorExit:
+        # Handle cancellation gracefully
+        print("Chat response cancelled.")
+        # Don't yield anything - let the cancellation propagate
+        return
+    except Exception as e:
+        history.append({'role': 'assistant', 'content': f"Error: {e}"})
+        yield history, debug
     finally:
-        # In ZeroGPU: move model back to CPU to free GPU memory
-        if IS_ZEROGPU and torch.cuda.is_available():
-            _MODEL = model.to(torch.device("cpu"))
-            torch.cuda.empty_cache()
-    
-    return result
+        gc.collect()
 
+
+def update_default_prompt(enable_search):
+    return f"You are a helpful assistant."
+
+def update_duration_estimate(model_name, enable_search, max_results, max_chars, max_tokens, search_timeout):
+    """Calculate and format the estimated GPU duration for current settings."""
+    try:
+        dummy_msg, dummy_history, dummy_system_prompt = "", [], ""
+        duration = get_duration(dummy_msg, dummy_history, dummy_system_prompt, 
+                              enable_search, max_results, max_chars, model_name, 
+                              max_tokens, 0.7, 40, 0.9, 1.2, search_timeout)
+        model_size = MODELS[model_name].get("params_b", 4.0)
+        return (f"⏱️ **Estimated GPU Time: {duration:.1f} seconds**\n\n"
+                f"📊 **Model Size:** {model_size:.1f}B parameters\n"
+                f"🔍 **Web Search:** {'Enabled' if enable_search else 'Disabled'}")
+    except Exception as e:
+        return f"⚠️ Error calculating estimate: {e}"
+
+# ------------------------------
 # Gradio UI
+# ------------------------------
 with gr.Blocks(
-    title="Router Model API - ZeroGPU",
+    title="LLM Inference with ZeroGPU",
     theme=gr.themes.Soft(
         primary_hue="indigo",
         secondary_hue="purple",
         neutral_hue="slate",
         radius_size="lg",
+        font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
     ),
     css="""
-        .main-header {
-            text-align: center;
-            padding: 20px;
-            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-            color: white;
-            border-radius: 10px;
-            margin-bottom: 20px;
-        }
-        .info-box {
-            background: #f0f0f0;
-            padding: 15px;
-            border-radius: 8px;
-            margin-bottom: 20px;
-        }
+        .duration-estimate { background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%); border-left: 4px solid #667eea; padding: 12px; border-radius: 8px; margin: 16px 0; }
+        .chatbot { border-radius: 12px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); }
+        button.primary { font-weight: 600; }
+        .gradio-accordion { margin-bottom: 12px; }
     """
 ) as demo:
     # Header
     gr.Markdown("""
-    <div class="main-header">
-        <h1>🚀 Router Model API - ZeroGPU</h1>
-        <p>Intelligent routing agent for coordinating specialized AI agents</p>
-    </div>
+    # 🧠 ZeroGPU LLM Inference
+    ### Powered by Hugging Face ZeroGPU with Web Search Integration
     """)
     
     with gr.Row():
-        # Left Panel - Input
-        with gr.Column(scale=1):
-            gr.Markdown("### 📝 Input")
-            prompt_input = gr.Textbox(
-                label="Router Prompt",
-                lines=10,
-                placeholder="Enter your router prompt here...\n\nExample:\nYou are the Router Agent coordinating Math, Code, and General-Search specialists.\nUser query: Solve the integral of x^2 from 0 to 1",
+        # Left Panel - Configuration
+        with gr.Column(scale=3):
+            # Core Settings (Always Visible)
+            with gr.Group():
+                gr.Markdown("### ⚙️ Core Settings")
+                model_dd = gr.Dropdown(
+                    label="🤖 Model",
+                    choices=list(MODELS.keys()),
+                    value="Qwen3-1.7B",
+                    info="Select the language model to use"
+                )
+                search_chk = gr.Checkbox(
+                    label="🔍 Enable Web Search",
+                    value=False,
+                    info="Augment responses with real-time web data"
+                )
+                sys_prompt = gr.Textbox(
+                    label="📝 System Prompt",
+                    lines=3,
+                    value=update_default_prompt(search_chk.value),
+                    placeholder="Define the assistant's behavior and personality..."
+                )
+            
+            # Duration Estimate
+            duration_display = gr.Markdown(
+                value=update_duration_estimate("Qwen3-1.7B", False, 4, 50, 1024, 5.0),
+                elem_classes="duration-estimate"
             )
             
-            with gr.Accordion("⚙️ Generation Parameters", open=True):
-                max_tokens_input = gr.Slider(
-                    minimum=64,
-                    maximum=2048,
-                    value=MAX_NEW_TOKENS,
-                    step=16,
-                    label="Max New Tokens",
-                    info="Maximum number of tokens to generate"
+            # Advanced Settings (Collapsible)
+            with gr.Accordion("🎛️ Advanced Generation Parameters", open=False):
+                max_tok = gr.Slider(
+                    64, 16384, value=1024, step=32,
+                    label="Max Tokens",
+                    info="Maximum length of generated response"
                 )
-                temp_input = gr.Slider(
-                    minimum=0.0,
-                    maximum=2.0,
-                    value=DEFAULT_TEMPERATURE,
-                    step=0.05,
+                temp = gr.Slider(
+                    0.1, 2.0, value=0.7, step=0.1,
                     label="Temperature",
-                    info="Controls randomness: lower = more deterministic"
+                    info="Higher = more creative, Lower = more focused"
                 )
-                top_p_input = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=DEFAULT_TOP_P,
-                    step=0.05,
-                    label="Top-p (Nucleus Sampling)",
-                    info="Probability mass to consider for sampling"
+                with gr.Row():
+                    k = gr.Slider(
+                        1, 100, value=40, step=1,
+                        label="Top-K",
+                        info="Number of top tokens to consider"
+                    )
+                    p = gr.Slider(
+                        0.1, 1.0, value=0.9, step=0.05,
+                        label="Top-P",
+                        info="Nucleus sampling threshold"
+                    )
+                rp = gr.Slider(
+                    1.0, 2.0, value=1.2, step=0.1,
+                    label="Repetition Penalty",
+                    info="Penalize repeated tokens"
                 )
             
-            generate_btn = gr.Button("🚀 Generate", variant="primary")
-            clear_btn = gr.Button("🗑️ Clear", variant="secondary")
+            # Web Search Settings (Collapsible)
+            with gr.Accordion("🌐 Web Search Settings", open=False, visible=False) as search_settings:
+                mr = gr.Number(
+                    value=4, precision=0,
+                    label="Max Results",
+                    info="Number of search results to retrieve"
+                )
+                mc = gr.Number(
+                    value=50, precision=0,
+                    label="Max Chars/Result",
+                    info="Character limit per search result"
+                )
+                st = gr.Slider(
+                    minimum=0.0, maximum=30.0, step=0.5, value=5.0,
+                    label="Search Timeout (s)",
+                    info="Maximum time to wait for search results"
+                )
+            
+            # Actions
+            with gr.Row():
+                clr = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1)
         
-        # Right Panel - Output
-        with gr.Column(scale=1):
-            gr.Markdown("### 📤 Output")
-            output = gr.Textbox(
-                label="Generated Response",
-                lines=25,
-                placeholder="Generated response will appear here...",
+        # Right Panel - Chat Interface
+        with gr.Column(scale=7):
+            chat = gr.Chatbot(
+                type="messages",
+                height=600,
+                label="💬 Conversation",
                 show_copy_button=True,
+                avatar_images=(None, "🤖"),
+                bubble_full_width=False
             )
             
-            with gr.Accordion("📚 Model Information", open=False):
-                gr.Markdown(f"""
-                **Model:** `{MODEL_ID}`  
-                **Strategy:** `{ACTIVE_STRATEGY or 'pending'}`  
-                **ZeroGPU:** `{IS_ZEROGPU}`  
-                **Max Tokens:** `{MAX_NEW_TOKENS}`  
-                **Default Temperature:** `{DEFAULT_TEMPERATURE}`  
-                **Default Top-p:** `{DEFAULT_TOP_P}`
-                """)
+            # Input Area
+            with gr.Row():
+                txt = gr.Textbox(
+                    placeholder="💭 Type your message here... (Press Enter to send)",
+                    scale=9,
+                    container=False,
+                    show_label=False,
+                    lines=1,
+                    max_lines=5
+                )
+                with gr.Column(scale=1, min_width=120):
+                    submit_btn = gr.Button("📤 Send", variant="primary", size="lg")
+                    cancel_btn = gr.Button("⏹️ Stop", variant="stop", visible=False, size="lg")
+            
+            # Example Prompts
+            gr.Examples(
+                examples=[
+                    ["Explain quantum computing in simple terms"],
+                    ["Write a Python function to calculate fibonacci numbers"],
+                    ["What are the latest developments in AI? (Enable web search)"],
+                    ["Tell me a creative story about a time traveler"],
+                    ["Help me debug this code: def add(a,b): return a+b+1"]
+                ],
+                inputs=txt,
+                label="💡 Example Prompts"
+            )
+            
+            # Debug/Status Info (Collapsible)
+            with gr.Accordion("🔍 Debug Info", open=False):
+                dbg = gr.Markdown()
     
-    # Event handlers
-    generate_btn.click(
-        fn=generate_response,
-        inputs=[prompt_input, max_tokens_input, temp_input, top_p_input],
-        outputs=output,
-    )
+    # Footer
+    gr.Markdown("""
+    ---
+    💡 **Tips:** 
+    - Use **Advanced Parameters** to fine-tune creativity and response length
+    - Enable **Web Search** for real-time, up-to-date information
+    - Try different **models** for various tasks (reasoning, coding, general chat)
+    - Click the **Copy** button on responses to save them to your clipboard
+    """, elem_classes="footer")
+
+    # --- Event Listeners ---
+
+    # Group all inputs for cleaner event handling
+    chat_inputs = [txt, chat, sys_prompt, search_chk, mr, mc, model_dd, max_tok, temp, k, p, rp, st]
+    # Group all UI components that can be updated.
+    ui_components = [chat, dbg, txt, submit_btn, cancel_btn]
+
+    def submit_and_manage_ui(user_msg, chat_history, *args):
+        """
+        Orchestrator function that manages UI state and calls the backend chat function.
+        It uses a try...finally block to ensure the UI is always reset.
+        """
+        if not user_msg.strip():
+            # If the message is empty, do nothing.
+            # We yield an empty dict to avoid any state changes.
+            yield {}
+            return
+
+        # 1. Update UI to "generating" state.
+        #    Crucially, we do NOT update the `chat` component here, as the backend
+        #    will provide the correctly formatted history in the first response chunk.
+        yield {
+            txt: gr.update(value="", interactive=False),
+            submit_btn: gr.update(interactive=False),
+            cancel_btn: gr.update(visible=True),
+        }
+
+        cancelled = False
+        try:
+            # 2. Call the backend and stream updates
+            backend_args = [user_msg, chat_history] + list(args)
+            for response_chunk in chat_response(*backend_args):
+                yield {
+                    chat: response_chunk[0],
+                    dbg: response_chunk[1],
+                }
+        except GeneratorExit:
+            # Mark as cancelled and re-raise to prevent "generator ignored GeneratorExit"
+            cancelled = True
+            print("Generation cancelled by user.")
+            raise
+        except Exception as e:
+            print(f"An error occurred during generation: {e}")
+            # If an error happens, add it to the chat history to inform the user.
+            error_history = (chat_history or []) + [
+                {'role': 'user', 'content': user_msg},
+                {'role': 'assistant', 'content': f"**An error occurred:** {str(e)}"}
+            ]
+            yield {chat: error_history}
+        finally:
+            # Only reset UI if not cancelled (to avoid "generator ignored GeneratorExit")
+            if not cancelled:
+                print("Resetting UI state.")
+                yield {
+                    txt: gr.update(interactive=True),
+                    submit_btn: gr.update(interactive=True),
+                    cancel_btn: gr.update(visible=False),
+                }
+
+    def set_cancel_flag():
+        """Called by the cancel button, sets the global event."""
+        cancel_event.set()
+        print("Cancellation signal sent.")
     
-    clear_btn.click(
-        fn=lambda: ("", ""),
-        inputs=None,
-        outputs=[prompt_input, output],
+    def reset_ui_after_cancel():
+        """Reset UI components after cancellation."""
+        cancel_event.clear()  # Clear the flag for next generation
+        print("UI reset after cancellation.")
+        return {
+            txt: gr.update(interactive=True),
+            submit_btn: gr.update(interactive=True),
+            cancel_btn: gr.update(visible=False),
+        }
+
+    # Event for submitting text via Enter key or Submit button
+    submit_event = txt.submit(
+        fn=submit_and_manage_ui,
+        inputs=chat_inputs,
+        outputs=ui_components,
+    )
+    submit_btn.click(
+        fn=submit_and_manage_ui,
+        inputs=chat_inputs,
+        outputs=ui_components,
     )
 
-# Launch the app
-if __name__ == "__main__":
-    print("Warm start skipped for ZeroGPU. Model will load on first request.")
-    demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    # Event for the "Cancel" button.
+    # It sets the cancel flag, cancels the submit event, then resets the UI.
+    cancel_btn.click(
+        fn=set_cancel_flag,
+        cancels=[submit_event]
+    ).then(
+        fn=reset_ui_after_cancel,
+        outputs=ui_components
+    )
+
+    # Listeners for updating the duration estimate
+    duration_inputs = [model_dd, search_chk, mr, mc, max_tok, st]
+    for component in duration_inputs:
+        component.change(fn=update_duration_estimate, inputs=duration_inputs, outputs=duration_display)
+
+    # Toggle web search settings visibility
+    def toggle_search_settings(enabled):
+        return gr.update(visible=enabled)
+    
+    search_chk.change(
+        fn=lambda enabled: (update_default_prompt(enabled), gr.update(visible=enabled)),
+        inputs=search_chk,
+        outputs=[sys_prompt, search_settings]
+    )
+    
+    # Clear chat action
+    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
+    
+    demo.launch()

apt.txt

@@ -0,0 +1,2 @@
+rustc
+cargo

requirements.txt

@@ -1,9 +1,12 @@
-bitsandbytes>=0.41.0
-fastapi>=0.110.0
-accelerate>=0.25.0
-spaces>=0.40.0
-torch>=2.1.0
-transformers>=4.40.0
-uvicorn>=0.22.0
-sentencepiece>=0.1.99
-gradio>=4.0.0
+wheel
+streamlit
+ddgs
+gradio>=5.0.0
+torch>=2.8.0
+transformers>=4.53.3
+spaces
+sentencepiece
+accelerate
+autoawq
+timm
+compressed-tensors

style.css

@@ -0,0 +1,150 @@
+/* Custom CSS for LLM Inference Interface */
+
+/* Header styling */
+.markdown h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    font-weight: 800;
+    margin-bottom: 0.5rem;
+}
+
+.markdown h3 {
+    color: #4a5568;
+    font-weight: 600;
+    margin-top: 0.25rem;
+}
+
+/* Duration estimate styling */
+.duration-estimate {
+    background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%);
+    border-left: 4px solid #667eea;
+    padding: 12px;
+    border-radius: 8px;
+    margin: 16px 0;
+    font-size: 0.9em;
+}
+
+/* Group styling for better visual separation */
+.gradio-group {
+    border: 1px solid #e2e8f0;
+    border-radius: 12px;
+    padding: 16px;
+    background: #f8fafc;
+    margin-bottom: 16px;
+}
+
+/* Accordion styling */
+.gradio-accordion {
+    border: 1px solid #e2e8f0;
+    border-radius: 8px;
+    margin-bottom: 12px;
+}
+
+.gradio-accordion .label-wrap {
+    background: #f1f5f9;
+    font-weight: 600;
+}
+
+/* Chat interface improvements */
+.chatbot {
+    border-radius: 12px;
+    border: 1px solid #e2e8f0;
+    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
+}
+
+/* Input area styling */
+.textbox-container {
+    border-radius: 24px;
+    border: 2px solid #e2e8f0;
+    transition: border-color 0.2s;
+}
+
+.textbox-container:focus-within {
+    border-color: #667eea;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1);
+}
+
+/* Button improvements */
+.gradio-button {
+    border-radius: 8px;
+    font-weight: 600;
+    transition: all 0.2s;
+}
+
+.gradio-button.primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    border: none;
+}
+
+.gradio-button.primary:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
+}
+
+.gradio-button.secondary {
+    border: 2px solid #e2e8f0;
+    background: white;
+}
+
+.gradio-button.secondary:hover {
+    border-color: #cbd5e0;
+    background: #f7fafc;
+}
+
+/* Slider styling */
+.gradio-slider {
+    margin: 8px 0;
+}
+
+.gradio-slider input[type="range"] {
+    accent-color: #667eea;
+}
+
+/* Info text styling */
+.info {
+    color: #718096;
+    font-size: 0.85em;
+    font-style: italic;
+}
+
+/* Footer styling */
+.footer .markdown {
+    text-align: center;
+    color: #718096;
+    font-size: 0.9em;
+    padding: 16px;
+    background: #f8fafc;
+    border-radius: 8px;
+}
+
+/* Responsive adjustments */
+@media (max-width: 768px) {
+    .gradio-row {
+        flex-direction: column;
+    }
+    
+    .chatbot {
+        height: 400px !important;
+    }
+}
+
+/* Loading animation */
+@keyframes pulse {
+    0%, 100% {
+        opacity: 1;
+    }
+    50% {
+        opacity: 0.5;
+    }
+}
+
+.generating {
+    animation: pulse 1.5s ease-in-out infinite;
+}
+
+/* Smooth transitions */
+* {
+    transition: background-color 0.2s, border-color 0.2s;
+}

test_api.py

@@ -14,14 +14,25 @@ def test_healthcheck():
     try:
         response = requests.get(f"{BASE_URL}/health", timeout=10)
         print(f"Status: {response.status_code}")
+        print(f"Headers: {dict(response.headers)}")
+        print(f"Content-Type: {response.headers.get('content-type', 'unknown')}")
+        print(f"Response length: {len(response.text)} bytes")
+        
         if response.status_code == 200:
-            print(f"Response: {json.dumps(response.json(), indent=2)}")
-            return True
+            try:
+                data = response.json()
+                print(f"Response: {json.dumps(data, indent=2)}")
+                return True
+            except json.JSONDecodeError:
+                print(f"Not JSON. First 200 chars: {response.text[:200]}")
+                return False
         else:
-            print(f"Error: {response.text}")
+            print(f"Error: {response.text[:200]}")
             return False
     except Exception as e:
         print(f"Exception: {e}")
+        import traceback
+        traceback.print_exc()
         return False
 
 def test_generate():
@@ -38,33 +49,57 @@ def test_generate():
             f"{BASE_URL}/v1/generate",
             json=payload,
             headers={"Content-Type": "application/json"},
-            timeout=60  # Longer timeout for model loading
+            timeout=120  # Longer timeout for model loading
         )
         print(f"Status: {response.status_code}")
+        print(f"Headers: {dict(response.headers)}")
+        
         if response.status_code == 200:
-            result = response.json()
-            print(f"Response keys: {list(result.keys())}")
-            if "text" in result:
-                print(f"Generated text (first 200 chars): {result['text'][:200]}...")
-            else:
-                print(f"Full response: {json.dumps(result, indent=2)}")
-            return True
+            try:
+                result = response.json()
+                print(f"Response keys: {list(result.keys())}")
+                if "text" in result:
+                    print(f"Generated text (first 300 chars): {result['text'][:300]}...")
+                else:
+                    print(f"Full response: {json.dumps(result, indent=2)}")
+                return True
+            except json.JSONDecodeError:
+                print(f"Not JSON. First 200 chars: {response.text[:200]}")
+                return False
         else:
-            print(f"Error: {response.text}")
+            print(f"Error: {response.text[:500]}")
             return False
     except Exception as e:
         print(f"Exception: {e}")
+        import traceback
+        traceback.print_exc()
         return False
 
-def test_gradio_ui():
-    """Test the Gradio UI endpoint."""
-    print("\nTesting GET /gradio (UI redirect target)...")
+def test_gradio_api():
+    """Test Gradio's built-in API endpoint."""
+    print("\nTesting Gradio API /api/predict...")
     try:
-        response = requests.get(f"{BASE_URL}/gradio", timeout=10)
+        # Gradio creates /api/predict endpoints automatically
+        # We need to find the function index - usually 0 for the first function
+        payload = {
+            "data": [
+                "You are a router agent. User query: What is 2+2?",
+                100,
+                0.2,
+                0.9
+            ],
+            "fn_index": 0
+        }
+        response = requests.post(
+            f"{BASE_URL}/api/predict",
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=120
+        )
         print(f"Status: {response.status_code}")
         if response.status_code == 200:
-            print(f"Response length: {len(response.text)} chars")
-            print(f"Response type: {response.headers.get('content-type', 'unknown')}")
+            result = response.json()
+            print(f"Gradio API Response: {json.dumps(result, indent=2)[:500]}...")
             return True
         else:
             print(f"Error: {response.text[:200]}")
@@ -73,6 +108,21 @@ def test_gradio_ui():
         print(f"Exception: {e}")
         return False
 
+def test_root():
+    """Test the root endpoint."""
+    print("\nTesting GET / (root)...")
+    try:
+        response = requests.get(f"{BASE_URL}/", timeout=10)
+        print(f"Status: {response.status_code}")
+        print(f"Content-Type: {response.headers.get('content-type', 'unknown')}")
+        print(f"Response length: {len(response.text)} chars")
+        if response.status_code == 200:
+            return True
+        return False
+    except Exception as e:
+        print(f"Exception: {e}")
+        return False
+
 def main():
     """Run all API tests."""
     print("=" * 60)
@@ -81,15 +131,16 @@ def main():
     print(f"Base URL: {BASE_URL}\n")
     
     # Wait a moment for Space to be ready
-    print("Waiting 5 seconds for Space to be ready...")
-    time.sleep(5)
+    print("Waiting 3 seconds for Space to be ready...")
+    time.sleep(3)
     
     results = []
     
     # Test endpoints
+    results.append(("Root", test_root()))
     results.append(("Health Check", test_healthcheck()))
     results.append(("Generate", test_generate()))
-    results.append(("Gradio UI", test_gradio_ui()))
+    results.append(("Gradio API", test_gradio_api()))
     
     # Summary
     print("\n" + "=" * 60)