app.py

import os
import time
import gc
import threading
from itertools import islice
from datetime import datetime
import re  # for parsing <think> blocks
import gradio as gr
import torch
from transformers import pipeline, TextIteratorStreamer
from transformers import AutoTokenizer
from ddgs import DDGS
import spaces  # Import spaces early to enable ZeroGPU support

access_token=os.environ['HF_TOKEN']

# Optional: Disable GPU visibility if you wish to force CPU usage
# os.environ["CUDA_VISIBLE_DEVICES"] = ""

# ------------------------------
# Global Cancellation Event
# ------------------------------
cancel_event = threading.Event()

# ------------------------------
# Torch-Compatible Model Definitions with Adjusted Descriptions
# ------------------------------
MODELS = {
    # ~30.5B total parameters (MoE: 3.3B activated)
    "Qwen3-30B-A3B-Thinking-2507-FP8": {
        "repo_id": "Qwen/Qwen3-30B-A3B-Thinking-2507-FP8",
        "description": "FP8-quantized MoE model with 30.5B total parameters (3.3B activated), 128 experts (8 activated), 48 layers, and native 262,144-token context. Optimized for complex reasoning tasks with enhanced thinking capabilities in mathematics, coding, science, and agent benchmarks. Supports only thinking mode; includes automatic reasoning delimiters."
    },
    "Qwen3-30B-A3B-Instruct-2507-FP8": {
        "repo_id": "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8",
        "description": "FP8-quantized instruct-tuned variant of Qwen3-30B-A3B (30.5B total params, 3.3B activated), featuring strong general capabilities in instruction following, tool usage, text generation, and 256K long-context understanding. Ideal for agentic and multi-turn dialogue applications."
    },

    # ~235B total parameters (MoE: 22B activated) — included for reference if added later
    # "Qwen3-235B-A22B-Thinking": { ... },

    # 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."
    },
    "Qwen/Qwen3-14B-FP8": {
        "repo_id": "Qwen/Qwen3-14B-FP8",
        "description": "FP8-quantized version of Qwen3-14B for efficient inference."
    },

    # ~15B (commented out in original, but larger than 14B)
    # "Apriel-1.5-15b-Thinker": { ... },

    # 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)"
    },
    "Phi-4-mini-Instruct": {
        "repo_id": "microsoft/Phi-4-mini-instruct",
        "description": "Phi-4-mini-Instruct (4.3B parameters)"
    },

    # 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."
    },
    "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."
    },
    "Gemma-3-4B-IT": {
        "repo_id": "unsloth/gemma-3-4b-it",
        "description": "Gemma-3-4B-IT"
    },
    "MiniCPM3-4B": {
        "repo_id": "openbmb/MiniCPM3-4B",
        "description": "MiniCPM3-4B"
    },
    "Gemma-3n-E4B": {
        "repo_id": "google/gemma-3n-E4B",
        "description": "Gemma 3n base model with effective 4 B parameters (≈3 GB VRAM)"
    },
    "SmallThinker-4BA0.6B-Instruct": {
        "repo_id": "PowerInfer/SmallThinker-4BA0.6B-Instruct",
        "description": "SmallThinker 4 B backbone with 0.6 B activated parameters, instruction‑tuned"
    },

    # ~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"
    },
    "Llama-3.2-Taiwan-3B-Instruct": {
        "repo_id": "lianghsun/Llama-3.2-Taiwan-3B-Instruct",
        "description": "Llama-3.2-Taiwan-3B-Instruct"
    },
    "Qwen2.5-3B-Instruct": {
        "repo_id": "Qwen/Qwen2.5-3B-Instruct",
        "description": "Qwen2.5-3B-Instruct"
    },
    "Qwen2.5-Omni-3B": {
        "repo_id": "Qwen/Qwen2.5-Omni-3B",
        "description": "Qwen2.5-Omni-3B"
    },
    "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."
    },

    # 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."
    },

    # 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."
    },

    # ~2B (effective)
    "Gemma-3n-E2B": {
        "repo_id": "google/gemma-3n-E2B",
        "description": "Gemma 3n base model with effective 2 B parameters (≈2 GB VRAM)"
    },

    # 1.5B
    "Nemotron-Research-Reasoning-Qwen-1.5B": {
        "repo_id": "nvidia/Nemotron-Research-Reasoning-Qwen-1.5B",
        "description": "Nemotron-Research-Reasoning-Qwen-1.5B"
    },
    "Falcon-H1-1.5B-Instruct": {
        "repo_id": "tiiuae/Falcon-H1-1.5B-Instruct",
        "description": "Falcon‑H1 model with 1.5 B parameters, instruction‑tuned"
    },
    "Qwen2.5-Taiwan-1.5B-Instruct": {
        "repo_id": "benchang1110/Qwen2.5-Taiwan-1.5B-Instruct",
        "description": "Qwen2.5-Taiwan-1.5B-Instruct"
    },

    # 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."
    },

    # 1.1B
    "Taiwan-ELM-1_1B-Instruct": {
        "repo_id": "liswei/Taiwan-ELM-1_1B-Instruct",
        "description": "Taiwan-ELM-1_1B-Instruct"
    },

    # 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"
    },

    # 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."
    },

    # 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."
    },
    "Qwen3-0.6B-Taiwan": {
        "repo_id": "ShengweiPeng/Qwen3-0.6B-Taiwan",
        "description": "Qwen3-Taiwan model with 0.6 B parameters"
    },

    # 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"
    },

    # 360M
    "SmolLM2-360M-Instruct": {
        "repo_id": "HuggingFaceTB/SmolLM2-360M-Instruct",
        "description": "Original SmolLM2‑360M Instruct"
    },
    "SmolLM2-360M-Instruct-TaiwanChat": {
        "repo_id": "Luigi/SmolLM2-360M-Instruct-TaiwanChat",
        "description": "SmolLM2‑360M Instruct fine-tuned on TaiwanChat"
    },

    # 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."
    },

    # 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."
    },
    "Gemma-3-Taiwan-270M-it": {
        "repo_id": "lianghsun/Gemma-3-Taiwan-270M-it",
        "description": "google/gemma-3-270m-it fintuned on Taiwan Chinese dataset"
    },
    "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.",
    },
    "Taiwan-ELM-270M-Instruct": {
        "repo_id": "liswei/Taiwan-ELM-270M-Instruct",
        "description": "Taiwan-ELM-270M-Instruct"
    },

    # 135M
    "SmolLM2-135M-multilingual-base": {
        "repo_id": "agentlans/SmolLM2-135M-multilingual-base",
        "description": "SmolLM2-135M-multilingual-base"
    },
    "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"
    },
    "SmolLM2_135M_Grpo_Gsm8k": {
        "repo_id": "prithivMLmods/SmolLM2_135M_Grpo_Gsm8k",
        "description": "SmolLM2_135M_Grpo_Gsm8k"
    },
    "SmolLM2-135M-Instruct": {
        "repo_id": "HuggingFaceTB/SmolLM2-135M-Instruct",
        "description": "Original SmolLM2‑135M Instruct"
    },
    "SmolLM2-135M-Instruct-TaiwanChat": {
        "repo_id": "Luigi/SmolLM2-135M-Instruct-TaiwanChat",
        "description": "SmolLM2‑135M Instruct fine-tuned on TaiwanChat"
    },
}

# 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,
                torch_dtype=dtype,
                device_map="auto",
                use_cache=False,      # ← disable 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"
    )
    PIPELINES[model_name] = pipe
    return pipe


def retrieve_context(query, max_results=6, max_chars=50):
    """
    Retrieve search snippets from DuckDuckGo (runs in background).
    Returns a list of result strings.
    """
    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

@spaces.GPU(duration=60)
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.
    """
    cancel_event.clear()
    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))
            )
        )
        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

        # Stream tokens
        for chunk in streamer:
            if cancel_event.is_set():
                break
            text = chunk

            # Detect start of thinking
            if not in_thought and '<think>' in text:
                in_thought = True
                # Insert thought placeholder
                history.append({
                    'role': 'assistant',
                    'content': '',
                    'metadata': {'title': '💭 Thought'}
                })
                # Capture after opening tag
                after = text.split('<think>', 1)[1]
                thought_buf += after
                # If closing tag in same chunk
                if '</think>' in thought_buf:
                    before, after2 = thought_buf.split('</think>', 1)
                    history[-1]['content'] = before.strip()
                    in_thought = False
                    # Start answer buffer
                    answer_buf = after2
                    history.append({'role': 'assistant', 'content': answer_buf})
                else:
                    history[-1]['content'] = thought_buf
                yield history, debug
                continue

            # Continue thought streaming
            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
                    # Start answer buffer
                    answer_buf = after2
                    history.append({'role': 'assistant', 'content': answer_buf})
                else:
                    history[-1]['content'] = thought_buf
                yield history, debug
                continue

            # Stream answer
            if not answer_buf:
                history.append({'role': 'assistant', 'content': ''})
            answer_buf += text
            history[-1]['content'] = answer_buf
            yield history, debug

        gen_thread.join()
        yield history, debug + prompt_debug
    except Exception as e:
        history.append({'role': 'assistant', 'content': f"Error: {e}"})
        yield history, debug
    finally:
        gc.collect()


def cancel_generation():
    cancel_event.set()
    return 'Generation cancelled.'


def update_default_prompt(enable_search):
    return f"You are a helpful assistant."

# ------------------------------
# Gradio UI
# ------------------------------
with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
    gr.Markdown("## 🧠 ZeroGPU LLM Inference with Web Search")
    gr.Markdown("Interact with the model. Select parameters and chat below.")
    with gr.Row():
        with gr.Column(scale=3):
            model_dd = gr.Dropdown(label="Select Model", choices=list(MODELS.keys()), value=list(MODELS.keys())[0])
            search_chk = gr.Checkbox(label="Enable Web Search", value=True)
            sys_prompt = gr.Textbox(label="System Prompt", lines=3, value=update_default_prompt(search_chk.value))
            gr.Markdown("### Generation Parameters")
            max_tok = gr.Slider(64, 16384, value=2048, step=32, label="Max Tokens")
            temp = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
            k = gr.Slider(1, 100, value=40, step=1, label="Top-K")
            p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
            rp = gr.Slider(1.0, 2.0, value=1.2, step=0.1, label="Repetition Penalty")
            gr.Markdown("### Web Search Settings")
            mr = gr.Number(value=4, precision=0, label="Max Results")
            mc = gr.Number(value=50, precision=0, label="Max Chars/Result")
            st = gr.Slider(minimum=0.0, maximum=30.0, step=0.5, value=5.0, label="Search Timeout (s)")
            clr = gr.Button("Clear Chat")
            cnl = gr.Button("Cancel Generation")
        with gr.Column(scale=7):
            chat = gr.Chatbot(type="messages")
            txt = gr.Textbox(placeholder="Type your message and press Enter...")
            dbg = gr.Markdown()

    search_chk.change(fn=update_default_prompt, inputs=search_chk, outputs=sys_prompt)
    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
    cnl.click(fn=cancel_generation, outputs=dbg)
    txt.submit(fn=chat_response,
               inputs=[txt, chat, sys_prompt, search_chk, mr, mc,
                       model_dd, max_tok, temp, k, p, rp, st],
               outputs=[chat, dbg])
    demo.launch()